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

@Test
public void testFillInRacks() {
  AMRMClientImpl<ContainerRequest> client =
      new AMRMClientImpl<ContainerRequest>();
  
  Configuration conf = new Configuration();
  conf.setClass(
      CommonConfigurationKeysPublic.NET_TOPOLOGY_NODE_SWITCH_MAPPING_IMPL_KEY,
      MyResolver.class, DNSToSwitchMapping.class);
  client.init(conf);
 
  Resource capability = Resource.newInstance(1024, 1, 1);
  ContainerRequest request =
      new ContainerRequest(capability, new String[] {"host1", "host2"},
          new String[] {"/rack2"}, Priority.newInstance(1));
  client.addContainerRequest(request);
  verifyResourceRequest(client, request, "host1", true);
  verifyResourceRequest(client, request, "host2", true);
  verifyResourceRequest(client, request, "/rack1", true);
  verifyResourceRequest(client, request, "/rack2", true);
  verifyResourceRequest(client, request, ResourceRequest.ANY, true);
}
 

@Test (expected = InvalidContainerRequestException.class)
public void testDifferentLocalityRelaxationSamePriority() {
  AMRMClientImpl<ContainerRequest> client =
      new AMRMClientImpl<ContainerRequest>();
  Configuration conf = new Configuration();
  conf.setClass(
      CommonConfigurationKeysPublic.NET_TOPOLOGY_NODE_SWITCH_MAPPING_IMPL_KEY,
      MyResolver.class, DNSToSwitchMapping.class);
  client.init(conf);
  
  Resource capability = Resource.newInstance(1024, 1);
  ContainerRequest request1 =
      new ContainerRequest(capability, new String[] {"host1", "host2"},
          null, Priority.newInstance(1), false);
  client.addContainerRequest(request1);
  ContainerRequest request2 =
      new ContainerRequest(capability, new String[] {"host3"},
          null, Priority.newInstance(1), true);
  client.addContainerRequest(request2);
}
 

/**
 * Get the list of matching container requests for the specified resource memory and cores.
 */
public List<? extends Collection<AMRMClient.ContainerRequest>> getMatchingRequestsList(int memory, int cores) {
  Resource resource = Resource.newInstance(memory, cores);
  Priority priority = Priority.newInstance(0);

  return getAmrmClientAsync().getMatchingRequests(priority, ResourceRequest.ANY, resource);
}
 

@Test
public void testInvalidValidWhenOldRemoved() {
  AMRMClientImpl<ContainerRequest> client =
      new AMRMClientImpl<ContainerRequest>();
  Configuration conf = new Configuration();
  conf.setClass(
      CommonConfigurationKeysPublic.NET_TOPOLOGY_NODE_SWITCH_MAPPING_IMPL_KEY,
      MyResolver.class, DNSToSwitchMapping.class);
  client.init(conf);
  
  Resource capability = Resource.newInstance(1024, 1, 1);
  ContainerRequest request1 =
      new ContainerRequest(capability, new String[] {"host1", "host2"},
          null, Priority.newInstance(1), false);
  client.addContainerRequest(request1);
  
  client.removeContainerRequest(request1);

  ContainerRequest request2 =
      new ContainerRequest(capability, new String[] {"host3"},
          null, Priority.newInstance(1), true);
  client.addContainerRequest(request2);
  
  client.removeContainerRequest(request2);
  
  ContainerRequest request3 =
      new ContainerRequest(capability, new String[] {"host1", "host2"},
          null, Priority.newInstance(1), false);
  client.addContainerRequest(request3);
  
  client.removeContainerRequest(request3);
  
  ContainerRequest request4 =
      new ContainerRequest(capability, null,
          new String[] {"rack1"}, Priority.newInstance(1), true);
  client.addContainerRequest(request4);

}
 

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

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

@Test
public void testComparePriorities(){
  Priority high = Priority.newInstance(1);
  Priority low = Priority.newInstance(2);
  assertTrue(high.compareTo(low) > 0);
}
 

@Test
public void testMove() {
  final String user = "user1";
  Queue parentQueue = createQueue("parent", null);
  Queue oldQueue = createQueue("old", parentQueue);
  Queue newQueue = createQueue("new", parentQueue);
  QueueMetrics parentMetrics = parentQueue.getMetrics();
  QueueMetrics oldMetrics = oldQueue.getMetrics();
  QueueMetrics newMetrics = newQueue.getMetrics();

  ApplicationAttemptId appAttId = createAppAttemptId(0, 0);
  RMContext rmContext = mock(RMContext.class);
  when(rmContext.getEpoch()).thenReturn(3L);
  SchedulerApplicationAttempt app = new SchedulerApplicationAttempt(appAttId,
      user, oldQueue, oldQueue.getActiveUsersManager(), rmContext);
  oldMetrics.submitApp(user);
  
  // confirm that containerId is calculated based on epoch.
  assertEquals(0x30000000001L, app.getNewContainerId());
  
  // Resource request
  Resource requestedResource = Resource.newInstance(1536, 2);
  Priority requestedPriority = Priority.newInstance(2);
  ResourceRequest request = ResourceRequest.newInstance(requestedPriority,
      ResourceRequest.ANY, requestedResource, 3);
  app.updateResourceRequests(Arrays.asList(request));

  // Allocated container
  RMContainer container1 = createRMContainer(appAttId, 1, requestedResource);
  app.liveContainers.put(container1.getContainerId(), container1);
  SchedulerNode node = createNode();
  app.appSchedulingInfo.allocate(NodeType.OFF_SWITCH, node, requestedPriority,
      request, container1.getContainer());
  
  // Reserved container
  Priority prio1 = Priority.newInstance(1);
  Resource reservedResource = Resource.newInstance(2048, 3);
  RMContainer container2 = createReservedRMContainer(appAttId, 1, reservedResource,
      node.getNodeID(), prio1);
  Map<NodeId, RMContainer> reservations = new HashMap<NodeId, RMContainer>();
  reservations.put(node.getNodeID(), container2);
  app.reservedContainers.put(prio1, reservations);
  oldMetrics.reserveResource(user, reservedResource);
  
  checkQueueMetrics(oldMetrics, 1, 1, 1536, 2, 2048, 3, 3072, 4);
  checkQueueMetrics(newMetrics, 0, 0, 0, 0, 0, 0, 0, 0);
  checkQueueMetrics(parentMetrics, 1, 1, 1536, 2, 2048, 3, 3072, 4);
  
  app.move(newQueue);
  
  checkQueueMetrics(oldMetrics, 0, 0, 0, 0, 0, 0, 0, 0);
  checkQueueMetrics(newMetrics, 1, 1, 1536, 2, 2048, 3, 3072, 4);
  checkQueueMetrics(parentMetrics, 1, 1, 1536, 2, 2048, 3, 3072, 4);
}
 

@Before
public void setup() throws YarnException, IOException {
  // start minicluster
  conf = new YarnConfiguration();
  yarnCluster =
      new MiniYARNCluster(TestAMRMClient.class.getName(), nodeCount, 1, 1);
  yarnCluster.init(conf);
  yarnCluster.start();
  assertNotNull(yarnCluster);
  assertEquals(STATE.STARTED, yarnCluster.getServiceState());

  // start rm client
  yarnClient = (YarnClientImpl) YarnClient.createYarnClient();
  yarnClient.init(conf);
  yarnClient.start();
  assertNotNull(yarnClient);
  assertEquals(STATE.STARTED, yarnClient.getServiceState());

  // get node info
  nodeReports = yarnClient.getNodeReports(NodeState.RUNNING);

  // submit new app
  ApplicationSubmissionContext appContext = 
      yarnClient.createApplication().getApplicationSubmissionContext();
  ApplicationId appId = appContext.getApplicationId();
  // set the application name
  appContext.setApplicationName("Test");
  // Set the priority for the application master
  Priority pri = Priority.newInstance(0);
  appContext.setPriority(pri);
  // Set the queue to which this application is to be submitted in the RM
  appContext.setQueue("default");
  // Set up the container launch context for the application master
  ContainerLaunchContext amContainer = Records
      .newRecord(ContainerLaunchContext.class);
  appContext.setAMContainerSpec(amContainer);
  // unmanaged AM
  appContext.setUnmanagedAM(true);
  // Create the request to send to the applications manager
  SubmitApplicationRequest appRequest = Records
      .newRecord(SubmitApplicationRequest.class);
  appRequest.setApplicationSubmissionContext(appContext);
  // Submit the application to the applications manager
  yarnClient.submitApplication(appContext);

  // wait for app to start
  int iterationsLeft = 30;
  RMAppAttempt appAttempt = null;
  while (iterationsLeft > 0) {
    ApplicationReport appReport = yarnClient.getApplicationReport(appId);
    if (appReport.getYarnApplicationState() ==
        YarnApplicationState.ACCEPTED) {
      attemptId = appReport.getCurrentApplicationAttemptId();
      appAttempt =
          yarnCluster.getResourceManager().getRMContext().getRMApps()
            .get(attemptId.getApplicationId()).getCurrentAppAttempt();
      while (true) {
        if (appAttempt.getAppAttemptState() == RMAppAttemptState.LAUNCHED) {
          break;
        }
      }
      break;
    }
    sleep(1000);
    --iterationsLeft;
  }
  if (iterationsLeft == 0) {
    fail("Application hasn't bee started");
  }

  // Just dig into the ResourceManager and get the AMRMToken just for the sake
  // of testing.
  UserGroupInformation.setLoginUser(UserGroupInformation
    .createRemoteUser(UserGroupInformation.getCurrentUser().getUserName()));
  UserGroupInformation.getCurrentUser().addToken(appAttempt.getAMRMToken());

  //creating an instance NMTokenCase
  nmTokenCache = new NMTokenCache();
  
  // start am rm client
  rmClient =
      (AMRMClientImpl<ContainerRequest>) AMRMClient
        .<ContainerRequest> createAMRMClient();

  //setting an instance NMTokenCase
  rmClient.setNMTokenCache(nmTokenCache);
  rmClient.init(conf);
  rmClient.start();
  assertNotNull(rmClient);
  assertEquals(STATE.STARTED, rmClient.getServiceState());

  // start am nm client
  nmClient = (NMClientImpl) NMClient.createNMClient();
  
  //propagating the AMRMClient NMTokenCache instance
  nmClient.setNMTokenCache(rmClient.getNMTokenCache());
  nmClient.init(conf);
  nmClient.start();
  assertNotNull(nmClient);
  assertEquals(STATE.STARTED, nmClient.getServiceState());
}
 

WrappedApplication(int id, long timestamp, String user, int numContainers) {
  Configuration conf = new Configuration();
  
  dispatcher = new DrainDispatcher();
  containerTokenIdentifierMap =
      new HashMap<ContainerId, ContainerTokenIdentifier>();
  dispatcher.init(conf);

  localizerBus = mock(EventHandler.class);
  launcherBus = mock(EventHandler.class);
  monitorBus = mock(EventHandler.class);
  auxBus = mock(EventHandler.class);
  containerBus = mock(EventHandler.class);
  logAggregationBus = mock(EventHandler.class);

  dispatcher.register(LocalizationEventType.class, localizerBus);
  dispatcher.register(ContainersLauncherEventType.class, launcherBus);
  dispatcher.register(ContainersMonitorEventType.class, monitorBus);
  dispatcher.register(AuxServicesEventType.class, auxBus);
  dispatcher.register(ContainerEventType.class, containerBus);
  dispatcher.register(LogHandlerEventType.class, logAggregationBus);

  nmTokenSecretMgr = mock(NMTokenSecretManagerInNM.class);

  context = mock(Context.class);
  
  when(context.getContainerTokenSecretManager()).thenReturn(
    new NMContainerTokenSecretManager(conf));
  when(context.getApplicationACLsManager()).thenReturn(
    new ApplicationACLsManager(conf));
  when(context.getNMTokenSecretManager()).thenReturn(nmTokenSecretMgr);
  
  // Setting master key
  MasterKey masterKey = new MasterKeyPBImpl();
  masterKey.setKeyId(123);
  masterKey.setBytes(ByteBuffer.wrap(new byte[] { (new Integer(123)
    .byteValue()) }));
  context.getContainerTokenSecretManager().setMasterKey(masterKey);
  
  this.user = user;
  this.appId = BuilderUtils.newApplicationId(timestamp, id);

  app = new ApplicationImpl(dispatcher, this.user, appId, null, context);
  containers = new ArrayList<Container>();
  for (int i = 0; i < numContainers; i++) {
    Container container = createMockedContainer(this.appId, i);
    containers.add(container);
    long currentTime = System.currentTimeMillis();
    ContainerTokenIdentifier identifier =
        new ContainerTokenIdentifier(container.getContainerId(), "", "",
          null, currentTime + 2000, masterKey.getKeyId(), currentTime,
          Priority.newInstance(0), 0);
    containerTokenIdentifierMap
      .put(identifier.getContainerID(), identifier);
    context.getContainerTokenSecretManager().startContainerSuccessful(
      identifier);
    Assert.assertFalse(context.getContainerTokenSecretManager()
      .isValidStartContainerRequest(identifier));
  }

  dispatcher.start();
}
 

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

private Set<Container> allocateContainers(
    AMRMClientImpl<ContainerRequest> rmClient, int num)
    throws YarnException, IOException {
  // setup container request
  Resource capability = Resource.newInstance(1024, 0);
  Priority priority = Priority.newInstance(0);
  String node = nodeReports.get(0).getNodeId().getHost();
  String rack = nodeReports.get(0).getRackName();
  String[] nodes = new String[] {node};
  String[] racks = new String[] {rack};

  for (int i = 0; i < num; ++i) {
    rmClient.addContainerRequest(new ContainerRequest(capability, nodes,
        racks, priority));
  }

  int containersRequestedAny = rmClient.remoteRequestsTable.get(priority)
      .get(ResourceRequest.ANY).get(capability).remoteRequest
      .getNumContainers();

  // RM should allocate container within 2 calls to allocate()
  int allocatedContainerCount = 0;
  int iterationsLeft = 2;
  Set<Container> containers = new TreeSet<Container>();
  while (allocatedContainerCount < containersRequestedAny
      && iterationsLeft > 0) {
    AllocateResponse allocResponse = rmClient.allocate(0.1f);

    allocatedContainerCount += allocResponse.getAllocatedContainers().size();
    for(Container container : allocResponse.getAllocatedContainers()) {
      containers.add(container);
    }
    if (!allocResponse.getNMTokens().isEmpty()) {
      for (NMToken token : allocResponse.getNMTokens()) {
        rmClient.getNMTokenCache().setToken(token.getNodeId().toString(),
            token.getToken());
      }
    }
    if(allocatedContainerCount < containersRequestedAny) {
      // sleep to let NM's heartbeat to RM and trigger allocations
      sleep(1000);
    }

    --iterationsLeft;
  }
  return containers;
}
 

private ContainerRequest createReq(int priority, int memory, String[] hosts) {
  Resource capability = Resource.newInstance(memory, 1);
  Priority priorityOfContainer = Priority.newInstance(priority);
  return new ContainerRequest(capability, hosts,
      new String[] { NetworkTopology.DEFAULT_RACK }, priorityOfContainer);
}
 

WrappedApplication(int id, long timestamp, String user, int numContainers) {
  Configuration conf = new Configuration();
  
  dispatcher = new DrainDispatcher();
  containerTokenIdentifierMap =
      new HashMap<ContainerId, ContainerTokenIdentifier>();
  dispatcher.init(conf);

  localizerBus = mock(EventHandler.class);
  launcherBus = mock(EventHandler.class);
  monitorBus = mock(EventHandler.class);
  auxBus = mock(EventHandler.class);
  containerBus = mock(EventHandler.class);
  logAggregationBus = mock(EventHandler.class);

  dispatcher.register(LocalizationEventType.class, localizerBus);
  dispatcher.register(ContainersLauncherEventType.class, launcherBus);
  dispatcher.register(ContainersMonitorEventType.class, monitorBus);
  dispatcher.register(AuxServicesEventType.class, auxBus);
  dispatcher.register(ContainerEventType.class, containerBus);
  dispatcher.register(LogHandlerEventType.class, logAggregationBus);

  nmTokenSecretMgr = mock(NMTokenSecretManagerInNM.class);

  context = mock(Context.class);
  
  when(context.getContainerTokenSecretManager()).thenReturn(
    new NMContainerTokenSecretManager(conf));
  when(context.getApplicationACLsManager()).thenReturn(
    new ApplicationACLsManager(conf));
  when(context.getNMTokenSecretManager()).thenReturn(nmTokenSecretMgr);
  
  // Setting master key
  MasterKey masterKey = new MasterKeyPBImpl();
  masterKey.setKeyId(123);
  masterKey.setBytes(ByteBuffer.wrap(new byte[] { (new Integer(123)
    .byteValue()) }));
  context.getContainerTokenSecretManager().setMasterKey(masterKey);
  
  this.user = user;
  this.appId = BuilderUtils.newApplicationId(timestamp, id);

  app = new ApplicationImpl(dispatcher, this.user, appId, null, context);
  containers = new ArrayList<Container>();
  for (int i = 0; i < numContainers; i++) {
    Container container = createMockedContainer(this.appId, i);
    containers.add(container);
    long currentTime = System.currentTimeMillis();
    ContainerTokenIdentifier identifier =
        new ContainerTokenIdentifier(container.getContainerId(), "", "",
          null, currentTime + 2000, masterKey.getKeyId(), currentTime,
          Priority.newInstance(0), 0);
    containerTokenIdentifierMap
      .put(identifier.getContainerID(), identifier);
    context.getContainerTokenSecretManager().startContainerSuccessful(
      identifier);
    Assert.assertFalse(context.getContainerTokenSecretManager()
      .isValidStartContainerRequest(identifier));
  }

  dispatcher.start();
}
 

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

private void testRPCTimeout(String rpcClass) throws Exception {
  Configuration conf = new Configuration();
  // set timeout low for the test
  conf.setInt("yarn.rpc.nm-command-timeout", 3000);

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

    ContainerManagementProtocol proxy = (ContainerManagementProtocol) rpc.getProxy(
        ContainerManagementProtocol.class,
        server.getListenerAddress(), 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, 3);
    ContainerTokenIdentifier containerTokenIdentifier =
        new ContainerTokenIdentifier(containerId, "localhost", "user",
          resource, System.currentTimeMillis() + 10000, 42, 42,
          Priority.newInstance(0), 0);
    Token containerToken =
        TestRPC.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);
    try {
      proxy.startContainers(allRequests);
    } catch (Exception e) {
      LOG.info(StringUtils.stringifyException(e));
      Assert.assertEquals("Error, exception is not: "
          + SocketTimeoutException.class.getName(),
          SocketTimeoutException.class.getName(), e.getClass().getName());
      return;
    }
  } finally {
    server.stop();
  }

  Assert.fail("timeout exception should have occurred!");
}
 

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
public void testContainerTokenIdentifier() throws IOException {
  ContainerId containerID = ContainerId.newContainerId(
      ApplicationAttemptId.newInstance(ApplicationId.newInstance(
          1, 1), 1), 1);
  String hostName = "host0";
  String appSubmitter = "usr0";
  Resource r = Resource.newInstance(1024, 1);
  long expiryTimeStamp = 1000;
  int masterKeyId = 1;
  long rmIdentifier = 1;
  Priority priority = Priority.newInstance(1);
  long creationTime = 1000;
  
  ContainerTokenIdentifier token = new ContainerTokenIdentifier(
      containerID, hostName, appSubmitter, r, expiryTimeStamp, 
      masterKeyId, rmIdentifier, priority, creationTime);
  
  ContainerTokenIdentifier anotherToken = new ContainerTokenIdentifier();
  
  byte[] tokenContent = token.getBytes();
  DataInputBuffer dib = new DataInputBuffer();
  dib.reset(tokenContent, tokenContent.length);
  anotherToken.readFields(dib);
  
  // verify the whole record equals with original record
  Assert.assertEquals("Token is not the same after serialization " +
      "and deserialization.", token, anotherToken);
  
  Assert.assertEquals(
      "ContainerID from proto is not the same with original token",
      anotherToken.getContainerID(), containerID);
  
  Assert.assertEquals(
      "Hostname from proto is not the same with original token",
      anotherToken.getNmHostAddress(), hostName);
  
  Assert.assertEquals(
      "ApplicationSubmitter from proto is not the same with original token",
      anotherToken.getApplicationSubmitter(), appSubmitter);
  
  Assert.assertEquals(
      "Resource from proto is not the same with original token",
      anotherToken.getResource(), r);
  
  Assert.assertEquals(
      "expiryTimeStamp from proto is not the same with original token",
      anotherToken.getExpiryTimeStamp(), expiryTimeStamp);
  
  Assert.assertEquals(
      "KeyId from proto is not the same with original token",
      anotherToken.getMasterKeyId(), masterKeyId);
  
  Assert.assertEquals(
      "RMIdentifier from proto is not the same with original token",
      anotherToken.getRMIdentifier(), rmIdentifier);
  
  Assert.assertEquals(
      "Priority from proto is not the same with original token",
      anotherToken.getPriority(), priority);
  
  Assert.assertEquals(
      "CreationTime from proto is not the same with original token",
      anotherToken.getCreationTime(), creationTime);
  
  Assert.assertNull(anotherToken.getLogAggregationContext());
}
 

@Test
public void testComparePriorities(){
  Priority high = Priority.newInstance(1);
  Priority low = Priority.newInstance(2);
  assertTrue(high.compareTo(low) > 0);
}