Java Code Examples for org.apache.hadoop.mapreduce.v2.util.MRBuilderUtils#newTaskId()

private ContainerRequestEvent
    createReq(JobId jobId, int taskAttemptId, int memory, String[] hosts,
        boolean earlierFailedAttempt, boolean reduce) {
  TaskId taskId;
  if (reduce) {
    taskId = MRBuilderUtils.newTaskId(jobId, 0, TaskType.REDUCE);
  } else {
    taskId = MRBuilderUtils.newTaskId(jobId, 0, TaskType.MAP);
  }
  TaskAttemptId attemptId = MRBuilderUtils.newTaskAttemptId(taskId,
      taskAttemptId);
  Resource containerNeed = Resource.newInstance(memory, 1);
  if (earlierFailedAttempt) {
    return ContainerRequestEvent
        .createContainerRequestEventForFailedContainer(attemptId,
            containerNeed);
  }
  return new ContainerRequestEvent(attemptId, containerNeed, hosts,
      new String[] { NetworkTopology.DEFAULT_RACK });
}
 

private static TaskAttemptCompletionEvent createTce(int eventId,
    boolean isMap, TaskAttemptCompletionEventStatus status) {
  JobId jid = MRBuilderUtils.newJobId(12345, 1, 1);
  TaskId tid = MRBuilderUtils.newTaskId(jid, 0,
      isMap ? org.apache.hadoop.mapreduce.v2.api.records.TaskType.MAP
          : org.apache.hadoop.mapreduce.v2.api.records.TaskType.REDUCE);
  TaskAttemptId attemptId = MRBuilderUtils.newTaskAttemptId(tid, 0);
  RecordFactory recordFactory = RecordFactoryProvider.getRecordFactory(null);
  TaskAttemptCompletionEvent tce = recordFactory
      .newRecordInstance(TaskAttemptCompletionEvent.class);
  tce.setEventId(eventId);
  tce.setAttemptId(attemptId);
  tce.setStatus(status);
  return tce;
}
 

public static TaskAttemptId makeTaskAttemptId(long ts, int appId, int taskId, 
    TaskType taskType, int id) {
  ApplicationId aID = ApplicationId.newInstance(ts, appId);
  JobId jID = MRBuilderUtils.newJobId(aID, id);
  TaskId tID = MRBuilderUtils.newTaskId(jID, taskId, taskType);
  return MRBuilderUtils.newTaskAttemptId(tID, id);
}
 

@Test
public void testRenameMapOutputForReduce() throws Exception {
  final JobConf conf = new JobConf();

  final MROutputFiles mrOutputFiles = new MROutputFiles();
  mrOutputFiles.setConf(conf);

  // make sure both dirs are distinct
  //
  conf.set(MRConfig.LOCAL_DIR, localDirs[0].toString());
  final Path mapOut = mrOutputFiles.getOutputFileForWrite(1);
  conf.set(MRConfig.LOCAL_DIR, localDirs[1].toString());
  final Path mapOutIdx = mrOutputFiles.getOutputIndexFileForWrite(1);
  Assert.assertNotEquals("Paths must be different!",
      mapOut.getParent(), mapOutIdx.getParent());

  // make both dirs part of LOCAL_DIR
  conf.setStrings(MRConfig.LOCAL_DIR, localDirs);

  final FileContext lfc = FileContext.getLocalFSFileContext(conf);
  lfc.create(mapOut, EnumSet.of(CREATE)).close();
  lfc.create(mapOutIdx, EnumSet.of(CREATE)).close();

  final JobId jobId = MRBuilderUtils.newJobId(12345L, 1, 2);
  final TaskId tid = MRBuilderUtils.newTaskId(jobId, 0, TaskType.MAP);
  final TaskAttemptId taid = MRBuilderUtils.newTaskAttemptId(tid, 0);

  LocalContainerLauncher.renameMapOutputForReduce(conf, taid, mrOutputFiles);
}
 

private ContainerFailedEvent createFailEvent(JobId jobId, int taskAttemptId,
    String host, boolean reduce) {
  TaskId taskId;
  if (reduce) {
    taskId = MRBuilderUtils.newTaskId(jobId, 0, TaskType.REDUCE);
  } else {
    taskId = MRBuilderUtils.newTaskId(jobId, 0, TaskType.MAP);
  }
  TaskAttemptId attemptId = MRBuilderUtils.newTaskAttemptId(taskId,
      taskAttemptId);
  return new ContainerFailedEvent(attemptId, host);    
}
 

private ContainerAllocatorEvent createDeallocateEvent(JobId jobId,
    int taskAttemptId, boolean reduce) {
  TaskId taskId;
  if (reduce) {
    taskId = MRBuilderUtils.newTaskId(jobId, 0, TaskType.REDUCE);
  } else {
    taskId = MRBuilderUtils.newTaskId(jobId, 0, TaskType.MAP);
  }
  TaskAttemptId attemptId =
      MRBuilderUtils.newTaskAttemptId(taskId, taskAttemptId);
  return new ContainerAllocatorEvent(attemptId,
      ContainerAllocator.EventType.CONTAINER_DEALLOCATE);
}
 

private ContainerAllocatorEvent createDeallocateEvent(JobId jobId,
    int taskAttemptId, boolean reduce) {
  TaskId taskId;
  if (reduce) {
    taskId = MRBuilderUtils.newTaskId(jobId, 0, TaskType.REDUCE);
  } else {
    taskId = MRBuilderUtils.newTaskId(jobId, 0, TaskType.MAP);
  }
  TaskAttemptId attemptId =
      MRBuilderUtils.newTaskAttemptId(taskId, taskAttemptId);
  return new ContainerAllocatorEvent(attemptId,
      ContainerAllocator.EventType.CONTAINER_DEALLOCATE);
}
 

@Test (timeout=10000)
public void testCompletedTask() throws Exception {
  HistoryFileInfo info = mock(HistoryFileInfo.class);
  when(info.getConfFile()).thenReturn(fullConfPath);
  completedJob =
    new CompletedJob(conf, jobId, fullHistoryPath, loadTasks, "user",
        info, jobAclsManager);
  TaskId mt1Id = MRBuilderUtils.newTaskId(jobId, 0, TaskType.MAP);
  TaskId rt1Id = MRBuilderUtils.newTaskId(jobId, 0, TaskType.REDUCE);
  
  Map<TaskId, Task> mapTasks = completedJob.getTasks(TaskType.MAP);
  Map<TaskId, Task> reduceTasks = completedJob.getTasks(TaskType.REDUCE);
  assertEquals(10, mapTasks.size());
  assertEquals(2, reduceTasks.size());
  
  Task mt1 = mapTasks.get(mt1Id);
  assertEquals(1, mt1.getAttempts().size());
  assertEquals(TaskState.SUCCEEDED, mt1.getState());
  TaskReport mt1Report = mt1.getReport();
  assertEquals(TaskState.SUCCEEDED, mt1Report.getTaskState());
  assertEquals(mt1Id, mt1Report.getTaskId());
  Task rt1 = reduceTasks.get(rt1Id);
  assertEquals(1, rt1.getAttempts().size());
  assertEquals(TaskState.SUCCEEDED, rt1.getState());
  TaskReport rt1Report = rt1.getReport();
  assertEquals(TaskState.SUCCEEDED, rt1Report.getTaskState());
  assertEquals(rt1Id, rt1Report.getTaskId());
}
 

private TaskAttemptImpl createMapTaskAttemptImplForTest(
    EventHandler eventHandler, TaskSplitMetaInfo taskSplitMetaInfo, Clock clock) {
  ApplicationId appId = ApplicationId.newInstance(1, 1);
  JobId jobId = MRBuilderUtils.newJobId(appId, 1);
  TaskId taskId = MRBuilderUtils.newTaskId(jobId, 1, TaskType.MAP);
  TaskAttemptListener taListener = mock(TaskAttemptListener.class);
  Path jobFile = mock(Path.class);
  JobConf jobConf = new JobConf();
  TaskAttemptImpl taImpl =
      new MapTaskAttemptImpl(taskId, 1, eventHandler, jobFile, 1,
          taskSplitMetaInfo, jobConf, taListener, null,
          null, clock, null);
  return taImpl;
}
 

@Test
public void testRenameMapOutputForReduce() throws Exception {
  final JobConf conf = new JobConf();

  final MROutputFiles mrOutputFiles = new MROutputFiles();
  mrOutputFiles.setConf(conf);

  // make sure both dirs are distinct
  //
  conf.set(MRConfig.LOCAL_DIR, localDirs[0].toString());
  final Path mapOut = mrOutputFiles.getOutputFileForWrite(1);
  conf.set(MRConfig.LOCAL_DIR, localDirs[1].toString());
  final Path mapOutIdx = mrOutputFiles.getOutputIndexFileForWrite(1);
  Assert.assertNotEquals("Paths must be different!",
      mapOut.getParent(), mapOutIdx.getParent());

  // make both dirs part of LOCAL_DIR
  conf.setStrings(MRConfig.LOCAL_DIR, localDirs);

  final FileContext lfc = FileContext.getLocalFSFileContext(conf);
  lfc.create(mapOut, EnumSet.of(CREATE)).close();
  lfc.create(mapOutIdx, EnumSet.of(CREATE)).close();

  final JobId jobId = MRBuilderUtils.newJobId(12345L, 1, 2);
  final TaskId tid = MRBuilderUtils.newTaskId(jobId, 0, TaskType.MAP);
  final TaskAttemptId taid = MRBuilderUtils.newTaskAttemptId(tid, 0);

  LocalContainerLauncher.renameMapOutputForReduce(conf, taid, mrOutputFiles);
}
 

@Test
public void testContainerCleanedWhileCommitting() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId =
    ApplicationAttemptId.newInstance(appId, 0);
  JobId jobId = MRBuilderUtils.newJobId(appId, 1);
  TaskId taskId = MRBuilderUtils.newTaskId(jobId, 1, TaskType.MAP);
  TaskAttemptId attemptId = MRBuilderUtils.newTaskAttemptId(taskId, 0);
  Path jobFile = mock(Path.class);

  MockEventHandler eventHandler = new MockEventHandler();
  TaskAttemptListener taListener = mock(TaskAttemptListener.class);
  when(taListener.getAddress()).thenReturn(new InetSocketAddress("localhost", 0));

  JobConf jobConf = new JobConf();
  jobConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  jobConf.setBoolean("fs.file.impl.disable.cache", true);
  jobConf.set(JobConf.MAPRED_MAP_TASK_ENV, "");
  jobConf.set(MRJobConfig.APPLICATION_ATTEMPT_ID, "10");

  TaskSplitMetaInfo splits = mock(TaskSplitMetaInfo.class);
  when(splits.getLocations()).thenReturn(new String[] {});

  AppContext appCtx = mock(AppContext.class);
  ClusterInfo clusterInfo = mock(ClusterInfo.class);
  Resource resource = mock(Resource.class);
  when(appCtx.getClusterInfo()).thenReturn(clusterInfo);
  when(resource.getMemory()).thenReturn(1024);

  TaskAttemptImpl taImpl =
    new MapTaskAttemptImpl(taskId, 1, eventHandler, jobFile, 1,
        splits, jobConf, taListener,
        new Token(), new Credentials(),
        new SystemClock(), appCtx);

  NodeId nid = NodeId.newInstance("127.0.0.1", 0);
  ContainerId contId = ContainerId.newContainerId(appAttemptId, 3);
  Container container = mock(Container.class);
  when(container.getId()).thenReturn(contId);
  when(container.getNodeId()).thenReturn(nid);
  when(container.getNodeHttpAddress()).thenReturn("localhost:0");

  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_SCHEDULE));
  taImpl.handle(new TaskAttemptContainerAssignedEvent(attemptId,
      container, mock(Map.class)));
  taImpl.handle(new TaskAttemptContainerLaunchedEvent(attemptId, 0));
  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_COMMIT_PENDING));

  assertEquals("Task attempt is not in commit pending state", taImpl.getState(),
      TaskAttemptState.COMMIT_PENDING);
  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_CONTAINER_CLEANED));
  assertFalse("InternalError occurred trying to handle TA_CONTAINER_CLEANED",
      eventHandler.internalError);
  assertEquals("Task attempt is assigned locally", Locality.OFF_SWITCH,
      taImpl.getLocality());
}
 

@Test
public void testAppDiognosticEventOnUnassignedTask() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(
      appId, 0);
  JobId jobId = MRBuilderUtils.newJobId(appId, 1);
  TaskId taskId = MRBuilderUtils.newTaskId(jobId, 1, TaskType.MAP);
  TaskAttemptId attemptId = MRBuilderUtils.newTaskAttemptId(taskId, 0);
  Path jobFile = mock(Path.class);

  MockEventHandler eventHandler = new MockEventHandler();
  TaskAttemptListener taListener = mock(TaskAttemptListener.class);
  when(taListener.getAddress()).thenReturn(
      new InetSocketAddress("localhost", 0));

  JobConf jobConf = new JobConf();
  jobConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  jobConf.setBoolean("fs.file.impl.disable.cache", true);
  jobConf.set(JobConf.MAPRED_MAP_TASK_ENV, "");
  jobConf.set(MRJobConfig.APPLICATION_ATTEMPT_ID, "10");

  TaskSplitMetaInfo splits = mock(TaskSplitMetaInfo.class);
  when(splits.getLocations()).thenReturn(new String[] { "127.0.0.1" });

  AppContext appCtx = mock(AppContext.class);
  ClusterInfo clusterInfo = mock(ClusterInfo.class);
  Resource resource = mock(Resource.class);
  when(appCtx.getClusterInfo()).thenReturn(clusterInfo);
  when(resource.getMemory()).thenReturn(1024);

  TaskAttemptImpl taImpl = new MapTaskAttemptImpl(taskId, 1, eventHandler,
      jobFile, 1, splits, jobConf, taListener,
      new Token(), new Credentials(), new SystemClock(), appCtx);

  NodeId nid = NodeId.newInstance("127.0.0.1", 0);
  ContainerId contId = ContainerId.newContainerId(appAttemptId, 3);
  Container container = mock(Container.class);
  when(container.getId()).thenReturn(contId);
  when(container.getNodeId()).thenReturn(nid);
  when(container.getNodeHttpAddress()).thenReturn("localhost:0");
  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_SCHEDULE));
  taImpl.handle(new TaskAttemptDiagnosticsUpdateEvent(attemptId,
      "Task got killed"));
  assertFalse(
      "InternalError occurred trying to handle TA_DIAGNOSTICS_UPDATE on assigned task",
      eventHandler.internalError);
}
 

@Test
public void testTooManyFetchFailureAfterKill() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId =
    ApplicationAttemptId.newInstance(appId, 0);
  JobId jobId = MRBuilderUtils.newJobId(appId, 1);
  TaskId taskId = MRBuilderUtils.newTaskId(jobId, 1, TaskType.MAP);
  TaskAttemptId attemptId = MRBuilderUtils.newTaskAttemptId(taskId, 0);
  Path jobFile = mock(Path.class);

  MockEventHandler eventHandler = new MockEventHandler();
  TaskAttemptListener taListener = mock(TaskAttemptListener.class);
  when(taListener.getAddress()).thenReturn(new InetSocketAddress("localhost", 0));

  JobConf jobConf = new JobConf();
  jobConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  jobConf.setBoolean("fs.file.impl.disable.cache", true);
  jobConf.set(JobConf.MAPRED_MAP_TASK_ENV, "");
  jobConf.set(MRJobConfig.APPLICATION_ATTEMPT_ID, "10");

  TaskSplitMetaInfo splits = mock(TaskSplitMetaInfo.class);
  when(splits.getLocations()).thenReturn(new String[] {"127.0.0.1"});

  AppContext appCtx = mock(AppContext.class);
  ClusterInfo clusterInfo = mock(ClusterInfo.class);
  Resource resource = mock(Resource.class);
  when(appCtx.getClusterInfo()).thenReturn(clusterInfo);
  when(resource.getMemory()).thenReturn(1024);

  TaskAttemptImpl taImpl =
    new MapTaskAttemptImpl(taskId, 1, eventHandler, jobFile, 1,
      splits, jobConf, taListener,
      mock(Token.class), new Credentials(),
      new SystemClock(), appCtx);

  NodeId nid = NodeId.newInstance("127.0.0.1", 0);
  ContainerId contId = ContainerId.newContainerId(appAttemptId, 3);
  Container container = mock(Container.class);
  when(container.getId()).thenReturn(contId);
  when(container.getNodeId()).thenReturn(nid);
  when(container.getNodeHttpAddress()).thenReturn("localhost:0");

  taImpl.handle(new TaskAttemptEvent(attemptId,
    TaskAttemptEventType.TA_SCHEDULE));
  taImpl.handle(new TaskAttemptContainerAssignedEvent(attemptId,
    container, mock(Map.class)));
  taImpl.handle(new TaskAttemptContainerLaunchedEvent(attemptId, 0));
  taImpl.handle(new TaskAttemptEvent(attemptId,
    TaskAttemptEventType.TA_DONE));
  taImpl.handle(new TaskAttemptEvent(attemptId,
    TaskAttemptEventType.TA_CONTAINER_CLEANED));

  assertEquals("Task attempt is not in succeeded state", taImpl.getState(),
    TaskAttemptState.SUCCEEDED);
  taImpl.handle(new TaskAttemptEvent(attemptId,
    TaskAttemptEventType.TA_KILL));
  assertEquals("Task attempt is not in KILLED state", taImpl.getState(),
    TaskAttemptState.KILLED);
  taImpl.handle(new TaskAttemptEvent(attemptId,
    TaskAttemptEventType.TA_TOO_MANY_FETCH_FAILURE));
  assertEquals("Task attempt is not in KILLED state, still", taImpl.getState(),
    TaskAttemptState.KILLED);
  assertFalse("InternalError occurred trying to handle TA_CONTAINER_CLEANED",
    eventHandler.internalError);
}
 

private TaskId createTaskId(long clusterTimestamp, int jobIdInt,
    int taskIdInt, TaskType taskType) {
  return MRBuilderUtils.newTaskId(createJobId(clusterTimestamp, jobIdInt),
      taskIdInt, taskType);
}
 

@Test
public void testAppDiognosticEventOnUnassignedTask() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(
      appId, 0);
  JobId jobId = MRBuilderUtils.newJobId(appId, 1);
  TaskId taskId = MRBuilderUtils.newTaskId(jobId, 1, TaskType.MAP);
  TaskAttemptId attemptId = MRBuilderUtils.newTaskAttemptId(taskId, 0);
  Path jobFile = mock(Path.class);

  MockEventHandler eventHandler = new MockEventHandler();
  TaskAttemptListener taListener = mock(TaskAttemptListener.class);
  when(taListener.getAddress()).thenReturn(
      new InetSocketAddress("localhost", 0));

  JobConf jobConf = new JobConf();
  jobConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  jobConf.setBoolean("fs.file.impl.disable.cache", true);
  jobConf.set(JobConf.MAPRED_MAP_TASK_ENV, "");
  jobConf.set(MRJobConfig.APPLICATION_ATTEMPT_ID, "10");

  TaskSplitMetaInfo splits = mock(TaskSplitMetaInfo.class);
  when(splits.getLocations()).thenReturn(new String[] { "127.0.0.1" });

  AppContext appCtx = mock(AppContext.class);
  ClusterInfo clusterInfo = mock(ClusterInfo.class);
  Resource resource = mock(Resource.class);
  when(appCtx.getClusterInfo()).thenReturn(clusterInfo);
  when(resource.getMemory()).thenReturn(1024);

  TaskAttemptImpl taImpl = new MapTaskAttemptImpl(taskId, 1, eventHandler,
      jobFile, 1, splits, jobConf, taListener,
      new Token(), new Credentials(), new SystemClock(), appCtx);

  NodeId nid = NodeId.newInstance("127.0.0.1", 0);
  ContainerId contId = ContainerId.newContainerId(appAttemptId, 3);
  Container container = mock(Container.class);
  when(container.getId()).thenReturn(contId);
  when(container.getNodeId()).thenReturn(nid);
  when(container.getNodeHttpAddress()).thenReturn("localhost:0");
  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_SCHEDULE));
  taImpl.handle(new TaskAttemptDiagnosticsUpdateEvent(attemptId,
      "Task got killed"));
  assertFalse(
      "InternalError occurred trying to handle TA_DIAGNOSTICS_UPDATE on assigned task",
      eventHandler.internalError);
}
 

@Test
public void testContainerKillWhileCommitPending() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId,
      0);
  JobId jobId = MRBuilderUtils.newJobId(appId, 1);
  TaskId taskId = MRBuilderUtils.newTaskId(jobId, 1, TaskType.MAP);
  TaskAttemptId attemptId = MRBuilderUtils.newTaskAttemptId(taskId, 0);
  Path jobFile = mock(Path.class);

  MockEventHandler eventHandler = new MockEventHandler();
  TaskAttemptListener taListener = mock(TaskAttemptListener.class);
  when(taListener.getAddress()).thenReturn(
      new InetSocketAddress("localhost", 0));

  JobConf jobConf = new JobConf();
  jobConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  jobConf.setBoolean("fs.file.impl.disable.cache", true);
  jobConf.set(JobConf.MAPRED_MAP_TASK_ENV, "");
  jobConf.set(MRJobConfig.APPLICATION_ATTEMPT_ID, "10");

  TaskSplitMetaInfo splits = mock(TaskSplitMetaInfo.class);
  when(splits.getLocations()).thenReturn(new String[] { "127.0.0.1" });

  AppContext appCtx = mock(AppContext.class);
  ClusterInfo clusterInfo = mock(ClusterInfo.class);
  Resource resource = mock(Resource.class);
  when(appCtx.getClusterInfo()).thenReturn(clusterInfo);
  when(resource.getMemory()).thenReturn(1024);

  TaskAttemptImpl taImpl = new MapTaskAttemptImpl(taskId, 1, eventHandler,
      jobFile, 1, splits, jobConf, taListener, new Token(),
      new Credentials(), new SystemClock(), appCtx);

  NodeId nid = NodeId.newInstance("127.0.0.2", 0);
  ContainerId contId = ContainerId.newContainerId(appAttemptId, 3);
  Container container = mock(Container.class);
  when(container.getId()).thenReturn(contId);
  when(container.getNodeId()).thenReturn(nid);
  when(container.getNodeHttpAddress()).thenReturn("localhost:0");

  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_SCHEDULE));
  taImpl.handle(new TaskAttemptContainerAssignedEvent(attemptId, container,
      mock(Map.class)));
  taImpl.handle(new TaskAttemptContainerLaunchedEvent(attemptId, 0));
  assertEquals("Task attempt is not in running state", taImpl.getState(),
      TaskAttemptState.RUNNING);
  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_COMMIT_PENDING));
  assertEquals("Task should be in COMMIT_PENDING state",
      TaskAttemptStateInternal.COMMIT_PENDING, taImpl.getInternalState());
  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_KILL));
  assertFalse("InternalError occurred trying to handle TA_KILL",
      eventHandler.internalError);
  assertEquals("Task should be in KILLED state",
      TaskAttemptStateInternal.KILL_CONTAINER_CLEANUP,
      taImpl.getInternalState());
}
 

@Test
public void testAverageReduceTime() {
 
  Job job = mock(CompletedJob.class);
  final Task task1 = mock(Task.class);
  final Task task2 = mock(Task.class);

  JobId  jobId = MRBuilderUtils.newJobId(1L, 1, 1);

  final TaskId taskId1 = MRBuilderUtils.newTaskId(jobId, 1, TaskType.REDUCE);
  final TaskId taskId2 = MRBuilderUtils.newTaskId(jobId, 2, TaskType.REDUCE);

  final TaskAttemptId taskAttemptId1  = MRBuilderUtils.
  		newTaskAttemptId(taskId1, 1);
  final TaskAttemptId taskAttemptId2  = MRBuilderUtils.
  		newTaskAttemptId(taskId2, 2);

  final TaskAttempt taskAttempt1 = mock(TaskAttempt.class);
  final TaskAttempt taskAttempt2 = mock(TaskAttempt.class);

  JobReport jobReport = mock(JobReport.class);

  when(taskAttempt1.getState()).thenReturn(TaskAttemptState.SUCCEEDED);
  when(taskAttempt1.getLaunchTime()).thenReturn(0L);
  when(taskAttempt1.getShuffleFinishTime()).thenReturn(4L);
  when(taskAttempt1.getSortFinishTime()).thenReturn(6L);
  when(taskAttempt1.getFinishTime()).thenReturn(8L);

  when(taskAttempt2.getState()).thenReturn(TaskAttemptState.SUCCEEDED);
  when(taskAttempt2.getLaunchTime()).thenReturn(5L);
  when(taskAttempt2.getShuffleFinishTime()).thenReturn(10L);
  when(taskAttempt2.getSortFinishTime()).thenReturn(22L);
  when(taskAttempt2.getFinishTime()).thenReturn(42L);


  when(task1.getType()).thenReturn(TaskType.REDUCE);
  when(task2.getType()).thenReturn(TaskType.REDUCE);
  when(task1.getAttempts()).thenReturn
    (new HashMap<TaskAttemptId, TaskAttempt>() 
  		{{put(taskAttemptId1,taskAttempt1); }});
  when(task2.getAttempts()).thenReturn
    (new HashMap<TaskAttemptId, TaskAttempt>() 
  		  {{put(taskAttemptId2,taskAttempt2); }});

  when(job.getTasks()).thenReturn
    (new HashMap<TaskId, Task>() 
  		{{ put(taskId1,task1); put(taskId2, task2);  }});
  when(job.getID()).thenReturn(jobId);

  when(job.getReport()).thenReturn(jobReport);

  when(job.getName()).thenReturn("TestJobInfo");	  
  when(job.getState()).thenReturn(JobState.SUCCEEDED);	  

  JobInfo jobInfo = new JobInfo(job);

  Assert.assertEquals(11L, jobInfo.getAvgReduceTime().longValue());
}
 

@Test
public void testContainerKillWhileCommitPending() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId,
      0);
  JobId jobId = MRBuilderUtils.newJobId(appId, 1);
  TaskId taskId = MRBuilderUtils.newTaskId(jobId, 1, TaskType.MAP);
  TaskAttemptId attemptId = MRBuilderUtils.newTaskAttemptId(taskId, 0);
  Path jobFile = mock(Path.class);

  MockEventHandler eventHandler = new MockEventHandler();
  TaskAttemptListener taListener = mock(TaskAttemptListener.class);
  when(taListener.getAddress()).thenReturn(
      new InetSocketAddress("localhost", 0));

  JobConf jobConf = new JobConf();
  jobConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  jobConf.setBoolean("fs.file.impl.disable.cache", true);
  jobConf.set(JobConf.MAPRED_MAP_TASK_ENV, "");
  jobConf.set(MRJobConfig.APPLICATION_ATTEMPT_ID, "10");

  TaskSplitMetaInfo splits = mock(TaskSplitMetaInfo.class);
  when(splits.getLocations()).thenReturn(new String[] { "127.0.0.1" });

  AppContext appCtx = mock(AppContext.class);
  ClusterInfo clusterInfo = mock(ClusterInfo.class);
  Resource resource = mock(Resource.class);
  when(appCtx.getClusterInfo()).thenReturn(clusterInfo);
  when(resource.getMemory()).thenReturn(1024);

  TaskAttemptImpl taImpl = new MapTaskAttemptImpl(taskId, 1, eventHandler,
      jobFile, 1, splits, jobConf, taListener, new Token(),
      new Credentials(), new SystemClock(), appCtx);

  NodeId nid = NodeId.newInstance("127.0.0.2", 0);
  ContainerId contId = ContainerId.newContainerId(appAttemptId, 3);
  Container container = mock(Container.class);
  when(container.getId()).thenReturn(contId);
  when(container.getNodeId()).thenReturn(nid);
  when(container.getNodeHttpAddress()).thenReturn("localhost:0");

  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_SCHEDULE));
  taImpl.handle(new TaskAttemptContainerAssignedEvent(attemptId, container,
      mock(Map.class)));
  taImpl.handle(new TaskAttemptContainerLaunchedEvent(attemptId, 0));
  assertEquals("Task attempt is not in running state", taImpl.getState(),
      TaskAttemptState.RUNNING);
  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_COMMIT_PENDING));
  assertEquals("Task should be in COMMIT_PENDING state",
      TaskAttemptStateInternal.COMMIT_PENDING, taImpl.getInternalState());
  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_KILL));
  assertFalse("InternalError occurred trying to handle TA_KILL",
      eventHandler.internalError);
  assertEquals("Task should be in KILLED state",
      TaskAttemptStateInternal.KILL_CONTAINER_CLEANUP,
      taImpl.getInternalState());
}
 

@Test
public void testShuffleProviders() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 1);
  JobId jobId = MRBuilderUtils.newJobId(appId, 1);
  TaskId taskId = MRBuilderUtils.newTaskId(jobId, 1, TaskType.MAP);
  Path jobFile = mock(Path.class);

  EventHandler eventHandler = mock(EventHandler.class);
  TaskAttemptListener taListener = mock(TaskAttemptListener.class);
  when(taListener.getAddress()).thenReturn(new InetSocketAddress("localhost", 0));

  JobConf jobConf = new JobConf();
  jobConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  jobConf.setBoolean("fs.file.impl.disable.cache", true);
  jobConf.set(JobConf.MAPRED_MAP_TASK_ENV, "");

  jobConf.set(YarnConfiguration.NM_AUX_SERVICES,
    TestShuffleHandler1.MAPREDUCE_TEST_SHUFFLE_SERVICEID + "," +
    TestShuffleHandler2.MAPREDUCE_TEST_SHUFFLE_SERVICEID);

  String serviceName = TestShuffleHandler1.MAPREDUCE_TEST_SHUFFLE_SERVICEID;
  String serviceStr = String.format(YarnConfiguration.NM_AUX_SERVICE_FMT, serviceName);
  jobConf.set(serviceStr, TestShuffleHandler1.class.getName());

  serviceName = TestShuffleHandler2.MAPREDUCE_TEST_SHUFFLE_SERVICEID;
  serviceStr = String.format(YarnConfiguration.NM_AUX_SERVICE_FMT, serviceName);
  jobConf.set(serviceStr, TestShuffleHandler2.class.getName());

  jobConf.set(MRJobConfig.MAPREDUCE_JOB_SHUFFLE_PROVIDER_SERVICES,
                TestShuffleHandler1.MAPREDUCE_TEST_SHUFFLE_SERVICEID
                   + "," + TestShuffleHandler2.MAPREDUCE_TEST_SHUFFLE_SERVICEID);

  Credentials credentials = new Credentials();
  Token<JobTokenIdentifier> jobToken = new Token<JobTokenIdentifier>(
      ("tokenid").getBytes(), ("tokenpw").getBytes(),
      new Text("tokenkind"), new Text("tokenservice"));
  TaskAttemptImpl taImpl =
      new MapTaskAttemptImpl(taskId, 1, eventHandler, jobFile, 1,
          mock(TaskSplitMetaInfo.class), jobConf, taListener,
          jobToken, credentials,
          new SystemClock(), null);

  jobConf.set(MRJobConfig.APPLICATION_ATTEMPT_ID, taImpl.getID().toString());

  ContainerLaunchContext launchCtx =
      TaskAttemptImpl.createContainerLaunchContext(null,
          jobConf, jobToken, taImpl.createRemoteTask(),
          TypeConverter.fromYarn(jobId),
          mock(WrappedJvmID.class), taListener,
          credentials);

  Map<String, ByteBuffer> serviceDataMap = launchCtx.getServiceData();
  Assert.assertNotNull("TestShuffleHandler1 is missing", serviceDataMap.get(TestShuffleHandler1.MAPREDUCE_TEST_SHUFFLE_SERVICEID));
  Assert.assertNotNull("TestShuffleHandler2 is missing", serviceDataMap.get(TestShuffleHandler2.MAPREDUCE_TEST_SHUFFLE_SERVICEID));
  Assert.assertTrue("mismatch number of services in map", serviceDataMap.size() == 3); // 2 that we entered + 1 for the built-in shuffle-provider
}
 

/**
 * Prepares the Task Id instance
 * @param jobId
 * @param id
 * @param taskType
 * @return the TaskId
 */
public static TaskId getTaskId(JobId jobId, int id, TaskType taskType){
	return MRBuilderUtils.newTaskId(jobId, id, taskType);
}