Java Code Examples for org.apache.hadoop.yarn.api.records.ContainerId#newInstance()

@Test(timeout = 5000)
public void testGetTaskShouldDie() throws InterruptedException, ExecutionException {
  ListeningExecutorService executor = null;
  try {
    ExecutorService rawExecutor = Executors.newFixedThreadPool(1);
    executor = MoreExecutors.listeningDecorator(rawExecutor);
    ApplicationId appId = ApplicationId.newInstance(10000, 1);
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
    @SuppressWarnings("deprecation")
    ContainerId containerId = ContainerId.newInstance(appAttemptId, 1);

    TaskExecutionTestHelpers.TezTaskUmbilicalForTest
        umbilical = new TaskExecutionTestHelpers.TezTaskUmbilicalForTest();
    ContainerContext containerContext = new ContainerContext(containerId.toString());

    ContainerReporter containerReporter = new ContainerReporter(umbilical, containerContext, 100);
    ListenableFuture<ContainerTask> getTaskFuture = executor.submit(containerReporter);

    getTaskFuture.get();
    assertEquals(1, umbilical.getTaskInvocations);

  } finally {
    executor.shutdownNow();
  }
}
 

private void testContainerStoppedEvent() throws Exception {
  ContainerStoppedEvent event = new ContainerStoppedEvent(
      ContainerId.newInstance(ApplicationAttemptId.newInstance(
          ApplicationId.newInstance(0, 1), 1), 1001), 100034566,
      ContainerExitStatus.SUCCESS, ApplicationAttemptId.newInstance(
          ApplicationId.newInstance(0, 1), 1));
  ContainerStoppedEvent deserializedEvent = (ContainerStoppedEvent)
      testProtoConversion(event);
  Assert.assertEquals(event.getContainerId(),
      deserializedEvent.getContainerId());
  Assert.assertEquals(event.getStoppedTime(),
      deserializedEvent.getStoppedTime());
  Assert.assertEquals(event.getApplicationAttemptId(),
      deserializedEvent.getApplicationAttemptId());
  logEvents(event, deserializedEvent);
}
 

@Test
public void testGetTaskShouldDie() throws InterruptedException, ExecutionException {
  ListeningExecutorService executor = null;
  try {
    ExecutorService rawExecutor = Executors.newFixedThreadPool(1);
    executor = MoreExecutors.listeningDecorator(rawExecutor);
    ApplicationId appId = ApplicationId.newInstance(10000, 1);
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
    ContainerId containerId = ContainerId.newInstance(appAttemptId, 1);

    TezTaskUmbilicalForTest umbilical = new TezTaskUmbilicalForTest();
    ContainerContext containerContext = new ContainerContext(containerId.toString());

    ContainerReporter containerReporter = new ContainerReporter(umbilical, containerContext, 100);
    ListenableFuture<ContainerTask> getTaskFuture = executor.submit(containerReporter);

    getTaskFuture.get();
    assertEquals(1, umbilical.getTaskInvocations);

  } finally {
    executor.shutdownNow();
  }
}
 

@Before
public void setup() {
  applicationId = ApplicationId.newInstance(9999l, 1);
  applicationAttemptId = ApplicationAttemptId.newInstance(applicationId, 1);
  tezDAGID = TezDAGID.getInstance(applicationId, random.nextInt());
  tezVertexID = TezVertexID.getInstance(tezDAGID, random.nextInt());
  tezTaskID = TezTaskID.getInstance(tezVertexID, random.nextInt());
  tezTaskAttemptID = TezTaskAttemptID.getInstance(tezTaskID, random.nextInt());
  dagPlan = DAGPlan.newBuilder().setName("DAGPlanMock").build();
  containerId = ContainerId.newInstance(applicationAttemptId, 111);
  nodeId = NodeId.newInstance("node", 13435);
}
 

@SuppressWarnings("deprecation")
public Container createContainer(Resource capability, Priority priority, String hostname, int port) {
  ContainerId containerId = ContainerId.newInstance(customAppAttemptId, nextId.getAndIncrement());
  NodeId nodeId = NodeId.newInstance(hostname, port);
  String nodeHttpAddress = "hostname:0";

  Container container = Container.newInstance(containerId,
      nodeId,
      nodeHttpAddress,
      capability,
      priority,
      null);

  return container;
}
 

private Container createContainer(int id, String host, Resource resource, Priority priority) {
  @SuppressWarnings("deprecation")
  ContainerId containerID = ContainerId.newInstance(
      ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1),
      id);
  NodeId nodeID = NodeId.newInstance(host, 0);
  Container container = Container.newInstance(containerID, nodeID, host + ":0",
      resource, priority, null);
  return container;
}
 

private Container createContainer(int id, String host, Resource resource, Priority priority) {
  ContainerId containerID = ContainerId.newInstance(
      ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1),
      id);
  NodeId nodeID = NodeId.newInstance(host, 0);
  Container container = Container.newInstance(containerID, nodeID, host + ":0",
      resource, priority, null);
  return container;
}
 

@SuppressWarnings("deprecation")
@Before
public void setup() {
  applicationId = ApplicationId.newInstance(9999l, 1);
  applicationAttemptId = ApplicationAttemptId.newInstance(applicationId, 1);
  tezDAGID = TezDAGID.getInstance(applicationId, random.nextInt());
  tezVertexID = TezVertexID.getInstance(tezDAGID, random.nextInt());
  tezTaskID = TezTaskID.getInstance(tezVertexID, random.nextInt());
  tezTaskAttemptID = TezTaskAttemptID.getInstance(tezTaskID, random.nextInt());
  dagPlan = DAGPlan.newBuilder().setName("DAGPlanMock").build();
  containerId = ContainerId.newInstance(applicationAttemptId, 111);
  nodeId = NodeId.newInstance("node", 13435);
}
 

public Container createContainer(Resource capability, Priority priority) {
  ApplicationAttemptId appAttemptId = appContext.getApplicationAttemptId();
  ContainerId containerId = ContainerId.newInstance(appAttemptId, nextId.getAndIncrement());
  NodeId nodeId = NodeId.newInstance("127.0.0.1", 0);
  String nodeHttpAddress = "127.0.0.1:0";

  Container container = Container.newInstance(containerId,
      nodeId,
      nodeHttpAddress,
      capability,
      priority,
      null);

  return container;
}
 

@SuppressWarnings("deprecation")
public Container createContainer(Resource capability, Priority priority) {
  ContainerId containerId = ContainerId.newInstance(customAppAttemptId, nextId.getAndIncrement());
  NodeId nodeId = NodeId.newInstance("127.0.0.1", 0);
  String nodeHttpAddress = "127.0.0.1:0";

  Container container = Container.newInstance(containerId,
      nodeId,
      nodeHttpAddress,
      capability,
      priority,
      null);

  return container;
}
 

@Test(timeout = 5000)
// Ensure the dag does not go into an error state if a attempt kill is
// received while STARTING
public void testLaunchFailedWhileKilling() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(
      appId, 0);
  TezDAGID dagID = TezDAGID.getInstance(appId, 1);
  TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
  TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);
  TezTaskAttemptID taskAttemptID = TezTaskAttemptID.getInstance(taskID, 0);

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

  Configuration taskConf = new Configuration();
  taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  taskConf.setBoolean("fs.file.impl.disable.cache", true);

  TaskLocationHint locationHint = new TaskLocationHint(
      new HashSet<String>(Arrays.asList(new String[] {"127.0.0.1"})), null);
  Resource resource = Resource.newInstance(1024, 1);

  AppContext mockAppContext = mock(AppContext.class);
  doReturn(new ClusterInfo()).when(mockAppContext).getClusterInfo();

  TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler,
      taListener, taskConf, new SystemClock(),
      mock(TaskHeartbeatHandler.class), mockAppContext, locationHint, false,
      resource, createFakeContainerContext(), false);

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

  taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, Priority
      .newInstance(3)));
  // At state STARTING.
  taImpl.handle(new TaskAttemptEventKillRequest(taskAttemptID, null));
  // At some KILLING state.
  taImpl.handle(new TaskAttemptEventKillRequest(taskAttemptID, null));
  // taImpl.handle(new TaskAttemptEventContainerTerminating(taskAttemptID,
  // null));
  assertFalse(eventHandler.internalError);
}
 

@Test(timeout = 5000)
// Ensure ContainerTerminating and ContainerTerminated is handled correctly by
// the TaskAttempt
public void testContainerTerminationWhileRunning() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(
      appId, 0);
  TezDAGID dagID = TezDAGID.getInstance(appId, 1);
  TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
  TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);
  TezTaskAttemptID taskAttemptID = TezTaskAttemptID.getInstance(taskID, 0);

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

  Configuration taskConf = new Configuration();
  taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  taskConf.setBoolean("fs.file.impl.disable.cache", true);

  TaskLocationHint locationHint = new TaskLocationHint(
      new HashSet<String>(Arrays.asList(new String[] {"127.0.0.1"})), null);
  Resource resource = Resource.newInstance(1024, 1);

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

  AppContext appCtx = mock(AppContext.class);
  AMContainerMap containers = new AMContainerMap(
      mock(ContainerHeartbeatHandler.class), mock(TaskAttemptListener.class),
      new ContainerContextMatcher(), appCtx);
  containers.addContainerIfNew(container);

  doReturn(new ClusterInfo()).when(appCtx).getClusterInfo();
  doReturn(containers).when(appCtx).getAllContainers();

  TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler,
      taListener, taskConf, new SystemClock(),
      mock(TaskHeartbeatHandler.class), appCtx, locationHint, false,
      resource, createFakeContainerContext(), false);

  ArgumentCaptor<Event> arg = ArgumentCaptor.forClass(Event.class);

  taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, null));
  // At state STARTING.
  taImpl.handle(new TaskAttemptEventStartedRemotely(taskAttemptID, contId,
      null));
  assertEquals("Task attempt is not in the RUNNING state", taImpl.getState(),
      TaskAttemptState.RUNNING);

  int expectedEventsAtRunning = 3;
  verify(eventHandler, times(expectedEventsAtRunning)).handle(arg.capture());

  taImpl.handle(new TaskAttemptEventContainerTerminating(taskAttemptID,
      "Terminating"));
  assertFalse(
      "InternalError occurred trying to handle TA_CONTAINER_TERMINATING",
      eventHandler.internalError);

  assertEquals("Task attempt is not in the  FAILED state", taImpl.getState(),
      TaskAttemptState.FAILED);

  assertEquals(1, taImpl.getDiagnostics().size());
  assertEquals("Terminating", taImpl.getDiagnostics().get(0));

  int expectedEvenstAfterTerminating = expectedEventsAtRunning + 3;
  arg = ArgumentCaptor.forClass(Event.class);
  verify(eventHandler, times(expectedEvenstAfterTerminating)).handle(arg.capture());

  verifyEventType(
      arg.getAllValues().subList(expectedEventsAtRunning,
          expectedEvenstAfterTerminating), TaskEventTAUpdate.class, 1);
  verifyEventType(
      arg.getAllValues().subList(expectedEventsAtRunning,
          expectedEvenstAfterTerminating), AMSchedulerEventTAEnded.class, 1);

  taImpl.handle(new TaskAttemptEventContainerTerminated(taskAttemptID,
      "Terminated"));
  int expectedEventAfterTerminated = expectedEvenstAfterTerminating + 0;
  arg = ArgumentCaptor.forClass(Event.class);
  verify(eventHandler, times(expectedEventAfterTerminated)).handle(arg.capture());

  assertEquals(2, taImpl.getDiagnostics().size());
  assertEquals("Terminated", taImpl.getDiagnostics().get(1));
}
 

private ContainerId createContainerId(ApplicationId appId) {
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
  ContainerId containerId = ContainerId.newInstance(appAttemptId, 1);
  return containerId;
}
 

@Test(timeout = 5000)
// Ensure ContainerTerminating and ContainerTerminated is handled correctly by
// the TaskAttempt
public void testContainerTerminatedAfterSuccess() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(
      appId, 0);
  TezDAGID dagID = TezDAGID.getInstance(appId, 1);
  TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
  TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);
  TezTaskAttemptID taskAttemptID = TezTaskAttemptID.getInstance(taskID, 0);

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

  Configuration taskConf = new Configuration();
  taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  taskConf.setBoolean("fs.file.impl.disable.cache", true);

  TaskLocationHint locationHint = new TaskLocationHint(
      new HashSet<String>(Arrays.asList(new String[] {"127.0.0.1"})), null);
  Resource resource = Resource.newInstance(1024, 1);

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

  AppContext appCtx = mock(AppContext.class);
  AMContainerMap containers = new AMContainerMap(
      mock(ContainerHeartbeatHandler.class), mock(TaskAttemptListener.class),
      new ContainerContextMatcher(), appCtx);
  containers.addContainerIfNew(container);

  doReturn(new ClusterInfo()).when(appCtx).getClusterInfo();
  doReturn(containers).when(appCtx).getAllContainers();

  TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler,
      taListener, taskConf, new SystemClock(),
      mock(TaskHeartbeatHandler.class), appCtx, locationHint, false,
      resource, createFakeContainerContext(), false);

  ArgumentCaptor<Event> arg = ArgumentCaptor.forClass(Event.class);

  taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, null));
  // At state STARTING.
  taImpl.handle(new TaskAttemptEventStartedRemotely(taskAttemptID, contId,
      null));
  assertEquals("Task attempt is not in the RUNNING state", taImpl.getState(),
      TaskAttemptState.RUNNING);

  int expectedEventsAtRunning = 3;
  verify(eventHandler, times(expectedEventsAtRunning)).handle(arg.capture());

  taImpl.handle(new TaskAttemptEvent(taskAttemptID, TaskAttemptEventType.TA_DONE));

  assertEquals("Task attempt is not in the  SUCCEEDED state", taImpl.getState(),
      TaskAttemptState.SUCCEEDED);

  assertEquals(0, taImpl.getDiagnostics().size());

  int expectedEvenstAfterTerminating = expectedEventsAtRunning + 3;
  arg = ArgumentCaptor.forClass(Event.class);
  verify(eventHandler, times(expectedEvenstAfterTerminating)).handle(arg.capture());

  verifyEventType(
      arg.getAllValues().subList(expectedEventsAtRunning,
          expectedEvenstAfterTerminating), TaskEventTAUpdate.class, 1);
  verifyEventType(
      arg.getAllValues().subList(expectedEventsAtRunning,
          expectedEvenstAfterTerminating), AMSchedulerEventTAEnded.class, 1);

  taImpl.handle(new TaskAttemptEventContainerTerminated(taskAttemptID,
      "Terminated"));
  int expectedEventAfterTerminated = expectedEvenstAfterTerminating + 0;
  arg = ArgumentCaptor.forClass(Event.class);
  verify(eventHandler, times(expectedEventAfterTerminated)).handle(arg.capture());

  // Verify that the diagnostic message included in the Terminated event is not
  // captured - TA already succeeded.
  assertEquals(0, taImpl.getDiagnostics().size());
}
 

@Test(timeout = 5000)
public void testCompletedAtSubmitted() throws ServicePluginException {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(
      appId, 0);
  TezDAGID dagID = TezDAGID.getInstance(appId, 1);
  TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
  TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);

  MockEventHandler eventHandler = spy(new MockEventHandler());
  TaskCommunicatorManagerInterface taListener = createMockTaskAttemptListener();

  Configuration taskConf = new Configuration();
  taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  taskConf.setBoolean("fs.file.impl.disable.cache", true);

  locationHint = TaskLocationHint.createTaskLocationHint(
      new HashSet<String>(Arrays.asList(new String[]{"127.0.0.1"})), null);
  Resource resource = Resource.newInstance(1024, 1);

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

  AMContainerMap containers = new AMContainerMap(
      mock(ContainerHeartbeatHandler.class), mock(TaskCommunicatorManagerInterface.class),
      new ContainerContextMatcher(), appCtx);
  containers.addContainerIfNew(container, 0, 0, 0);

  doReturn(new ClusterInfo()).when(appCtx).getClusterInfo();
  doReturn(containers).when(appCtx).getAllContainers();

  TaskHeartbeatHandler mockHeartbeatHandler = mock(TaskHeartbeatHandler.class);
  TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler,
      taListener, taskConf, new SystemClock(),
      mockHeartbeatHandler, appCtx, false,
      resource, createFakeContainerContext(), false);
  TezTaskAttemptID taskAttemptID = taImpl.getID();
  ArgumentCaptor<Event> arg = ArgumentCaptor.forClass(Event.class);

  taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, 0, 0));
  taImpl.handle(new TaskAttemptEventSubmitted(taskAttemptID, contId));
  assertEquals("Task attempt is not in the RUNNING state", taImpl.getState(),
      TaskAttemptState.STARTING);

  verify(mockHeartbeatHandler).register(taskAttemptID);

  int expectedEventsAtStarting = 4;
  verify(eventHandler, times(expectedEventsAtStarting)).handle(arg.capture());

  // Ensure status_updates are handled in the submitted state.
  taImpl.handle(new TaskAttemptEventStatusUpdate(taskAttemptID,
      new TaskStatusUpdateEvent(null, 0.1f, null, false)));

  taImpl.handle(new TaskAttemptEvent(taskAttemptID, TaskAttemptEventType.TA_DONE));

  assertEquals("Task attempt is not in the  SUCCEEDED state", taImpl.getState(),
      TaskAttemptState.SUCCEEDED);
  verify(mockHeartbeatHandler).unregister(taskAttemptID);
  assertEquals(0, taImpl.getDiagnostics().size());

  int expectedEvenstAfterTerminating = expectedEventsAtStarting + 3;
  arg = ArgumentCaptor.forClass(Event.class);
  verify(eventHandler, times(expectedEvenstAfterTerminating)).handle(arg.capture());


  Event e = verifyEventType(
      arg.getAllValues().subList(expectedEventsAtStarting,
          expectedEvenstAfterTerminating), TaskEventTASucceeded.class, 1);
  assertEquals(TaskEventType.T_ATTEMPT_SUCCEEDED, e.getType());
  verifyEventType(
      arg.getAllValues().subList(expectedEventsAtStarting,
          expectedEvenstAfterTerminating), AMSchedulerEventTAEnded.class, 1);
  verifyEventType(
      arg.getAllValues().subList(expectedEventsAtStarting,
          expectedEvenstAfterTerminating), DAGEventCounterUpdate.class, 1);
}
 

@Test(timeout = 5000)
// Ensure node failure on Successful Non-Leaf tasks cause them to be marked as KILLED
public void testNodeFailedNonLeafVertex() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(
      appId, 0);
  TezDAGID dagID = TezDAGID.getInstance(appId, 1);
  TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
  TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);
  TezTaskAttemptID taskAttemptID = TezTaskAttemptID.getInstance(taskID, 0);

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

  Configuration taskConf = new Configuration();
  taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  taskConf.setBoolean("fs.file.impl.disable.cache", true);

  TaskLocationHint locationHint = new TaskLocationHint(
      new HashSet<String>(Arrays.asList(new String[] {"127.0.0.1"})), null);
  Resource resource = Resource.newInstance(1024, 1);

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

  AppContext appCtx = mock(AppContext.class);
  AMContainerMap containers = new AMContainerMap(
      mock(ContainerHeartbeatHandler.class), mock(TaskAttemptListener.class),
      new ContainerContextMatcher(), appCtx);
  containers.addContainerIfNew(container);

  doReturn(new ClusterInfo()).when(appCtx).getClusterInfo();
  doReturn(containers).when(appCtx).getAllContainers();

  MockTaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler,
      taListener, taskConf, new SystemClock(),
      mock(TaskHeartbeatHandler.class), appCtx, locationHint, false,
      resource, createFakeContainerContext(), false);

  ArgumentCaptor<Event> arg = ArgumentCaptor.forClass(Event.class);

  taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, null));
  // At state STARTING.
  taImpl.handle(new TaskAttemptEventStartedRemotely(taskAttemptID, contId,
      null));
  assertEquals("Task attempt is not in the RUNNING state", TaskAttemptState.RUNNING,
      taImpl.getState());

  int expectedEventsAtRunning = 3;
  verify(eventHandler, times(expectedEventsAtRunning)).handle(arg.capture());

  taImpl.handle(new TaskAttemptEvent(taskAttemptID, TaskAttemptEventType.TA_DONE));

  assertEquals("Task attempt is not in the  SUCCEEDED state", TaskAttemptState.SUCCEEDED,
      taImpl.getState());

  assertEquals(0, taImpl.getDiagnostics().size());

  int expectedEvenstAfterTerminating = expectedEventsAtRunning + 3;
  arg = ArgumentCaptor.forClass(Event.class);
  verify(eventHandler, times(expectedEvenstAfterTerminating)).handle(arg.capture());

  verifyEventType(
      arg.getAllValues().subList(expectedEventsAtRunning,
          expectedEvenstAfterTerminating), TaskEventTAUpdate.class, 1);
  verifyEventType(
      arg.getAllValues().subList(expectedEventsAtRunning,
          expectedEvenstAfterTerminating), AMSchedulerEventTAEnded.class, 1);

  // Send out a Node Failure.
  taImpl.handle(new TaskAttemptEventNodeFailed(taskAttemptID, "NodeDecomissioned"));
  // Verify in KILLED state
  assertEquals("Task attempt is not in the  KILLED state", TaskAttemptState.KILLED,
      taImpl.getState());
  assertEquals(true, taImpl.inputFailedReported);
  // Verify one event to the Task informing it about FAILURE. No events to scheduler. Counter event.
  int expectedEventsNodeFailure = expectedEvenstAfterTerminating + 2;
  arg = ArgumentCaptor.forClass(Event.class);
  verify(eventHandler, times(expectedEventsNodeFailure)).handle(arg.capture());
  verifyEventType(
      arg.getAllValues().subList(expectedEvenstAfterTerminating,
          expectedEventsNodeFailure), TaskEventTAUpdate.class, 1);

  // Verify still in KILLED state
  assertEquals("Task attempt is not in the  KILLED state", TaskAttemptState.KILLED,
      taImpl.getState());
}
 

@Test(timeout = 5000)
// Ensure the dag does not go into an error state if a attempt kill is
// received while STARTING
public void testLaunchFailedWhileKilling() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(
      appId, 0);
  TezDAGID dagID = TezDAGID.getInstance(appId, 1);
  TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
  TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);
  TezTaskAttemptID taskAttemptID = TezTaskAttemptID.getInstance(taskID, 0);

  MockEventHandler eventHandler = new MockEventHandler();
  TaskCommunicatorManagerInterface taListener = createMockTaskAttemptListener();

  Configuration taskConf = new Configuration();
  taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  taskConf.setBoolean("fs.file.impl.disable.cache", true);

  locationHint = TaskLocationHint.createTaskLocationHint(
      new HashSet<String>(Arrays.asList(new String[]{"127.0.0.1"})), null);
  Resource resource = Resource.newInstance(1024, 1);

  AppContext mockAppContext = appCtx;
  doReturn(new ClusterInfo()).when(mockAppContext).getClusterInfo();

  TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler,
      taListener, taskConf, new SystemClock(),
      mock(TaskHeartbeatHandler.class), mockAppContext, false,
      resource, createFakeContainerContext(), false);

  NodeId nid = NodeId.newInstance("127.0.0.1", 0);
  @SuppressWarnings("deprecation")
  ContainerId contId = ContainerId.newInstance(appAttemptId, 3);
  Container container = mock(Container.class);
  when(container.getId()).thenReturn(contId);
  when(container.getNodeId()).thenReturn(nid);

  taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, 0, 0));
  // At state STARTING.
  taImpl.handle(new TaskAttemptEventKillRequest(taskAttemptID, null,
      TaskAttemptTerminationCause.TERMINATED_BY_CLIENT));
  assertEquals(TaskAttemptStateInternal.KILL_IN_PROGRESS, taImpl.getInternalState());
  taImpl.handle(new TaskAttemptEventTezEventUpdate(taImpl.getID(), Collections.EMPTY_LIST));
  assertFalse(
      "InternalError occurred trying to handle TA_TEZ_EVENT_UPDATE in KILL_IN_PROGRESS state",
      eventHandler.internalError);
  // At some KILLING state.
  taImpl.handle(new TaskAttemptEventKillRequest(taskAttemptID, null,
      TaskAttemptTerminationCause.TERMINATED_BY_CLIENT));
  // taImpl.handle(new TaskAttemptEventContainerTerminating(taskAttemptID,
  // null));
  assertFalse(eventHandler.internalError);
}
 

@SuppressWarnings("deprecation")
private ContainerId createContainerId(int id) {
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
  ContainerId containerId = ContainerId.newInstance(appAttemptId, id);
  return containerId;
}
 

@Test(timeout = 60000)
public void testTAFailureBasedOnRunningTasks() throws Exception {
  ApplicationId appId = ApplicationId.newInstance(1, 2);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(
      appId, 0);
  TezDAGID dagID = TezDAGID.getInstance(appId, 1);
  TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
  TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);

  MockEventHandler mockEh = new MockEventHandler();
  MockEventHandler eventHandler = spy(mockEh);
  TaskCommunicatorManagerInterface taListener = createMockTaskAttemptListener();

  Configuration taskConf = new Configuration();
  taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
  taskConf.setBoolean("fs.file.impl.disable.cache", true);

  locationHint = TaskLocationHint.createTaskLocationHint(
      new HashSet<String>(Arrays.asList(new String[]{"127.0.0.1"})), null);
  Resource resource = Resource.newInstance(1024, 1);

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

  AMContainerMap containers = new AMContainerMap(
      mock(ContainerHeartbeatHandler.class), mock(TaskCommunicatorManagerInterface.class),
      new ContainerContextMatcher(), appCtx);
  containers.addContainerIfNew(container, 0, 0, 0);

  doReturn(new ClusterInfo()).when(appCtx).getClusterInfo();
  doReturn(containers).when(appCtx).getAllContainers();
  HistoryEventHandler mockHistHandler = mock(HistoryEventHandler.class);
  doReturn(mockHistHandler).when(appCtx).getHistoryHandler();
  DAGImpl mockDAG = mock(DAGImpl.class);

  TaskHeartbeatHandler mockHeartbeatHandler = mock(TaskHeartbeatHandler.class);
  MockTaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler,
      taListener, taskConf, new SystemClock(),
      mockHeartbeatHandler, appCtx, false,
      resource, createFakeContainerContext(), false);
  TezTaskAttemptID taskAttemptID = taImpl.getID();

  taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, 0, 0));
  taImpl.handle(new TaskAttemptEventSubmitted(taskAttemptID, contId));
  taImpl.handle(new TaskAttemptEventStartedRemotely(taskAttemptID));
  verify(mockHeartbeatHandler).register(taskAttemptID);
  taImpl.handle(new TaskAttemptEvent(taskAttemptID,
      TaskAttemptEventType.TA_DONE));
  assertEquals("Task attempt is not in succeeded state", taImpl.getState(),
      TaskAttemptState.SUCCEEDED);
  verify(mockHeartbeatHandler).unregister(taskAttemptID);

  int expectedEventsTillSucceeded = 8;
  ArgumentCaptor<Event> arg = ArgumentCaptor.forClass(Event.class);
  ArgumentCaptor<DAGHistoryEvent> histArg = ArgumentCaptor.forClass(DAGHistoryEvent.class);
  verify(eventHandler, times(expectedEventsTillSucceeded)).handle(arg.capture());
  verify(mockHistHandler, times(2)).handle(histArg.capture()); // start and finish
  DAGHistoryEvent histEvent = histArg.getValue();
  TaskAttemptFinishedEvent finishEvent = (TaskAttemptFinishedEvent)histEvent.getHistoryEvent();
  long finishTime = finishEvent.getFinishTime();
  verifyEventType(arg.getAllValues(), TaskEventTAUpdate.class, 2);

  InputReadErrorEvent mockReEvent = InputReadErrorEvent.create("", 0, 1);
  EventMetaData mockMeta = mock(EventMetaData.class);
  TezTaskAttemptID mockDestId1 = mock(TezTaskAttemptID.class);
  when(mockMeta.getTaskAttemptID()).thenReturn(mockDestId1);
  TezTaskID destTaskID = mock(TezTaskID.class);
  TezVertexID destVertexID = mock(TezVertexID.class);
  when(mockDestId1.getTaskID()).thenReturn(destTaskID);
  when(destTaskID.getVertexID()).thenReturn(destVertexID);
  Vertex destVertex = mock(VertexImpl.class);
  when(destVertex.getRunningTasks()).thenReturn(5);
  when(mockDAG.getVertex(destVertexID)).thenReturn(destVertex);
  when(appCtx.getCurrentDAG()).thenReturn(mockDAG);
  TezEvent tzEvent = new TezEvent(mockReEvent, mockMeta);
  taImpl.handle(new TaskAttemptEventOutputFailed(taskAttemptID, tzEvent, 11));

  // failure threshold is met due to running tasks. state is FAILED
  assertEquals("Task attempt is not in FAILED state", taImpl.getState(),
      TaskAttemptState.FAILED);
}
 

@SuppressWarnings("deprecation")
public static ContainerId createContainerId(ApplicationId appId) {
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
  ContainerId containerId = ContainerId.newInstance(appAttemptId, 1);
  return containerId;
}