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

@Test
public void testMinMaxContainerIdleMillisAreEqual() throws Exception {
  AMRMClientAsyncWrapperForTest mockRMClient = new AMRMClientAsyncWrapperForTest();
  Configuration conf = new Configuration();
  conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 10000);
  conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 10000);

  TaskSchedulerContext mockApp = setupMockTaskSchedulerContext("host", 0, "url", conf);
  TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
  MockClock clock = new MockClock(1000);
  NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);
  scheduler.initialize();

  NodeId host1 = NodeId.newInstance("host1", 1);
  Container container1 = Container.newInstance(null, host1, null, null, null, null);
  HeldContainer heldContainer = scheduler.new HeldContainer(container1);
  long now = clock.getTime();
  assertEquals(now + 10000, heldContainer.getIdleExpirationTimestamp(now));
}
 

public ContainerId assignContainer(String nodeName, boolean isReduce) {
  ContainerId containerId =
      ContainerId.newContainerId(attemptId, nextContainerId++);
  Priority priority = isReduce ? RMContainerAllocator.PRIORITY_REDUCE
      : RMContainerAllocator.PRIORITY_MAP;
  Container container = Container.newInstance(containerId,
      NodeId.newInstance(nodeName, 1234), nodeName + ":5678",
    Resource.newInstance(1024, 1), priority, null);
  containersToAllocate.add(container);
  return containerId;
}
 

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

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

@Test (timeout=3000)
public void testTaskSchedulerHeldContainersReleaseAfterExpired() throws Exception {
  final TezAMRMClientAsync<CookieContainerRequest> mockRMClient = spy(
    new AMRMClientAsyncForTest(new AMRMClientForTest(), 100));
  final TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(DEFAULT_APP_HOST, DEFAULT_APP_PORT,
    DEFAULT_APP_URL, true, new Configuration());
  final TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
  final TaskSchedulerWithDrainableContext scheduler =
    new TaskSchedulerWithDrainableContext(drainableAppCallback, mockRMClient);

  scheduler.initialize();
  scheduler.start();

  Resource mockCapability = Resource.newInstance(1024, 1);
  NodeId emptyHost = NodeId.newInstance("", 1);
  ApplicationAttemptId appId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(0, 0), 0);
  ContainerId containerId = ContainerId.newInstance(appId, 0);
  Container c1 = Container.newInstance(containerId, emptyHost, null, mockCapability, null, null);

  HeldContainer hc1 = new HeldContainer(c1, -1, -1, null, containerSignatureMatcher);

  // containerExpiryTime = 0
  scheduler.heldContainers.put(containerId, hc1);

  long currTime = System.currentTimeMillis();
  scheduler.delayedContainerManager.addDelayedContainer(hc1.getContainer(), currTime);
  // sleep and wait for mainLoop() check-in to release this expired held container.
  Thread.sleep(1000);

  verify(mockRMClient, times(1)).releaseAssignedContainer((ContainerId)any());
  Assert.assertEquals(0, scheduler.heldContainers.size());

  AppFinalStatus finalStatus =
    new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, SUCCEED_APP_MESSAGE, DEFAULT_APP_URL);
  when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
  scheduler.shutdown();
}
 

/**
 * @param host Host.
 * @param cpu Cpu count.
 * @param mem Memory.
 * @return Container.
 */
private Container createContainer(String host, int cpu, int mem) {
    return Container.newInstance(
        ContainerId.newContainerId(ApplicationAttemptId.newInstance(ApplicationId.newInstance(0l, 0), 0),
            ThreadLocalRandom.current().nextLong()),
        NodeId.newInstance(host, 0),
        "example.com",
        new MockResource(mem, cpu),
        Priority.newInstance(0),
        null
    );
}
 

private RMContainer recoverAndCreateContainer(NMContainerStatus status,
    RMNode node) {
  Container container =
      Container.newInstance(status.getContainerId(), node.getNodeID(),
        node.getHttpAddress(), status.getAllocatedResource(),
        status.getPriority(), null);
  ApplicationAttemptId attemptId =
      container.getId().getApplicationAttemptId();
  RMContainer rmContainer =
      new RMContainerImpl(container, attemptId, node.getNodeID(),
        applications.get(attemptId.getApplicationId()).getUser(), rmContext,
        status.getCreationTime());
  return rmContainer;
}
 

@Override
public void handle(ContainerAllocatorEvent event) {
  ContainerId cId =
      ContainerId.newContainerId(getContext().getApplicationAttemptId(),
        containerCount++);
  NodeId nodeId = NodeId.newInstance(NM_HOST, NM_PORT);
  Resource resource = Resource.newInstance(1234, 2);
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier(cId, nodeId.toString(), "user",
      resource, System.currentTimeMillis() + 10000, 42, 42,
      Priority.newInstance(0), 0);
  Token containerToken = newContainerToken(nodeId, "password".getBytes(),
        containerTokenIdentifier);
  Container container = Container.newInstance(cId, nodeId,
      NM_HOST + ":" + NM_HTTP_PORT, resource, null, containerToken);
  JobID id = TypeConverter.fromYarn(applicationId);
  JobId jobId = TypeConverter.toYarn(id);
  getContext().getEventHandler().handle(new JobHistoryEvent(jobId, 
      new NormalizedResourceEvent(
          org.apache.hadoop.mapreduce.TaskType.REDUCE,
      100)));
  getContext().getEventHandler().handle(new JobHistoryEvent(jobId, 
      new NormalizedResourceEvent(
          org.apache.hadoop.mapreduce.TaskType.MAP,
      100)));
  getContext().getEventHandler().handle(
      new TaskAttemptContainerAssignedEvent(event.getAttemptID(),
          container, null));
}
 

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

private RMContainer recoverAndCreateContainer(NMContainerStatus status,
    RMNode node) {
  Container container =
      Container.newInstance(status.getContainerId(), node.getNodeID(),
        node.getHttpAddress(), status.getAllocatedResource(),
        status.getPriority(), null);
  ApplicationAttemptId attemptId =
      container.getId().getApplicationAttemptId();
  RMContainer rmContainer =
      new RMContainerImpl(container, attemptId, node.getNodeID(),
        applications.get(attemptId.getApplicationId()).getUser(), rmContext,
        status.getCreationTime());
  return rmContainer;
}
 

public ContainerId assignContainer(String nodeName, boolean isReduce) {
  ContainerId containerId =
      ContainerId.newContainerId(attemptId, nextContainerId++);
  Priority priority = isReduce ? RMContainerAllocator.PRIORITY_REDUCE
      : RMContainerAllocator.PRIORITY_MAP;
  Container container = Container.newInstance(containerId,
      NodeId.newInstance(nodeName, 1234), nodeName + ":5678",
    Resource.newInstance(1024, 1, 1), priority, null);
  containersToAllocate.add(container);
  return containerId;
}
 

@Override
public void handle(ContainerAllocatorEvent event) {
  ContainerId cId =
      ContainerId.newContainerId(getContext().getApplicationAttemptId(),
        containerCount++);
  NodeId nodeId = NodeId.newInstance(NM_HOST, NM_PORT);
  Resource resource = Resource.newInstance(1234, 2, 2);
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier(cId, nodeId.toString(), "user",
      resource, System.currentTimeMillis() + 10000, 42, 42,
      Priority.newInstance(0), 0);
  Token containerToken = newContainerToken(nodeId, "password".getBytes(),
        containerTokenIdentifier);
  Container container = Container.newInstance(cId, nodeId,
      NM_HOST + ":" + NM_HTTP_PORT, resource, null, containerToken);
  JobID id = TypeConverter.fromYarn(applicationId);
  JobId jobId = TypeConverter.toYarn(id);
  getContext().getEventHandler().handle(new JobHistoryEvent(jobId, 
      new NormalizedResourceEvent(
          org.apache.hadoop.mapreduce.TaskType.REDUCE,
      100)));
  getContext().getEventHandler().handle(new JobHistoryEvent(jobId, 
      new NormalizedResourceEvent(
          org.apache.hadoop.mapreduce.TaskType.MAP,
      100)));
  getContext().getEventHandler().handle(
      new TaskAttemptContainerAssignedEvent(event.getAttemptID(),
          container, null));
}
 

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

@Test(timeout=50000)
public void testContainerAssignmentReleaseNewContainers() 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, false);
  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 new container is released is not immediately allocated
  scheduler.onContainersAllocated(Collections.singletonList(container1));
  drainableAppCallback.drain();
  // app is not notified of the container being released since it never launched
  verify(mockApp, never()).containerBeingReleased(cid1);
  verify(mockRMClient).releaseAssignedContainer(eq(cid1));
}
 

@Test
public void testNMSimulator() throws Exception {
  // Register one node
  NMSimulator node1 = new NMSimulator();
  node1.init("rack1/node1", GB * 10, 10, 0, 1000, rm);
  node1.middleStep();

  int numClusterNodes = rm.getResourceScheduler().getNumClusterNodes();
  int cumulativeSleepTime = 0;
  int sleepInterval = 100;

  while(numClusterNodes != 1 && cumulativeSleepTime < 5000) {
    Thread.sleep(sleepInterval);
    cumulativeSleepTime = cumulativeSleepTime + sleepInterval;
    numClusterNodes = rm.getResourceScheduler().getNumClusterNodes();
  }

  Assert.assertEquals(1, rm.getResourceScheduler().getNumClusterNodes());
  Assert.assertEquals(GB * 10,
      rm.getResourceScheduler().getRootQueueMetrics().getAvailableMB());
  Assert.assertEquals(10,
      rm.getResourceScheduler().getRootQueueMetrics()
          .getAvailableVirtualCores());

  // Allocate one container on node1
  ContainerId cId1 = newContainerId(1, 1, 1);
  Container container1 = Container.newInstance(cId1, null, null,
      Resources.createResource(GB, 1), null, null);
  node1.addNewContainer(container1, 100000l);
  Assert.assertTrue("Node1 should have one running container.",
      node1.getRunningContainers().containsKey(cId1));

  // Allocate one AM container on node1
  ContainerId cId2 = newContainerId(2, 1, 1);
  Container container2 = Container.newInstance(cId2, null, null,
      Resources.createResource(GB, 1), null, null);
  node1.addNewContainer(container2, -1l);
  Assert.assertTrue("Node1 should have one running AM container",
      node1.getAMContainers().contains(cId2));

  // Remove containers
  node1.cleanupContainer(cId1);
  Assert.assertTrue("Container1 should be removed from Node1.",
      node1.getCompletedContainers().contains(cId1));
  node1.cleanupContainer(cId2);
  Assert.assertFalse("Container2 should be removed from Node1.",
      node1.getAMContainers().contains(cId2));
}
 

@Test
public void testContainerExpired() throws Exception {
  TezAMRMClientAsync<CookieContainerRequest> mockRMClient = spy(
      new AMRMClientAsyncForTest(new AMRMClientForTest(), 100));

  Configuration conf = new Configuration();
  // to match all in the same pass
  conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
  // to release immediately after deallocate
  conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
  conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 0);

  TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(DEFAULT_APP_HOST, DEFAULT_APP_PORT, DEFAULT_APP_URL, conf);
  final TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
  TaskSchedulerWithDrainableContext scheduler =
      new TaskSchedulerWithDrainableContext(drainableAppCallback, mockRMClient);

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

  Object mockTask1 = new MockTask("task1");
  Object mockCookie1 = new Object();
  Resource mockCapability = Resource.newInstance(1024, 1);
  String[] hosts = {"host1", "host5"};
  String[] racks = {"/default-rack", "/default-rack"};
  final Priority mockPriority1 = Priority.newInstance(1);
  final Priority mockPriority2 = Priority.newInstance(2);
  Object mockTask2 = new MockTask("task2");
  Object mockCookie2 = new Object();
  ArgumentCaptor<CookieContainerRequest> requestCaptor =
      ArgumentCaptor.forClass(CookieContainerRequest.class);

  scheduler.allocateTask(mockTask2, mockCapability, hosts,
      racks, mockPriority2, null, mockCookie2);
  drainableAppCallback.drain();
  verify(mockRMClient, times(1)).
                              addContainerRequest(requestCaptor.capture());
  CookieContainerRequest request2 = requestCaptor.getValue();

  scheduler.allocateTask(mockTask1, mockCapability, hosts,
      racks, mockPriority1, null, mockCookie1);
  drainableAppCallback.drain();
  verify(mockRMClient, times(2)).
                              addContainerRequest(requestCaptor.capture());

  List<Container> containers = new ArrayList<Container>();
  // sending only lower priority container to make sure its not matched
  NodeId host2 = NodeId.newInstance("host2", 2);
  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
  ContainerId mockCId2 = ContainerId.newContainerId(attemptId, 2);
  Container mockContainer2 = Container.newInstance(mockCId2, host2, null, mockCapability, mockPriority2, null);
  containers.add(mockContainer2);

  scheduler.onContainersAllocated(containers);

  List<ContainerStatus> statuses = new ArrayList<ContainerStatus>();
  ContainerStatus mockStatus2 = mock(ContainerStatus.class);
  when(mockStatus2.getContainerId()).thenReturn(mockCId2);
  statuses.add(mockStatus2);
  scheduler.onContainersCompleted(statuses);

  verify(mockApp, times(0)).taskAllocated(any(), any(), any(Container.class));
  verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
  CookieContainerRequest resubmitRequest = requestCaptor.getValue();
  assertEquals(request2.getCookie().getTask(), resubmitRequest.getCookie().getTask());
  assertEquals(request2.getCookie().getAppCookie(), resubmitRequest.getCookie().getAppCookie());
  assertEquals(request2.getCookie().getContainerSignature(), resubmitRequest.getCookie().getContainerSignature());
  assertEquals(request2.getCapability(), resubmitRequest.getCapability());
  assertEquals(request2.getPriority(), resubmitRequest.getPriority());

  // verify container is not re-requested when nothing at that priority
  assertFalse(scheduler.deallocateTask(mockTask2, true, null, null));
  scheduler.onContainersAllocated(containers);
  scheduler.onContainersCompleted(statuses);
  verify(mockApp, times(0)).taskAllocated(any(), any(), any(Container.class));
  verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
}
 

@Test(timeout=30000)
public void testReuseWithAffinity() 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, true);
  conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
  conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

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

  Priority priorityv0 = Priority.newInstance(1);
  Priority priorityv1 = Priority.newInstance(2);
  String[] hostsv0t0 = { "host1", "host2" };
  MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
  MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, hostsv0t0);

  TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient,
      scheduler, drainableAppCallback);
  TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
  taskRequestCaptor.scheduleTask(taskv0t1);

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

  // allocate one container at v0 priority
  scheduler.onContainersAllocated(Collections.singletonList(container1));
  drainableAppCallback.drain();
  verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
  verify(mockRMClient).removeContainerRequest(reqv0t0);

  // add a new request for this container
  MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
  TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0, cid1);

  // finish v0t0 successfully, verify v0t1 is skipped even though it is node-local
  // and v1t0 assigned instead for affinity
  assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
  clock.incrementTime(10000);
  drainableAppCallback.drain();
  verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
  verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
  verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
  verify(mockRMClient).removeContainerRequest(reqv1t0);

  String appMsg = "success";
  AppFinalStatus finalStatus =
      new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
  when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
  scheduler.shutdown();
  drainableAppCallback.drain();
  verify(mockRMClient).
      unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
          appMsg, appUrl);
  verify(mockRMClient).stop();
}
 

@Test
public void testNMSimulator() throws Exception {
  // Register one node
  NMSimulator node1 = new NMSimulator();
  node1.init("rack1/node1", GB * 10, 10, 0, 1000, rm);
  node1.middleStep();

  int numClusterNodes = rm.getResourceScheduler().getNumClusterNodes();
  int cumulativeSleepTime = 0;
  int sleepInterval = 100;

  while(numClusterNodes != 1 && cumulativeSleepTime < 5000) {
    Thread.sleep(sleepInterval);
    cumulativeSleepTime = cumulativeSleepTime + sleepInterval;
    numClusterNodes = rm.getResourceScheduler().getNumClusterNodes();
  }

  Assert.assertEquals(1, rm.getResourceScheduler().getNumClusterNodes());
  Assert.assertEquals(GB * 10,
      rm.getResourceScheduler().getRootQueueMetrics().getAvailableMB());
  Assert.assertEquals(10,
      rm.getResourceScheduler().getRootQueueMetrics()
          .getAvailableVirtualCores());

  // Allocate one container on node1
  ContainerId cId1 = newContainerId(1, 1, 1);
  Container container1 = Container.newInstance(cId1, null, null,
      Resources.createResource(GB, 1), null, null);
  node1.addNewContainer(container1, 100000l);
  Assert.assertTrue("Node1 should have one running container.",
      node1.getRunningContainers().containsKey(cId1));

  // Allocate one AM container on node1
  ContainerId cId2 = newContainerId(2, 1, 1);
  Container container2 = Container.newInstance(cId2, null, null,
      Resources.createResource(GB, 1), null, null);
  node1.addNewContainer(container2, -1l);
  Assert.assertTrue("Node1 should have one running AM container",
      node1.getAMContainers().contains(cId2));

  // Remove containers
  node1.cleanupContainer(cId1);
  Assert.assertTrue("Container1 should be removed from Node1.",
      node1.getCompletedContainers().contains(cId1));
  node1.cleanupContainer(cId2);
  Assert.assertFalse("Container2 should be removed from Node1.",
      node1.getAMContainers().contains(cId2));
}
 

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