Java Code Examples for org.apache.hadoop.yarn.server.utils.BuilderUtils#newResource()

public RegisterNodeManagerResponse registerNode(
    List<NMContainerStatus> containerReports,
    List<ApplicationId> runningApplications) throws Exception {
  RegisterNodeManagerRequest req = Records.newRecord(
      RegisterNodeManagerRequest.class);
  req.setNodeId(nodeId);
  req.setHttpPort(httpPort);
  Resource resource = BuilderUtils.newResource(memory, vCores);
  req.setResource(resource);
  req.setContainerStatuses(containerReports);
  req.setNMVersion(version);
  req.setRunningApplications(runningApplications);
  RegisterNodeManagerResponse registrationResponse =
      resourceTracker.registerNodeManager(req);
  this.currentContainerTokenMasterKey =
      registrationResponse.getContainerTokenMasterKey();
  this.currentNMTokenMasterKey = registrationResponse.getNMTokenMasterKey();
  return registrationResponse;    
}
 

public static Token createContainerToken(ContainerId cId, long rmIdentifier,
    NodeId nodeId, String user,
    NMContainerTokenSecretManager containerTokenSecretManager,
    LogAggregationContext logAggregationContext)
    throws IOException {
  Resource r = BuilderUtils.newResource(1024, 1);
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier(cId, nodeId.toString(), user, r,
        System.currentTimeMillis() + 100000L, 123, rmIdentifier,
        Priority.newInstance(0), 0, logAggregationContext);
  Token containerToken =
      BuilderUtils
        .newContainerToken(nodeId, containerTokenSecretManager
          .retrievePassword(containerTokenIdentifier),
          containerTokenIdentifier);
  return containerToken;
}
 

@SuppressWarnings("unchecked")
private void startAM() throws YarnException, IOException {
  // application/container configuration
  int heartbeatInterval = conf.getInt(
          SLSConfiguration.AM_HEARTBEAT_INTERVAL_MS,
          SLSConfiguration.AM_HEARTBEAT_INTERVAL_MS_DEFAULT);
  int containerMemoryMB = conf.getInt(SLSConfiguration.CONTAINER_MEMORY_MB,
          SLSConfiguration.CONTAINER_MEMORY_MB_DEFAULT);
  int containerVCores = conf.getInt(SLSConfiguration.CONTAINER_VCORES,
          SLSConfiguration.CONTAINER_VCORES_DEFAULT);
  Resource containerResource =
          BuilderUtils.newResource(containerMemoryMB, containerVCores);

  // application workload
  if (isSLS) {
    startAMFromSLSTraces(containerResource, heartbeatInterval);
  } else {
    startAMFromRumenTraces(containerResource, heartbeatInterval);
  }
  numAMs = amMap.size();
  remainingApps = numAMs;
}
 

@Test
public void testNodeRegistrationSuccess() throws Exception {
  writeToHostsFile("host2");
  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.RM_NODES_INCLUDE_FILE_PATH, hostFile
      .getAbsolutePath());
  rm = new MockRM(conf);
  rm.start();

  ResourceTrackerService resourceTrackerService = rm.getResourceTrackerService();
  RegisterNodeManagerRequest req = Records.newRecord(
      RegisterNodeManagerRequest.class);
  NodeId nodeId = NodeId.newInstance("host2", 1234);
  Resource capability = BuilderUtils.newResource(1024, 1, 1);
  req.setResource(capability);
  req.setNodeId(nodeId);
  req.setHttpPort(1234);
  req.setNMVersion(YarnVersionInfo.getVersion());
  // trying to register a invalid node.
  RegisterNodeManagerResponse response = resourceTrackerService.registerNodeManager(req);
  Assert.assertEquals(NodeAction.NORMAL,response.getNodeAction());
}
 

private NodeReport createNodeReports(RMNode rmNode) {    
  SchedulerNodeReport schedulerNodeReport = 
      scheduler.getNodeReport(rmNode.getNodeID());
  Resource used = BuilderUtils.newResource(0, 0);
  int numContainers = 0;
  if (schedulerNodeReport != null) {
    used = schedulerNodeReport.getUsedResource();
    numContainers = schedulerNodeReport.getNumContainers();
  } 
  
  NodeReport report =
      BuilderUtils.newNodeReport(rmNode.getNodeID(), rmNode.getState(),
          rmNode.getHttpAddress(), rmNode.getRackName(), used,
          rmNode.getTotalCapability(), numContainers,
          rmNode.getHealthReport(), rmNode.getLastHealthReportTime(),
          rmNode.getNodeLabels());

  return report;
}
 

protected Token createContainerToken(ContainerId cId, Priority priority,
    long createTime) throws InvalidToken {
  Resource r = BuilderUtils.newResource(1024, 1);
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier(cId, context.getNodeId().toString(), user,
        r, System.currentTimeMillis() + 10000L, 123, DUMMY_RM_IDENTIFIER,
        priority, createTime);
  Token containerToken =
      BuilderUtils.newContainerToken(
        context.getNodeId(),
        context.getContainerTokenSecretManager().retrievePassword(
          containerTokenIdentifier), containerTokenIdentifier);
  return containerToken;
}
 

private static ContainerTokenIdentifier createContainerTokenId(
    ContainerId cid, NodeId nodeId, String user,
    NMContainerTokenSecretManager secretMgr) throws IOException {
  long rmid = cid.getApplicationAttemptId().getApplicationId()
      .getClusterTimestamp();
  ContainerTokenIdentifier ctid = new ContainerTokenIdentifier(cid,
      nodeId.toString(), user, BuilderUtils.newResource(1024, 1),
      System.currentTimeMillis() + 100000L,
      secretMgr.getCurrentKey().getKeyId(), rmid,
      Priority.newInstance(0), 0);
  Token token = BuilderUtils.newContainerToken(nodeId,
      secretMgr.createPassword(ctid), ctid);
  return BuilderUtils.newContainerTokenIdentifier(token);
}
 

private Schedulable createSchedulable(int memUsage, int cpuUsage,
    ResourceWeights weights, int minMemShare, int minCpuShare) {
  Resource usage = BuilderUtils.newResource(memUsage, cpuUsage);
  Resource minShare = BuilderUtils.newResource(minMemShare, minCpuShare);
  return new FakeSchedulable(minShare,
      Resources.createResource(Integer.MAX_VALUE, Integer.MAX_VALUE),
      weights, Resources.none(), usage, 0l);
}
 

private ApplicationId submitAppAndGetAppId(String submitter,
    String queueName, boolean setupACLs) throws Exception {

  GetNewApplicationRequest newAppRequest =
      GetNewApplicationRequest.newInstance();

  ApplicationClientProtocol submitterClient = getRMClientForUser(submitter);
  ApplicationId applicationId =
      submitterClient.getNewApplication(newAppRequest).getApplicationId();

  Resource resource = BuilderUtils.newResource(1024, 1);
  Map<ApplicationAccessType, String> acls = createACLs(submitter, setupACLs);
  ContainerLaunchContext amContainerSpec =
      ContainerLaunchContext.newInstance(null, null, null, null, null, acls);

  ApplicationSubmissionContext appSubmissionContext =
      ApplicationSubmissionContext.newInstance(applicationId,
        "applicationName", queueName, null, amContainerSpec, false, true, 1,
        resource, "applicationType");
  appSubmissionContext.setApplicationId(applicationId);
  appSubmissionContext.setQueue(queueName);

  SubmitApplicationRequest submitRequest =
      SubmitApplicationRequest.newInstance(appSubmissionContext);
  submitterClient.submitApplication(submitRequest);
  resourceManager.waitForState(applicationId, RMAppState.ACCEPTED);
  return applicationId;
}
 

@SuppressWarnings("deprecation")
private ApplicationId submitAppAndGetAppId(AccessControlList viewACL,
    AccessControlList modifyACL) throws Exception {
  SubmitApplicationRequest submitRequest = recordFactory
      .newRecordInstance(SubmitApplicationRequest.class);
  ApplicationSubmissionContext context = recordFactory
      .newRecordInstance(ApplicationSubmissionContext.class);

  ApplicationId applicationId = rmClient.getNewApplication(
      recordFactory.newRecordInstance(GetNewApplicationRequest.class))
      .getApplicationId();
  context.setApplicationId(applicationId);

  Map<ApplicationAccessType, String> acls
      = new HashMap<ApplicationAccessType, String>();
  acls.put(ApplicationAccessType.VIEW_APP, viewACL.getAclString());
  acls.put(ApplicationAccessType.MODIFY_APP, modifyACL.getAclString());

  ContainerLaunchContext amContainer = recordFactory
      .newRecordInstance(ContainerLaunchContext.class);
  Resource resource = BuilderUtils.newResource(1024, 1);
  context.setResource(resource);
  amContainer.setApplicationACLs(acls);
  context.setAMContainerSpec(amContainer);
  submitRequest.setApplicationSubmissionContext(context);
  rmClient.submitApplication(submitRequest);
  resourceManager.waitForState(applicationId, RMAppState.ACCEPTED);
  return applicationId;
}
 

@SuppressWarnings("unchecked")
private Container allocateApplicationAttempt() {
  scheduleApplicationAttempt();
  
  // Mock the allocation of AM container 
  Container container = mock(Container.class);
  Resource resource = BuilderUtils.newResource(2048, 1, 1);
  when(container.getId()).thenReturn(
      BuilderUtils.newContainerId(applicationAttempt.getAppAttemptId(), 1));
  when(container.getResource()).thenReturn(resource);
  Allocation allocation = mock(Allocation.class);
  when(allocation.getContainers()).
      thenReturn(Collections.singletonList(container));
  when(
      scheduler.allocate(
          any(ApplicationAttemptId.class), 
          any(List.class), 
          any(List.class), 
          any(List.class), 
          any(List.class))).
  thenReturn(allocation);
  RMContainer rmContainer = mock(RMContainerImpl.class);
  when(scheduler.getRMContainer(container.getId())).
      thenReturn(rmContainer);
  
  applicationAttempt.handle(
      new RMAppAttemptContainerAllocatedEvent(
          applicationAttempt.getAppAttemptId()));
  
  assertEquals(RMAppAttemptState.ALLOCATED_SAVING, 
      applicationAttempt.getAppAttemptState());
  applicationAttempt.handle(
      new RMAppAttemptEvent(applicationAttempt.getAppAttemptId(),
          RMAppAttemptEventType.ATTEMPT_NEW_SAVED));
  
  testAppAttemptAllocatedState(container);
  
  return container;
}
 

private Schedulable createSchedulable(int memUsage, int cpuUsage,
    ResourceWeights weights, int minMemShare, int minCpuShare) {
  Resource usage = BuilderUtils.newResource(memUsage, cpuUsage);
  Resource minShare = BuilderUtils.newResource(minMemShare, minCpuShare);
  return new FakeSchedulable(minShare,
      Resources.createResource(Integer.MAX_VALUE, Integer.MAX_VALUE),
      weights, Resources.none(), usage, 0l);
}
 

@Test
public void testReconnect() throws Exception {
  String hostname1 = "localhost1";
  Resource capability = BuilderUtils.newResource(1024, 1, 1);

  RegisterNodeManagerRequest request1 = recordFactory
      .newRecordInstance(RegisterNodeManagerRequest.class);
  NodeId nodeId1 = NodeId.newInstance(hostname1, 0);
  request1.setNodeId(nodeId1);
  request1.setHttpPort(0);
  request1.setResource(capability);
  resourceTrackerService.registerNodeManager(request1);

  Assert.assertEquals(RMNodeEventType.STARTED, rmNodeEvents.get(0).getType());

  rmNodeEvents.clear();
  resourceTrackerService.registerNodeManager(request1);
  Assert.assertEquals(RMNodeEventType.RECONNECTED,
      rmNodeEvents.get(0).getType());

  rmNodeEvents.clear();
  resourceTrackerService.registerNodeManager(request1);
  capability = BuilderUtils.newResource(1024, 2, 2);
  request1.setResource(capability);
  Assert.assertEquals(RMNodeEventType.RECONNECTED,
      rmNodeEvents.get(0).getType());
}
 

@Test
public void testRPCResponseId() throws IOException, YarnException {
  String node = "localhost";
  Resource capability = BuilderUtils.newResource(1024, 1);
  RegisterNodeManagerRequest request = recordFactory.newRecordInstance(RegisterNodeManagerRequest.class);
  nodeId = NodeId.newInstance(node, 1234);
  request.setNodeId(nodeId);
  request.setHttpPort(0);
  request.setResource(capability);

  RegisterNodeManagerRequest request1 = recordFactory
      .newRecordInstance(RegisterNodeManagerRequest.class);
  request1.setNodeId(nodeId);
  request1.setHttpPort(0);
  request1.setResource(capability);
  resourceTrackerService.registerNodeManager(request1);

  org.apache.hadoop.yarn.server.api.records.NodeStatus nodeStatus = recordFactory.
    newRecordInstance(org.apache.hadoop.yarn.server.api.records.NodeStatus.class);
  nodeStatus.setNodeId(nodeId);
  NodeHealthStatus nodeHealthStatus = recordFactory.newRecordInstance(NodeHealthStatus.class);
  nodeHealthStatus.setIsNodeHealthy(true);
  nodeStatus.setNodeHealthStatus(nodeHealthStatus);
  NodeHeartbeatRequest nodeHeartBeatRequest = recordFactory
      .newRecordInstance(NodeHeartbeatRequest.class);
  nodeHeartBeatRequest.setNodeStatus(nodeStatus);

  nodeStatus.setResponseId(0);
  NodeHeartbeatResponse response = resourceTrackerService.nodeHeartbeat(
      nodeHeartBeatRequest);
  Assert.assertTrue(response.getResponseId() == 1);

  nodeStatus.setResponseId(response.getResponseId());
  response = resourceTrackerService.nodeHeartbeat(nodeHeartBeatRequest);
  Assert.assertTrue(response.getResponseId() == 2);   

  /* try calling with less response id */
  response = resourceTrackerService.nodeHeartbeat(nodeHeartBeatRequest);
  Assert.assertTrue(response.getResponseId() == 2);

  nodeStatus.setResponseId(0);
  response = resourceTrackerService.nodeHeartbeat(nodeHeartBeatRequest);
  Assert.assertTrue(NodeAction.RESYNC.equals(response.getNodeAction()));
  Assert.assertEquals("Too far behind rm response id:2 nm response id:0",
    response.getDiagnosticsMessage());
}
 

@SuppressWarnings("deprecation")
@Before
public void setUp() throws Exception {
  AuthenticationMethod authMethod = AuthenticationMethod.SIMPLE;
  if (isSecurityEnabled) {
    authMethod = AuthenticationMethod.KERBEROS;
  }
  SecurityUtil.setAuthenticationMethod(authMethod, conf);
  UserGroupInformation.setConfiguration(conf);
  InlineDispatcher rmDispatcher = new InlineDispatcher();

  ContainerAllocationExpirer containerAllocationExpirer =
      mock(ContainerAllocationExpirer.class);
  amLivelinessMonitor = mock(AMLivelinessMonitor.class);
  amFinishingMonitor = mock(AMLivelinessMonitor.class);
  writer = mock(RMApplicationHistoryWriter.class);
  MasterKeyData masterKeyData = amRMTokenManager.createNewMasterKey();
  when(amRMTokenManager.getMasterKey()).thenReturn(masterKeyData);
  rmContext =
      new RMContextImpl(rmDispatcher,
        containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
        null, amRMTokenManager,
        new RMContainerTokenSecretManager(conf),
        nmTokenManager,
        clientToAMTokenManager,
        writer);
  
  store = mock(RMStateStore.class);
  ((RMContextImpl) rmContext).setStateStore(store);
  publisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl) rmContext).setSystemMetricsPublisher(publisher);
  
  scheduler = mock(YarnScheduler.class);
  masterService = mock(ApplicationMasterService.class);
  applicationMasterLauncher = mock(ApplicationMasterLauncher.class);
  
  rmDispatcher.register(RMAppAttemptEventType.class,
      new TestApplicationAttemptEventDispatcher());

  rmDispatcher.register(RMAppEventType.class,
      new TestApplicationEventDispatcher());
  
  rmDispatcher.register(SchedulerEventType.class, 
      new TestSchedulerEventDispatcher());
  
  rmDispatcher.register(AMLauncherEventType.class, 
      new TestAMLauncherEventDispatcher());

  rmnodeEventHandler = mock(RMNodeImpl.class);
  rmDispatcher.register(RMNodeEventType.class, rmnodeEventHandler);

  rmDispatcher.init(conf);
  rmDispatcher.start();
  

  ApplicationId applicationId = MockApps.newAppID(appId++);
  ApplicationAttemptId applicationAttemptId =
      ApplicationAttemptId.newInstance(applicationId, 0);

  resourceScheduler = mock(ResourceScheduler.class);

  ApplicationResourceUsageReport appResUsgRpt =
      mock(ApplicationResourceUsageReport.class);
  when(appResUsgRpt.getMemorySeconds()).thenReturn(0L);
  when(appResUsgRpt.getVcoreSeconds()).thenReturn(0L);
  when(resourceScheduler
      .getAppResourceUsageReport((ApplicationAttemptId)Matchers.any()))
   .thenReturn(appResUsgRpt);
  spyRMContext = spy(rmContext);
  Mockito.doReturn(resourceScheduler).when(spyRMContext).getScheduler();


  final String user = MockApps.newUserName();
  final String queue = MockApps.newQueue();
  submissionContext = mock(ApplicationSubmissionContext.class);
  when(submissionContext.getQueue()).thenReturn(queue);
  Resource resource = BuilderUtils.newResource(1536, 1);
  ContainerLaunchContext amContainerSpec =
      BuilderUtils.newContainerLaunchContext(null, null,
          null, null, null, null);
  when(submissionContext.getAMContainerSpec()).thenReturn(amContainerSpec);
  when(submissionContext.getResource()).thenReturn(resource);

  unmanagedAM = false;
  
  application = mock(RMAppImpl.class);
  applicationAttempt =
      new RMAppAttemptImpl(applicationAttemptId, spyRMContext, scheduler,
          masterService, submissionContext, new Configuration(), false,
          BuilderUtils.newResourceRequest(
              RMAppAttemptImpl.AM_CONTAINER_PRIORITY, ResourceRequest.ANY,
              submissionContext.getResource(), 1));

  when(application.getCurrentAppAttempt()).thenReturn(applicationAttempt);
  when(application.getApplicationId()).thenReturn(applicationId);
  spyRMContext.getRMApps().put(application.getApplicationId(), application);

  testAppAttemptNewState();
}
 

@Test
public void testNMExpiry() throws Exception {
  String hostname1 = "localhost1";
  String hostname2 = "localhost2";
  String hostname3 = "localhost3";
  Resource capability = BuilderUtils.newResource(1024, 1, 1);

  RegisterNodeManagerRequest request1 = recordFactory
      .newRecordInstance(RegisterNodeManagerRequest.class);
  NodeId nodeId1 = NodeId.newInstance(hostname1, 0);
  request1.setNodeId(nodeId1);
  request1.setHttpPort(0);
  request1.setResource(capability);
  resourceTrackerService.registerNodeManager(request1);

  RegisterNodeManagerRequest request2 = recordFactory
      .newRecordInstance(RegisterNodeManagerRequest.class);
  NodeId nodeId2 = NodeId.newInstance(hostname2, 0);
  request2.setNodeId(nodeId2);
  request2.setHttpPort(0);
  request2.setResource(capability);
  resourceTrackerService.registerNodeManager(request2);
  
  int waitCount = 0;
  while(ClusterMetrics.getMetrics().getNumLostNMs()!=2 && waitCount ++<20){
    synchronized (this) {
      wait(100);
    }
  }
  Assert.assertEquals(2, ClusterMetrics.getMetrics().getNumLostNMs());

  request3 = recordFactory
      .newRecordInstance(RegisterNodeManagerRequest.class);
  NodeId nodeId3 = NodeId.newInstance(hostname3, 0);
  request3.setNodeId(nodeId3);
  request3.setHttpPort(0);
  request3.setResource(capability);
  resourceTrackerService
      .registerNodeManager(request3);

  /* test to see if hostanme 3 does not expire */
  stopT = false;
  new ThirdNodeHeartBeatThread().start();
  Assert.assertEquals(2,ClusterMetrics.getMetrics().getNumLostNMs());
  stopT = true;
}
 

@Test
public void testRPCResponseId() throws IOException, YarnException {
  String node = "localhost";
  Resource capability = BuilderUtils.newResource(1024, 1, 1);
  RegisterNodeManagerRequest request = recordFactory.newRecordInstance(RegisterNodeManagerRequest.class);
  nodeId = NodeId.newInstance(node, 1234);
  request.setNodeId(nodeId);
  request.setHttpPort(0);
  request.setResource(capability);

  RegisterNodeManagerRequest request1 = recordFactory
      .newRecordInstance(RegisterNodeManagerRequest.class);
  request1.setNodeId(nodeId);
  request1.setHttpPort(0);
  request1.setResource(capability);
  resourceTrackerService.registerNodeManager(request1);

  org.apache.hadoop.yarn.server.api.records.NodeStatus nodeStatus = recordFactory.
    newRecordInstance(org.apache.hadoop.yarn.server.api.records.NodeStatus.class);
  nodeStatus.setNodeId(nodeId);
  NodeHealthStatus nodeHealthStatus = recordFactory.newRecordInstance(NodeHealthStatus.class);
  nodeHealthStatus.setIsNodeHealthy(true);
  nodeStatus.setNodeHealthStatus(nodeHealthStatus);
  NodeHeartbeatRequest nodeHeartBeatRequest = recordFactory
      .newRecordInstance(NodeHeartbeatRequest.class);
  nodeHeartBeatRequest.setNodeStatus(nodeStatus);

  nodeStatus.setResponseId(0);
  NodeHeartbeatResponse response = resourceTrackerService.nodeHeartbeat(
      nodeHeartBeatRequest);
  Assert.assertTrue(response.getResponseId() == 1);

  nodeStatus.setResponseId(response.getResponseId());
  response = resourceTrackerService.nodeHeartbeat(nodeHeartBeatRequest);
  Assert.assertTrue(response.getResponseId() == 2);   

  /* try calling with less response id */
  response = resourceTrackerService.nodeHeartbeat(nodeHeartBeatRequest);
  Assert.assertTrue(response.getResponseId() == 2);

  nodeStatus.setResponseId(0);
  response = resourceTrackerService.nodeHeartbeat(nodeHeartBeatRequest);
  Assert.assertTrue(NodeAction.RESYNC.equals(response.getNodeAction()));
  Assert.assertEquals("Too far behind rm response id:2 nm response id:0",
    response.getDiagnosticsMessage());
}
 

@SuppressWarnings("deprecation")
@Before
public void setUp() throws Exception {
  AuthenticationMethod authMethod = AuthenticationMethod.SIMPLE;
  if (isSecurityEnabled) {
    authMethod = AuthenticationMethod.KERBEROS;
  }
  SecurityUtil.setAuthenticationMethod(authMethod, conf);
  UserGroupInformation.setConfiguration(conf);
  InlineDispatcher rmDispatcher = new InlineDispatcher();

  ContainerAllocationExpirer containerAllocationExpirer =
      mock(ContainerAllocationExpirer.class);
  amLivelinessMonitor = mock(AMLivelinessMonitor.class);
  amFinishingMonitor = mock(AMLivelinessMonitor.class);
  writer = mock(RMApplicationHistoryWriter.class);
  MasterKeyData masterKeyData = amRMTokenManager.createNewMasterKey();
  when(amRMTokenManager.getMasterKey()).thenReturn(masterKeyData);
  rmContext =
      new RMContextImpl(rmDispatcher,
        containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
        null, amRMTokenManager,
        new RMContainerTokenSecretManager(conf),
        nmTokenManager,
        clientToAMTokenManager,
        writer);
  
  store = mock(RMStateStore.class);
  ((RMContextImpl) rmContext).setStateStore(store);
  publisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl) rmContext).setSystemMetricsPublisher(publisher);
  
  scheduler = mock(YarnScheduler.class);
  masterService = mock(ApplicationMasterService.class);
  applicationMasterLauncher = mock(ApplicationMasterLauncher.class);
  
  rmDispatcher.register(RMAppAttemptEventType.class,
      new TestApplicationAttemptEventDispatcher());

  rmDispatcher.register(RMAppEventType.class,
      new TestApplicationEventDispatcher());
  
  rmDispatcher.register(SchedulerEventType.class, 
      new TestSchedulerEventDispatcher());
  
  rmDispatcher.register(AMLauncherEventType.class, 
      new TestAMLauncherEventDispatcher());

  rmnodeEventHandler = mock(RMNodeImpl.class);
  rmDispatcher.register(RMNodeEventType.class, rmnodeEventHandler);

  rmDispatcher.init(conf);
  rmDispatcher.start();
  

  ApplicationId applicationId = MockApps.newAppID(appId++);
  ApplicationAttemptId applicationAttemptId =
      ApplicationAttemptId.newInstance(applicationId, 0);

  resourceScheduler = mock(ResourceScheduler.class);

  ApplicationResourceUsageReport appResUsgRpt =
      mock(ApplicationResourceUsageReport.class);
  when(appResUsgRpt.getMemorySeconds()).thenReturn(0L);
  when(appResUsgRpt.getVcoreSeconds()).thenReturn(0L);
  when(appResUsgRpt.getGcoreSeconds()).thenReturn(0L);
  when(resourceScheduler
      .getAppResourceUsageReport((ApplicationAttemptId)Matchers.any()))
   .thenReturn(appResUsgRpt);
  spyRMContext = spy(rmContext);
  Mockito.doReturn(resourceScheduler).when(spyRMContext).getScheduler();


  final String user = MockApps.newUserName();
  final String queue = MockApps.newQueue();
  submissionContext = mock(ApplicationSubmissionContext.class);
  when(submissionContext.getQueue()).thenReturn(queue);
  Resource resource = BuilderUtils.newResource(1536, 1, 1);
  ContainerLaunchContext amContainerSpec =
      BuilderUtils.newContainerLaunchContext(null, null,
          null, null, null, null);
  when(submissionContext.getAMContainerSpec()).thenReturn(amContainerSpec);
  when(submissionContext.getResource()).thenReturn(resource);

  unmanagedAM = false;
  
  application = mock(RMAppImpl.class);
  applicationAttempt =
      new RMAppAttemptImpl(applicationAttemptId, spyRMContext, scheduler,
          masterService, submissionContext, new Configuration(), false,
          BuilderUtils.newResourceRequest(
              RMAppAttemptImpl.AM_CONTAINER_PRIORITY, ResourceRequest.ANY,
              submissionContext.getResource(), 1));

  when(application.getCurrentAppAttempt()).thenReturn(applicationAttempt);
  when(application.getApplicationId()).thenReturn(applicationId);
  spyRMContext.getRMApps().put(application.getApplicationId(), application);

  testAppAttemptNewState();
}
 

@Test
public void testReleaseWhileRunning() {

  DrainDispatcher drainDispatcher = new DrainDispatcher();
  EventHandler<RMAppAttemptEvent> appAttemptEventHandler = mock(EventHandler.class);
  EventHandler generic = mock(EventHandler.class);
  drainDispatcher.register(RMAppAttemptEventType.class,
      appAttemptEventHandler);
  drainDispatcher.register(RMNodeEventType.class, generic);
  drainDispatcher.init(new YarnConfiguration());
  drainDispatcher.start();
  NodeId nodeId = BuilderUtils.newNodeId("host", 3425);
  ApplicationId appId = BuilderUtils.newApplicationId(1, 1);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      appId, 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  ContainerAllocationExpirer expirer = mock(ContainerAllocationExpirer.class);

  Resource resource = BuilderUtils.newResource(512, 1);
  Priority priority = BuilderUtils.newPriority(5);

  Container container = BuilderUtils.newContainer(containerId, nodeId,
      "host:3465", resource, priority, null);
  ConcurrentMap<ApplicationId, RMApp> rmApps =
      spy(new ConcurrentHashMap<ApplicationId, RMApp>());
  RMApp rmApp = mock(RMApp.class);
  when(rmApp.getRMAppAttempt((ApplicationAttemptId)Matchers.any())).thenReturn(null);
  Mockito.doReturn(rmApp).when(rmApps).get((ApplicationId)Matchers.any());

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getRMApps()).thenReturn(rmApps);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  when(rmContext.getYarnConfiguration()).thenReturn(new YarnConfiguration());
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      nodeId, "user", rmContext);

  assertEquals(RMContainerState.NEW, rmContainer.getState());
  assertEquals(resource, rmContainer.getAllocatedResource());
  assertEquals(nodeId, rmContainer.getAllocatedNode());
  assertEquals(priority, rmContainer.getAllocatedPriority());
  verify(writer).containerStarted(any(RMContainer.class));
  verify(publisher).containerCreated(any(RMContainer.class), anyLong());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.START));
  drainDispatcher.await();
  assertEquals(RMContainerState.ALLOCATED, rmContainer.getState());
  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.ACQUIRED));
  drainDispatcher.await();
  assertEquals(RMContainerState.ACQUIRED, rmContainer.getState());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.LAUNCHED));
  drainDispatcher.await();
  assertEquals(RMContainerState.RUNNING, rmContainer.getState());
  assertEquals("http://host:3465/node/containerlogs/container_1_0001_01_000001/user",
      rmContainer.getLogURL());

  // In RUNNING state. Verify RELEASED and associated actions.
  reset(appAttemptEventHandler);
  ContainerStatus containerStatus = SchedulerUtils
      .createAbnormalContainerStatus(containerId,
          SchedulerUtils.RELEASED_CONTAINER);
  rmContainer.handle(new RMContainerFinishedEvent(containerId,
      containerStatus, RMContainerEventType.RELEASED));
  drainDispatcher.await();
  assertEquals(RMContainerState.RELEASED, rmContainer.getState());
  assertEquals(SchedulerUtils.RELEASED_CONTAINER,
      rmContainer.getDiagnosticsInfo());
  assertEquals(ContainerExitStatus.ABORTED,
      rmContainer.getContainerExitStatus());
  assertEquals(ContainerState.COMPLETE, rmContainer.getContainerState());
  verify(writer).containerFinished(any(RMContainer.class));
  verify(publisher).containerFinished(any(RMContainer.class), anyLong());

  ArgumentCaptor<RMAppAttemptContainerFinishedEvent> captor = ArgumentCaptor
      .forClass(RMAppAttemptContainerFinishedEvent.class);
  verify(appAttemptEventHandler).handle(captor.capture());
  RMAppAttemptContainerFinishedEvent cfEvent = captor.getValue();
  assertEquals(appAttemptId, cfEvent.getApplicationAttemptId());
  assertEquals(containerStatus, cfEvent.getContainerStatus());
  assertEquals(RMAppAttemptEventType.CONTAINER_FINISHED, cfEvent.getType());
  
  // In RELEASED state. A FINIHSED event may come in.
  rmContainer.handle(new RMContainerFinishedEvent(containerId, SchedulerUtils
      .createAbnormalContainerStatus(containerId, "FinishedContainer"),
      RMContainerEventType.FINISHED));
  assertEquals(RMContainerState.RELEASED, rmContainer.getState());
}
 

@Test
public void testNodeRegistrationWithMinimumAllocations() throws Exception {
  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB, "2048");
  conf.set(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES, "4");
  rm = new MockRM(conf);
  rm.start();

  ResourceTrackerService resourceTrackerService
    = rm.getResourceTrackerService();
  RegisterNodeManagerRequest req = Records.newRecord(
      RegisterNodeManagerRequest.class);
  NodeId nodeId = BuilderUtils.newNodeId("host", 1234);
  req.setNodeId(nodeId);

  Resource capability = BuilderUtils.newResource(1024, 1);
  req.setResource(capability);
  RegisterNodeManagerResponse response1 =
      resourceTrackerService.registerNodeManager(req);
  Assert.assertEquals(NodeAction.SHUTDOWN,response1.getNodeAction());
  
  capability.setMemory(2048);
  capability.setVirtualCores(1);
  req.setResource(capability);
  RegisterNodeManagerResponse response2 =
      resourceTrackerService.registerNodeManager(req);
  Assert.assertEquals(NodeAction.SHUTDOWN,response2.getNodeAction());
  
  capability.setMemory(1024);
  capability.setVirtualCores(4);
  req.setResource(capability);
  RegisterNodeManagerResponse response3 =
      resourceTrackerService.registerNodeManager(req);
  Assert.assertEquals(NodeAction.SHUTDOWN,response3.getNodeAction());
  
  capability.setMemory(2048);
  capability.setVirtualCores(4);
  req.setResource(capability);
  RegisterNodeManagerResponse response4 =
      resourceTrackerService.registerNodeManager(req);
  Assert.assertEquals(NodeAction.NORMAL,response4.getNodeAction());
}