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

@Test
public void testNodeId() {
  NodeId nodeId1 = NodeId.newInstance("10.18.52.124", 8041);
  NodeId nodeId2 = NodeId.newInstance("10.18.52.125", 8038);
  NodeId nodeId3 = NodeId.newInstance("10.18.52.124", 8041);
  NodeId nodeId4 = NodeId.newInstance("10.18.52.124", 8039);

  Assert.assertTrue(nodeId1.equals(nodeId3));
  Assert.assertFalse(nodeId1.equals(nodeId2));
  Assert.assertFalse(nodeId3.equals(nodeId4));

  Assert.assertTrue(nodeId1.compareTo(nodeId3) == 0);
  Assert.assertTrue(nodeId1.compareTo(nodeId2) < 0);
  Assert.assertTrue(nodeId3.compareTo(nodeId4) > 0);

  Assert.assertTrue(nodeId1.hashCode() == nodeId3.hashCode());
  Assert.assertFalse(nodeId1.hashCode() == nodeId2.hashCode());
  Assert.assertFalse(nodeId3.hashCode() == nodeId4.hashCode());

  Assert.assertEquals("10.18.52.124:8041", nodeId1.toString());
}
 

@Test(timeout = 15000)
public void testResourceTrackerOnHA() throws Exception {
  NodeId nodeId = NodeId.newInstance("localhost", 0);
  Resource resource = Resource.newInstance(2048, 4);

  // make sure registerNodeManager works when failover happens
  RegisterNodeManagerRequest request =
      RegisterNodeManagerRequest.newInstance(nodeId, 0, resource,
          YarnVersionInfo.getVersion(), null, null);
  resourceTracker.registerNodeManager(request);
  Assert.assertTrue(waitForNodeManagerToConnect(10000, nodeId));

  // restart the failover thread, and make sure nodeHeartbeat works
  failoverThread = createAndStartFailoverThread();
  NodeStatus status =
      NodeStatus.newInstance(NodeId.newInstance("localhost", 0), 0, null,
          null, null);
  NodeHeartbeatRequest request2 =
      NodeHeartbeatRequest.newInstance(status, null, null);
  resourceTracker.nodeHeartbeat(request2);
}
 

private void testTaskAttemptStartedEvent() throws Exception {
  TaskAttemptStartedEvent event = new TaskAttemptStartedEvent(
      TezTaskAttemptID.getInstance(TezTaskID.getInstance(TezVertexID.getInstance(
          TezDAGID.getInstance(ApplicationId.newInstance(0, 1), 1), 111), 1), 1),
      "vertex1", 10009l, ContainerId.newInstance(
      ApplicationAttemptId.newInstance(
          ApplicationId.newInstance(0, 1), 1), 1001), NodeId.newInstance(
      "host1", 19999), "inProgress", "Completed");
  TaskAttemptStartedEvent deserializedEvent = (TaskAttemptStartedEvent)
      testProtoConversion(event);
  Assert.assertEquals(event.getTaskAttemptID(),
      deserializedEvent.getTaskAttemptID());
  Assert.assertEquals(event.getContainerId(),
      deserializedEvent.getContainerId());
  Assert.assertEquals(event.getNodeId(),
      deserializedEvent.getNodeId());
  Assert.assertEquals(event.getStartTime(),
      deserializedEvent.getStartTime());
  logEvents(event, deserializedEvent);
}
 

public NodeId toNodeId(String str) {
  if (str.contains(":")) {
    int idx = str.indexOf(':');
    NodeId id =
        NodeId.newInstance(str.substring(0, idx),
            Integer.valueOf(str.substring(idx + 1)));
    return id;
  } else {
    return NodeId.newInstance(str, CommonNodeLabelsManager.WILDCARD_PORT);
  }
}
 

Container createTestingContainerWithResource(Resource resource) {
	final ContainerId containerId = ContainerId.newInstance(
		ApplicationAttemptId.newInstance(
			ApplicationId.newInstance(System.currentTimeMillis(), 1),
			1),
		containerIdx++);
	final NodeId nodeId = NodeId.newInstance("container", 1234);
	return new TestingContainer(containerId, nodeId, resource, Priority.UNDEFINED);
}
 

@SuppressWarnings("unchecked")
@Override
public void handle(ContainerAllocatorEvent event) {
  if (event.getType() == ContainerAllocator.EventType.CONTAINER_REQ) {
    LOG.info("Processing the event " + event.toString());
    // Assign the same container ID as the AM
    ContainerId cID =
        ContainerId.newContainerId(getContext().getApplicationAttemptId(),
          this.containerId.getContainerId());
    Container container = recordFactory.newRecordInstance(Container.class);
    container.setId(cID);
    NodeId nodeId = NodeId.newInstance(this.nmHost, this.nmPort);
    container.setNodeId(nodeId);
    container.setContainerToken(null);
    container.setNodeHttpAddress(this.nmHost + ":" + this.nmHttpPort);
    // send the container-assigned event to task attempt

    if (event.getAttemptID().getTaskId().getTaskType() == TaskType.MAP) {
      JobCounterUpdateEvent jce =
          new JobCounterUpdateEvent(event.getAttemptID().getTaskId()
              .getJobId());
      // TODO Setting OTHER_LOCAL_MAP for now.
      jce.addCounterUpdate(JobCounter.OTHER_LOCAL_MAPS, 1);
      eventHandler.handle(jce);
    }
    eventHandler.handle(new TaskAttemptContainerAssignedEvent(
        event.getAttemptID(), container, applicationACLs));
  }
}
 

private NodeId buildNodeId(InetSocketAddress connectAddress,
    String hostOverride) {
  if (hostOverride != null) {
    connectAddress = NetUtils.getConnectAddress(
        new InetSocketAddress(hostOverride, connectAddress.getPort()));
  }
  return NodeId.newInstance(
      connectAddress.getAddress().getCanonicalHostName(),
      connectAddress.getPort());
}
 

@SuppressWarnings("deprecation")
@Test
public void testContainersCleanupForLastAttempt() {
  // create a failed attempt.
  applicationAttempt =
      new RMAppAttemptImpl(applicationAttempt.getAppAttemptId(), spyRMContext,
        scheduler, masterService, submissionContext, new Configuration(),
        true, BuilderUtils.newResourceRequest(
            RMAppAttemptImpl.AM_CONTAINER_PRIORITY, ResourceRequest.ANY,
            submissionContext.getResource(), 1));
  when(submissionContext.getKeepContainersAcrossApplicationAttempts())
    .thenReturn(true);
  when(submissionContext.getMaxAppAttempts()).thenReturn(1);
  Container amContainer = allocateApplicationAttempt();
  launchApplicationAttempt(amContainer);
  runApplicationAttempt(amContainer, "host", 8042, "oldtrackingurl", false);
  ContainerStatus cs1 =
      ContainerStatus.newInstance(amContainer.getId(),
        ContainerState.COMPLETE, "some error", 123);
  ApplicationAttemptId appAttemptId = applicationAttempt.getAppAttemptId();
  NodeId anyNodeId = NodeId.newInstance("host", 1234);
  applicationAttempt.handle(new RMAppAttemptContainerFinishedEvent(
    appAttemptId, cs1, anyNodeId));
  assertEquals(YarnApplicationAttemptState.RUNNING,
      applicationAttempt.createApplicationAttemptState());
  sendAttemptUpdateSavedEvent(applicationAttempt);
  assertEquals(RMAppAttemptState.FAILED,
    applicationAttempt.getAppAttemptState());
  assertFalse(transferStateFromPreviousAttempt);
  verifyApplicationAttemptFinished(RMAppAttemptState.FAILED);
}
 

@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);
}
 

@Test
public void
    testFinalSavingToFinishedWithContainerFinished() {
  Container amContainer = allocateApplicationAttempt();
  launchApplicationAttempt(amContainer);
  runApplicationAttempt(amContainer, "host", 8042, "oldtrackingurl", false);
  FinalApplicationStatus finalStatus = FinalApplicationStatus.SUCCEEDED;
  String trackingUrl = "mytrackingurl";
  String diagnostics = "Successful";
  applicationAttempt.handle(new RMAppAttemptUnregistrationEvent(
    applicationAttempt.getAppAttemptId(), trackingUrl, finalStatus,
    diagnostics));
  assertEquals(RMAppAttemptState.FINAL_SAVING,
    applicationAttempt.getAppAttemptState());
  assertEquals(YarnApplicationAttemptState.RUNNING,
      applicationAttempt.createApplicationAttemptState());
  // Container_finished event comes before Attempt_Saved event.
  NodeId anyNodeId = NodeId.newInstance("host", 1234);
  applicationAttempt.handle(new RMAppAttemptContainerFinishedEvent(
    applicationAttempt.getAppAttemptId(), BuilderUtils.newContainerStatus(
      amContainer.getId(), ContainerState.COMPLETE, "", 0), anyNodeId));
  assertEquals(RMAppAttemptState.FINAL_SAVING,
    applicationAttempt.getAppAttemptState());
  // send attempt_saved
  sendAttemptUpdateSavedEvent(applicationAttempt);
  testAppAttemptFinishedState(amContainer, finalStatus, trackingUrl,
    diagnostics, 0, false);
}
 

@Override
protected void serviceStart() throws Exception {
  thread = new Thread(new Runnable() {
    @Override
    public void run() {
      ContainerAllocatorEvent event = null;
      while (!Thread.currentThread().isInterrupted()) {
        try {
          if (concurrentRunningTasks < maxConcurrentRunningTasks) {
            event = eventQueue.take();
            ContainerId cId =
                ContainerId.newContainerId(getContext()
                  .getApplicationAttemptId(), containerCount++);

            //System.out.println("Allocating " + containerCount);
            
            Container container = 
                recordFactory.newRecordInstance(Container.class);
            container.setId(cId);
            NodeId nodeId = NodeId.newInstance("dummy", 1234);
            container.setNodeId(nodeId);
            container.setContainerToken(null);
            container.setNodeHttpAddress("localhost:8042");
            getContext().getEventHandler()
                .handle(
                new TaskAttemptContainerAssignedEvent(event
                    .getAttemptID(), container, null));
            concurrentRunningTasks++;
          } else {
            Thread.sleep(1000);
          }
        } catch (InterruptedException e) {
          System.out.println("Returning, interrupted");
          return;
        }
      }
    }
  });
  thread.start();
  super.serviceStart();
}
 

@Test
public void testContainerKillAfterAssigned() 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)));
  assertEquals("Task attempt is not in assinged state",
      taImpl.getInternalState(), TaskAttemptStateInternal.ASSIGNED);
  taImpl.handle(new TaskAttemptEvent(attemptId,
      TaskAttemptEventType.TA_KILL));
  assertEquals("Task should be in KILLED state",
      TaskAttemptStateInternal.KILL_CONTAINER_CLEANUP,
      taImpl.getInternalState());
}
 

@SuppressWarnings("deprecation") // ContainerId
public WrappedContainer(boolean shouldProfile, String profileString, int cIdInt) {
  applicationID = ApplicationId.newInstance(rmIdentifier, 1);
  appAttemptID = ApplicationAttemptId.newInstance(applicationID, 1);
  containerID = ContainerId.newInstance(appAttemptID, cIdInt);
  nodeID = NodeId.newInstance("host", 12500);
  nodeHttpAddress = "host:12501";
  resource = Resource.newInstance(1024, 1);
  priority = Priority.newInstance(1);
  container = Container.newInstance(containerID, nodeID,
      nodeHttpAddress, resource, priority, null);

  chh = mock(ContainerHeartbeatHandler.class);

  tal = mock(TaskCommunicatorManagerInterface.class);
  TaskCommunicator taskComm = mock(TaskCommunicator.class);
  try {
    doReturn(new InetSocketAddress("localhost", 0)).when(taskComm).getAddress();
  } catch (ServicePluginException e) {
    throw new RuntimeException(e);
  }
  doReturn(new TaskCommunicatorWrapper(taskComm)).when(tal).getTaskCommunicator(0);

  dagID = TezDAGID.getInstance(applicationID, 1);
  vertexID = TezVertexID.getInstance(dagID, 1);
  taskID = TezTaskID.getInstance(vertexID, 1);
  taskAttemptID = TezTaskAttemptID.getInstance(taskID, 1);
  
  eventHandler = mock(EventHandler.class);
  historyEventHandler = mock(HistoryEventHandler.class);

  Configuration conf = new Configuration(false);
  appContext = mock(AppContext.class);
  doReturn(new HashMap<ApplicationAccessType, String>()).when(appContext)
  .getApplicationACLs();
  doReturn(eventHandler).when(appContext).getEventHandler();
  doReturn(appAttemptID).when(appContext).getApplicationAttemptId();
  doReturn(applicationID).when(appContext).getApplicationID();
  doReturn(new SystemClock()).when(appContext).getClock();
  doReturn(historyEventHandler).when(appContext).getHistoryHandler();
  doReturn(conf).when(appContext).getAMConf();
  mockDAGID();

  taskSpec = mock(TaskSpec.class);
  doReturn(taskAttemptID).when(taskSpec).getTaskAttemptID();

  amContainer = new AMContainerImpl(container, chh, tal,
      new ContainerContextMatcher(), appContext, 0, 0, 0, conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID,
      TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
}
 

@Test
public void testFailedToFailed() {
  // create a failed attempt.
  when(submissionContext.getKeepContainersAcrossApplicationAttempts())
    .thenReturn(true);
  Container amContainer = allocateApplicationAttempt();
  launchApplicationAttempt(amContainer);
  runApplicationAttempt(amContainer, "host", 8042, "oldtrackingurl", false);
  ContainerStatus cs1 =
      ContainerStatus.newInstance(amContainer.getId(),
        ContainerState.COMPLETE, "some error", 123);
  ApplicationAttemptId appAttemptId = applicationAttempt.getAppAttemptId();
  NodeId anyNodeId = NodeId.newInstance("host", 1234);
  applicationAttempt.handle(new RMAppAttemptContainerFinishedEvent(
    appAttemptId, cs1, anyNodeId));
  assertEquals(YarnApplicationAttemptState.RUNNING,
      applicationAttempt.createApplicationAttemptState());
  sendAttemptUpdateSavedEvent(applicationAttempt);
  assertEquals(RMAppAttemptState.FAILED,
    applicationAttempt.getAppAttemptState());
  // should not kill containers when attempt fails.
  assertTrue(transferStateFromPreviousAttempt);
  verifyApplicationAttemptFinished(RMAppAttemptState.FAILED);

  // failed attempt captured the container finished event.
  assertEquals(0, applicationAttempt.getJustFinishedContainers().size());
  ContainerStatus cs2 =
      ContainerStatus.newInstance(ContainerId.newContainerId(appAttemptId, 2),
        ContainerState.COMPLETE, "", 0);
  applicationAttempt.handle(new RMAppAttemptContainerFinishedEvent(
    appAttemptId, cs2, anyNodeId));
  assertEquals(1, applicationAttempt.getJustFinishedContainers().size());
  boolean found = false;
  for (ContainerStatus containerStatus:applicationAttempt
      .getJustFinishedContainers()) {
    if (cs2.getContainerId().equals(containerStatus.getContainerId())) {
      found = true;
    }
  }
  assertTrue(found);
}
 

private void test(String rpcClass) throws Exception {
  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.IPC_RPC_IMPL, rpcClass);
  YarnRPC rpc = YarnRPC.create(conf);
  String bindAddr = "localhost:0";
  InetSocketAddress addr = NetUtils.createSocketAddr(bindAddr);
  Server server = rpc.getServer(ContainerManagementProtocol.class, 
          new DummyContainerManager(), addr, conf, null, 1);
  server.start();
  RPC.setProtocolEngine(conf, ContainerManagementProtocolPB.class, ProtobufRpcEngine.class);
  ContainerManagementProtocol proxy = (ContainerManagementProtocol) 
      rpc.getProxy(ContainerManagementProtocol.class, 
          NetUtils.getConnectAddress(server), conf);
  ContainerLaunchContext containerLaunchContext = 
      recordFactory.newRecordInstance(ContainerLaunchContext.class);

  ApplicationId applicationId = ApplicationId.newInstance(0, 0);
  ApplicationAttemptId applicationAttemptId =
      ApplicationAttemptId.newInstance(applicationId, 0);
  ContainerId containerId =
      ContainerId.newContainerId(applicationAttemptId, 100);
  NodeId nodeId = NodeId.newInstance("localhost", 1234);
  Resource resource = Resource.newInstance(1234, 2);
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier(containerId, "localhost", "user",
        resource, System.currentTimeMillis() + 10000, 42, 42,
        Priority.newInstance(0), 0);
  Token containerToken = newContainerToken(nodeId, "password".getBytes(),
        containerTokenIdentifier);

  StartContainerRequest scRequest =
      StartContainerRequest.newInstance(containerLaunchContext,
        containerToken);
  List<StartContainerRequest> list = new ArrayList<StartContainerRequest>();
  list.add(scRequest);
  StartContainersRequest allRequests =
      StartContainersRequest.newInstance(list);
  proxy.startContainers(allRequests);

  List<ContainerId> containerIds = new ArrayList<ContainerId>();
  containerIds.add(containerId);
  GetContainerStatusesRequest gcsRequest =
      GetContainerStatusesRequest.newInstance(containerIds);
  GetContainerStatusesResponse response =
      proxy.getContainerStatuses(gcsRequest);
  List<ContainerStatus> statuses = response.getContainerStatuses();

  //test remote exception
  boolean exception = false;
  try {
    StopContainersRequest stopRequest =
        recordFactory.newRecordInstance(StopContainersRequest.class);
    stopRequest.setContainerIds(containerIds);
    proxy.stopContainers(stopRequest);
    } catch (YarnException e) {
    exception = true;
    Assert.assertTrue(e.getMessage().contains(EXCEPTION_MSG));
    Assert.assertTrue(e.getMessage().contains(EXCEPTION_CAUSE));
    System.out.println("Test Exception is " + e.getMessage());
  } catch (Exception ex) {
    ex.printStackTrace();
  }
  Assert.assertTrue(exception);
  
  server.stop();
  Assert.assertNotNull(statuses.get(0));
  Assert.assertEquals(ContainerState.RUNNING, statuses.get(0).getState());
}
 

public static NodeId newNodeID(String host, int port) {
  return NodeId.newInstance(host, port);
}
 

private void test(String rpcClass) throws Exception {
  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.IPC_RPC_IMPL, rpcClass);
  YarnRPC rpc = YarnRPC.create(conf);
  String bindAddr = "localhost:0";
  InetSocketAddress addr = NetUtils.createSocketAddr(bindAddr);
  Server server = rpc.getServer(ContainerManagementProtocol.class, 
          new DummyContainerManager(), addr, conf, null, 1);
  server.start();
  RPC.setProtocolEngine(conf, ContainerManagementProtocolPB.class, ProtobufRpcEngine.class);
  ContainerManagementProtocol proxy = (ContainerManagementProtocol) 
      rpc.getProxy(ContainerManagementProtocol.class, 
          NetUtils.getConnectAddress(server), conf);
  ContainerLaunchContext containerLaunchContext = 
      recordFactory.newRecordInstance(ContainerLaunchContext.class);

  ApplicationId applicationId = ApplicationId.newInstance(0, 0);
  ApplicationAttemptId applicationAttemptId =
      ApplicationAttemptId.newInstance(applicationId, 0);
  ContainerId containerId =
      ContainerId.newContainerId(applicationAttemptId, 100);
  NodeId nodeId = NodeId.newInstance("localhost", 1234);
  Resource resource = Resource.newInstance(1234, 2);
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier(containerId, "localhost", "user",
        resource, System.currentTimeMillis() + 10000, 42, 42,
        Priority.newInstance(0), 0);
  Token containerToken = newContainerToken(nodeId, "password".getBytes(),
        containerTokenIdentifier);

  StartContainerRequest scRequest =
      StartContainerRequest.newInstance(containerLaunchContext,
        containerToken);
  List<StartContainerRequest> list = new ArrayList<StartContainerRequest>();
  list.add(scRequest);
  StartContainersRequest allRequests =
      StartContainersRequest.newInstance(list);
  proxy.startContainers(allRequests);

  List<ContainerId> containerIds = new ArrayList<ContainerId>();
  containerIds.add(containerId);
  GetContainerStatusesRequest gcsRequest =
      GetContainerStatusesRequest.newInstance(containerIds);
  GetContainerStatusesResponse response =
      proxy.getContainerStatuses(gcsRequest);
  List<ContainerStatus> statuses = response.getContainerStatuses();

  //test remote exception
  boolean exception = false;
  try {
    StopContainersRequest stopRequest =
        recordFactory.newRecordInstance(StopContainersRequest.class);
    stopRequest.setContainerIds(containerIds);
    proxy.stopContainers(stopRequest);
    } catch (YarnException e) {
    exception = true;
    Assert.assertTrue(e.getMessage().contains(EXCEPTION_MSG));
    Assert.assertTrue(e.getMessage().contains(EXCEPTION_CAUSE));
    System.out.println("Test Exception is " + e.getMessage());
  } catch (Exception ex) {
    ex.printStackTrace();
  }
  Assert.assertTrue(exception);
  
  server.stop();
  Assert.assertNotNull(statuses.get(0));
  Assert.assertEquals(ContainerState.RUNNING, statuses.get(0).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);
}
 

@Test
public void testAppDiognosticEventOnNewTask() 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 TaskAttemptDiagnosticsUpdateEvent(attemptId,
      "Task got killed"));
  assertFalse(
      "InternalError occurred trying to handle TA_DIAGNOSTICS_UPDATE on assigned task",
      eventHandler.internalError);
}
 

private static Container createMockContainer() {
  Container container = mock(Container.class);
  NodeId nodeId = NodeId.newInstance("foo.bar.org", 1234);
  when(container.getNodeId()).thenReturn(nodeId);
  return container;
}