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

@Test
public void testGetApplicationAttemptReport() throws YarnException,
    IOException {
  ClientRMService rmService = createRMService();
  RecordFactory recordFactory = RecordFactoryProvider.getRecordFactory(null);
  GetApplicationAttemptReportRequest request = recordFactory
      .newRecordInstance(GetApplicationAttemptReportRequest.class);
  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(
      ApplicationId.newInstance(123456, 1), 1);
  request.setApplicationAttemptId(attemptId);

  try {
    GetApplicationAttemptReportResponse response = rmService
        .getApplicationAttemptReport(request);
    Assert.assertEquals(attemptId, response.getApplicationAttemptReport()
        .getApplicationAttemptId());
  } catch (ApplicationNotFoundException ex) {
    Assert.fail(ex.getMessage());
  }
}
 

@Test
public void testMassiveWriteContainerHistory() throws IOException {
  long mb = 1024 * 1024;
  Runtime runtime = Runtime.getRuntime();
  long usedMemoryBefore = (runtime.totalMemory() - runtime.freeMemory()) / mb;
  int numContainers = 100000;
  ApplicationId appId = ApplicationId.newInstance(0, 1);
  ApplicationAttemptId appAttemptId =
      ApplicationAttemptId.newInstance(appId, 1);
  for (int i = 1; i <= numContainers; ++i) {
    ContainerId containerId = ContainerId.newContainerId(appAttemptId, i);
    writeContainerStartData(containerId);
    writeContainerFinishData(containerId);
  }
  long usedMemoryAfter = (runtime.totalMemory() - runtime.freeMemory()) / mb;
  Assert.assertTrue((usedMemoryAfter - usedMemoryBefore) < 400);
}
 

private void verifyJobReport(JobReport jobReport, JobId jobId) {
  List<AMInfo> amInfos = jobReport.getAMInfos();
  Assert.assertEquals(1, amInfos.size());
  AMInfo amInfo = amInfos.get(0);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(jobId.getAppId(), 1);
  ContainerId amContainerId = ContainerId.newContainerId(appAttemptId, 1);
  Assert.assertEquals(appAttemptId, amInfo.getAppAttemptId());
  Assert.assertEquals(amContainerId, amInfo.getContainerId());
  Assert.assertTrue(jobReport.getSubmitTime() > 0);
  Assert.assertTrue(jobReport.getStartTime() > 0
      && jobReport.getStartTime() >= jobReport.getSubmitTime());
  Assert.assertTrue(jobReport.getFinishTime() > 0
      && jobReport.getFinishTime() >= jobReport.getStartTime());
}
 

@Test
public void testGetContainerReport() throws Exception {
  ApplicationCLI cli = createAndGetAppCLI();
  ApplicationId applicationId = ApplicationId.newInstance(1234, 5);
  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(
      applicationId, 1);
  ContainerId containerId = ContainerId.newContainerId(attemptId, 1);
  ContainerReport container = ContainerReport.newInstance(containerId, null,
      NodeId.newInstance("host", 1234), Priority.UNDEFINED, 1234, 5678,
      "diagnosticInfo", "logURL", 0, ContainerState.COMPLETE,
      "http://" + NodeId.newInstance("host", 2345).toString());
  when(client.getContainerReport(any(ContainerId.class))).thenReturn(
      container);
  int result = cli.run(new String[] { "container", "-status",
      containerId.toString() });
  assertEquals(0, result);
  verify(client).getContainerReport(containerId);
  ByteArrayOutputStream baos = new ByteArrayOutputStream();
  PrintWriter pw = new PrintWriter(baos);
  pw.println("Container Report : ");
  pw.println("\tContainer-Id : container_1234_0005_01_000001");
  pw.println("\tStart-Time : 1234");
  pw.println("\tFinish-Time : 5678");
  pw.println("\tState : COMPLETE");
  pw.println("\tLOG-URL : logURL");
  pw.println("\tHost : host:1234");
  pw.println("\tNodeHttpAddress : http://host:2345");
  pw.println("\tDiagnostics : diagnosticInfo");
  pw.close();
  String appReportStr = baos.toString("UTF-8");
  Assert.assertEquals(appReportStr, sysOutStream.toString());
  verify(sysOut, times(1)).println(isA(String.class));
}
 

@Test
public void testNodeUpdateBeforeAppAttemptInit() throws Exception {
  FifoScheduler scheduler = new FifoScheduler();
  MockRM rm = new MockRM(conf);
  scheduler.setRMContext(rm.getRMContext());
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, rm.getRMContext());

  RMNode node = MockNodes.newNodeInfo(1,
          Resources.createResource(1024, 4), 1, "127.0.0.1");
  scheduler.handle(new NodeAddedSchedulerEvent(node));

  ApplicationId appId = ApplicationId.newInstance(0, 1);
  scheduler.addApplication(appId, "queue1", "user1", false);

  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  try {
    scheduler.handle(updateEvent);
  } catch (NullPointerException e) {
      Assert.fail();
  }

  ApplicationAttemptId attId = ApplicationAttemptId.newInstance(appId, 1);
  scheduler.addApplicationAttempt(attId, false, false);

  rm.stop();
}
 

private void testWriteHistoryData(
    int num, boolean missingContainer, boolean missingApplicationAttempt)
        throws IOException {
  // write application history data
  for (int i = 1; i <= num; ++i) {
    ApplicationId appId = ApplicationId.newInstance(0, i);
    writeApplicationStartData(appId);

    // write application attempt history data
    for (int j = 1; j <= num; ++j) {
      ApplicationAttemptId appAttemptId =
          ApplicationAttemptId.newInstance(appId, j);
      writeApplicationAttemptStartData(appAttemptId);

      if (missingApplicationAttempt && j == num) {
        continue;
      }
      // write container history data
      for (int k = 1; k <= num; ++k) {
        ContainerId containerId = ContainerId.newContainerId(appAttemptId, k);
        writeContainerStartData(containerId);
        if (missingContainer && k == num) {
          continue;
        }
        writeContainerFinishData(containerId);
      }
      writeApplicationAttemptFinishData(appAttemptId);
    }
    writeApplicationFinishData(appId);
  }
}
 

private void testAMLaunchedEvent() throws Exception {
  AMLaunchedEvent event = new AMLaunchedEvent(
      ApplicationAttemptId.newInstance(
          ApplicationId.newInstance(0, 1), 1),
      100, 100, null);
  AMLaunchedEvent deserializedEvent = (AMLaunchedEvent)
      testProtoConversion(event);
  Assert.assertEquals(event.getApplicationAttemptId(),
      deserializedEvent.getApplicationAttemptId());
  Assert.assertEquals(event.getAppSubmitTime(),
      deserializedEvent.getAppSubmitTime());
  Assert.assertEquals(event.getLaunchTime(),
      deserializedEvent.getLaunchTime());
  logEvents(event, deserializedEvent);
}
 

@Test(timeout = 10000)
public void testPostExecuteAfterReacquisition() throws Exception {
  Assume.assumeTrue(shouldRun());
  // make up some bogus container ID
  ApplicationId appId = ApplicationId.newInstance(12345, 67890);
  ApplicationAttemptId attemptId =
      ApplicationAttemptId.newInstance(appId, 54321);
  ContainerId cid = ContainerId.newContainerId(attemptId, 9876);

  Configuration conf = new YarnConfiguration();
  conf.setClass(YarnConfiguration.NM_LINUX_CONTAINER_RESOURCES_HANDLER,
    TestResourceHandler.class, LCEResourcesHandler.class);
  LinuxContainerExecutor lce = new LinuxContainerExecutor();
  lce.setConf(conf);
  try {
    lce.init();
  } catch (IOException e) {
    // expected if LCE isn't setup right, but not necessary for this test
  }

  Container container = mock(Container.class);
  ContainerLaunchContext context = mock(ContainerLaunchContext.class);
  HashMap<String, String> env = new HashMap<>();

  when(container.getLaunchContext()).thenReturn(context);
  when(context.getEnvironment()).thenReturn(env);

  lce.reacquireContainer(new ContainerReacquisitionContext.Builder()
      .setContainer(container)
      .setUser("foouser")
      .setContainerId(cid)
      .build());
  assertTrue("postExec not called after reacquisition",
      TestResourceHandler.postExecContainers.contains(cid));
}
 

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

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

@Override
protected UserGroupInformation getRemoteUgi() throws YarnException {
  ApplicationId appId = ApplicationId.newInstance(0, 0);
  ApplicationAttemptId appAttemptId =
      ApplicationAttemptId.newInstance(appId, 1);
  UserGroupInformation ugi =
      UserGroupInformation.createRemoteUser(appAttemptId.toString());
  ugi.addTokenIdentifier(new NMTokenIdentifier(appAttemptId, getContext()
    .getNodeId(), "testuser", getContext().getNMTokenSecretManager().getCurrentKey()
    .getKeyId()));
  return ugi;
}
 

@Test(timeout = 5000)
// Ensure ContainerTerminated is handled correctly by the TaskAttempt
public void testContainerTerminatedWhileRunning() 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);

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

  taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, null));
  // At state STARTING.
  taImpl.handle(new TaskAttemptEventStartedRemotely(taskAttemptID, contId,
      null));
  assertEquals("Task attempt is not in running state", taImpl.getState(),
      TaskAttemptState.RUNNING);
  taImpl.handle(new TaskAttemptEventContainerTerminated(taskAttemptID, "Terminated"));
  assertFalse(
      "InternalError occurred trying to handle TA_CONTAINER_TERMINATED",
      eventHandler.internalError);

  assertEquals("Terminated", taImpl.getDiagnostics().get(0));

  // TODO Ensure TA_TERMINATING after this is ingored.
}
 

@Test
public void testMapReduceAllocationWithNodeLabelExpression() throws Exception {

  LOG.info("Running testMapReduceAllocationWithNodeLabelExpression");
  Configuration conf = new Configuration();
  /*
   * final int MAP_LIMIT = 3; final int REDUCE_LIMIT = 1;
   * conf.setInt(MRJobConfig.JOB_RUNNING_MAP_LIMIT, MAP_LIMIT);
   * conf.setInt(MRJobConfig.JOB_RUNNING_REDUCE_LIMIT, REDUCE_LIMIT);
   */
  conf.setFloat(MRJobConfig.COMPLETED_MAPS_FOR_REDUCE_SLOWSTART, 1.0f);
  conf.set(MRJobConfig.MAP_NODE_LABEL_EXP, "MapNodes");
  conf.set(MRJobConfig.REDUCE_NODE_LABEL_EXP, "ReduceNodes");
  ApplicationId appId = ApplicationId.newInstance(1, 1);
  ApplicationAttemptId appAttemptId =
      ApplicationAttemptId.newInstance(appId, 1);
  JobId jobId = MRBuilderUtils.newJobId(appAttemptId.getApplicationId(), 0);
  Job mockJob = mock(Job.class);
  when(mockJob.getReport()).thenReturn(
      MRBuilderUtils.newJobReport(jobId, "job", "user", JobState.RUNNING, 0,
          0, 0, 0, 0, 0, 0, "jobfile", null, false, ""));
  final MockScheduler mockScheduler = new MockScheduler(appAttemptId);
  MyContainerAllocator allocator =
      new MyContainerAllocator(null, conf, appAttemptId, mockJob) {
        @Override
        protected void register() {
        }

        @Override
        protected ApplicationMasterProtocol createSchedulerProxy() {
          return mockScheduler;
        }
      };

  // create some map requests
  ContainerRequestEvent reqMapEvents;
  reqMapEvents = createReq(jobId, 0, 1024, new String[] { "map" });
  allocator.sendRequests(Arrays.asList(reqMapEvents));

  // create some reduce requests
  ContainerRequestEvent reqReduceEvents;
  reqReduceEvents =
      createReq(jobId, 0, 2048, new String[] { "reduce" }, false, true);
  allocator.sendRequests(Arrays.asList(reqReduceEvents));
  allocator.schedule();
  // verify all of the host-specific asks were sent plus one for the
  // default rack and one for the ANY request
  Assert.assertEquals(3, mockScheduler.lastAsk.size());
  // verify ResourceRequest sent for MAP have appropriate node
  // label expression as per the configuration
  validateLabelsRequests(mockScheduler.lastAsk.get(0), false);
  validateLabelsRequests(mockScheduler.lastAsk.get(1), false);
  validateLabelsRequests(mockScheduler.lastAsk.get(2), false);

  // assign a map task and verify we do not ask for any more maps
  ContainerId cid0 = mockScheduler.assignContainer("map", false);
  allocator.schedule();
  // default rack and one for the ANY request
  Assert.assertEquals(3, mockScheduler.lastAsk.size());
  validateLabelsRequests(mockScheduler.lastAsk.get(0), true);
  validateLabelsRequests(mockScheduler.lastAsk.get(1), true);
  validateLabelsRequests(mockScheduler.lastAsk.get(2), true);

  // complete the map task and verify that we ask for one more
  allocator.close();
}
 

private static ApplicationAttemptId getApplicationAttemptId(
    ApplicationId applicationId, int startCount) {
  ApplicationAttemptId applicationAttemptId =
      ApplicationAttemptId.newInstance(applicationId, startCount);
  return applicationAttemptId;
}
 

@SuppressWarnings("deprecation")
@Test(timeout = 5000)
public void testKilledInNew() 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);
  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);
  MockTaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler,
      taListener, taskConf, new SystemClock(),
      mockHeartbeatHandler, appCtx, false,
      resource, createFakeContainerContext(), true);
  Assert.assertEquals(TaskAttemptStateInternal.NEW, taImpl.getInternalState());
  taImpl.handle(new TaskAttemptEventKillRequest(taImpl.getID(), "kill it",
      TaskAttemptTerminationCause.TERMINATED_BY_CLIENT));
  Assert.assertEquals(TaskAttemptStateInternal.KILLED, taImpl.getInternalState());

  Assert.assertEquals(0, taImpl.taskAttemptStartedEventLogged);
  Assert.assertEquals(1, taImpl.taskAttemptFinishedEventLogged);
}
 

@Test(timeout=50000)
public void testIdleContainerAssignmentReuseNewContainers() throws Exception {
  AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

  String appHost = "host";
  int appPort = 0;
  String appUrl = "url";

  Configuration conf = new Configuration();
  conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
  conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
  conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false);
  conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
  conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NEW_CONTAINERS_ENABLED, true);
  conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);
  conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 4000);
  conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 5000);
  conf.setInt(TezConfiguration.TEZ_AM_SESSION_MIN_HELD_CONTAINERS, 5);

  DagInfo mockDagInfo = mock(DagInfo.class);
  when(mockDagInfo.getTotalVertices()).thenReturn(10);
  when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
  TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
  when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
  when(mockApp.isSession()).thenReturn(true);
  TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

  MockClock clock = new MockClock(1000);
  NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
      mockRMClient, clock);

  scheduler.initialize();
  drainableAppCallback.drain();

  scheduler.start();
  drainableAppCallback.drain();
  verify(mockRMClient).start();
  verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
  RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
  verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
      regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
      regResponse.getQueue());

  assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

  final String rack1 = "/r1";
  final String rack2 = "/r2";
  final String node1Rack1 = "n1r1";
  final String node2Rack1 = "n2r1";
  final String node1Rack2 = "n1r2";
  MockDNSToSwitchMapping.addRackMapping(node1Rack1, rack1);
  MockDNSToSwitchMapping.addRackMapping(node2Rack1, rack1);
  MockDNSToSwitchMapping.addRackMapping(node1Rack2, rack2);

  Priority priorityv0 = Priority.newInstance(1);
  MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, node1Rack1, rack1);

  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
  ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
  NodeId n2r1 = NodeId.newInstance(node2Rack1, 1);
  Container container1 = Container.newInstance(cid1, n2r1, null, taskv0t0.capability, priorityv0, null);

  // verify idle container is kept for now
  scheduler.onContainersAllocated(Collections.singletonList(container1));
  clock.incrementTime(2000);
  drainableAppCallback.drain();
  verify(mockApp, never()).containerBeingReleased(cid1);
  verify(mockRMClient, never()).releaseAssignedContainer(cid1);

  // verify idle container is released without being assigned to a task because rack-local reuse is
  // disabled
  TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
      scheduler, drainableAppCallback);
  TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
  clock.incrementTime(10000);
  drainableAppCallback.drain();
  verify(mockApp, never()).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
  verify(mockRMClient, never()).removeContainerRequest(reqv0t0);
  verify(mockApp, never()).containerBeingReleased(cid1);
  verify(mockRMClient).releaseAssignedContainer(cid1);

  // cancel the task request
  assertFalse(scheduler.deallocateTask(taskv0t0.task, false, null, null));

  // allocate another container that's node-local
  ContainerId cid2 = ContainerId.newContainerId(attemptId, 2);
  NodeId n1r1 = NodeId.newInstance(node1Rack1, 1);
  Container container2 = Container.newInstance(cid2, n1r1, null, taskv0t0.capability, priorityv0, null);
  scheduler.onContainersAllocated(Collections.singletonList(container2));
  clock.incrementTime(2000);
  drainableAppCallback.drain();
  verify(mockApp, never()).containerBeingReleased(cid2);
  verify(mockRMClient, never()).releaseAssignedContainer(cid2);

  // reschedule the task, verify it's now scheduled without a container request
  // since node-local idle container is available
  reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0, false);
  verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container2);
  verify(mockRMClient).removeContainerRequest(reqv0t0);
}
 

@Test
public void testCleanedupApplicationContainerCleanup() throws IOException {
  NodeManager nm = new NodeManager();
  YarnConfiguration conf = new YarnConfiguration();
  conf.set(NodeStatusUpdaterImpl
          .YARN_NODEMANAGER_DURATION_TO_TRACK_STOPPED_CONTAINERS,
      "1000000");
  nm.init(conf);

  NodeStatusUpdaterImpl nodeStatusUpdater =
      (NodeStatusUpdaterImpl) nm.getNodeStatusUpdater();
  ApplicationId appId = ApplicationId.newInstance(0, 0);
  ApplicationAttemptId appAttemptId =
      ApplicationAttemptId.newInstance(appId, 0);

  ContainerId cId = ContainerId.newContainerId(appAttemptId, 1);
  Token containerToken =
      BuilderUtils.newContainerToken(cId, "anyHost", 1234, "anyUser",
          BuilderUtils.newResource(1024, 1), 0, 123,
          "password".getBytes(), 0);
  Container anyCompletedContainer = new ContainerImpl(conf, null,
      null, null, null, null,
      BuilderUtils.newContainerTokenIdentifier(containerToken)) {

    @Override
    public ContainerState getCurrentState() {
      return ContainerState.COMPLETE;
    }
  };

  Application application = mock(Application.class);
  when(application.getApplicationState()).thenReturn(ApplicationState.RUNNING);
  nm.getNMContext().getApplications().putIfAbsent(appId, application);
  nm.getNMContext().getContainers().put(cId, anyCompletedContainer);

  Assert.assertEquals(1, nodeStatusUpdater.getContainerStatuses().size());

  when(application.getApplicationState()).thenReturn(
      ApplicationState.FINISHING_CONTAINERS_WAIT);
  // The completed container will be saved in case of lost heartbeat.
  Assert.assertEquals(1, nodeStatusUpdater.getContainerStatuses().size());
  Assert.assertEquals(1, nodeStatusUpdater.getContainerStatuses().size());

  nm.getNMContext().getContainers().put(cId, anyCompletedContainer);
  nm.getNMContext().getApplications().remove(appId);
  // The completed container will be saved in case of lost heartbeat.
  Assert.assertEquals(1, nodeStatusUpdater.getContainerStatuses().size());
  Assert.assertEquals(1, nodeStatusUpdater.getContainerStatuses().size());
}
 

private static ApplicationAttemptId getApplicationAttemptId(int id) {
  return ApplicationAttemptId.newInstance(getApplicationId(id), 1);
}
 

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

@Test (timeout = 5000)
public void testContainerAliveOnGetTask() throws IOException {

  TaskCommunicatorContext context = mock(TaskCommunicatorContext.class);
  Configuration conf = new Configuration(false);
  UserPayload userPayload = TezUtils.createUserPayloadFromConf(conf);




  ApplicationId appId = ApplicationId.newInstance(1000, 1);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
  ContainerId containerId = createContainerId(appId, 1);

  doReturn(appAttemptId).when(context).getApplicationAttemptId();
  doReturn(userPayload).when(context).getInitialUserPayload();
  doReturn(new Credentials()).when(context).getAMCredentials();

  TezTaskCommunicatorImpl taskComm = new TezTaskCommunicatorImpl(context);

  ContainerContext containerContext = new ContainerContext(containerId.toString());
  taskComm.registerRunningContainer(containerId, "fakehost", 0);
  ContainerTask containerTask = taskComm.getUmbilical().getTask(containerContext);
  assertNull(containerTask);

  verify(context).containerAlive(containerId);
}