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

@Before
@SuppressWarnings("unchecked")
public void setup() {
  mockDelService = mock(DeletionService.class);
  conf = new YarnConfiguration();
  dispatcher = createDispatcher(conf);
  appEventHandler = mock(EventHandler.class);
  dispatcher.register(ApplicationEventType.class, appEventHandler);
  appId = BuilderUtils.newApplicationId(1234, 1);
  appAttemptId = BuilderUtils.newApplicationAttemptId(appId, 1);
  container11 = BuilderUtils.newContainerId(appAttemptId, 1);
  dirsHandler = new LocalDirsHandlerService();
}
 

/**
 * Create and return a container object reflecting an allocation for the
 * given appliction on the given node with the given capability and
 * priority.
 */
public Container createContainer(
    FSSchedulerNode node, Resource capability, Priority priority) {

  NodeId nodeId = node.getRMNode().getNodeID();
  ContainerId containerId = BuilderUtils.newContainerId(
      getApplicationAttemptId(), getNewContainerId());

  // Create the container
  Container container =
      BuilderUtils.newContainer(containerId, nodeId, node.getRMNode()
          .getHttpAddress(), capability, priority, null);

  return container;
}
 

private Container createMockedContainer(ApplicationId appId, int containerId) {
  ApplicationAttemptId appAttemptId =
      BuilderUtils.newApplicationAttemptId(appId, 1);
  ContainerId cId = BuilderUtils.newContainerId(appAttemptId, containerId);
  Container c = mock(Container.class);
  when(c.getContainerId()).thenReturn(cId);
  ContainerLaunchContext launchContext = mock(ContainerLaunchContext.class);
  when(c.getLaunchContext()).thenReturn(launchContext);
  when(launchContext.getApplicationACLs()).thenReturn(
      new HashMap<ApplicationAccessType, String>());
  return c;
}
 

@Test (timeout = 5000)
public void testStatusChange(){
  //Start the node
  node.handle(new RMNodeStartedEvent(null, null, null));
  //Add info to the queue first
  node.setNextHeartBeat(false);

  ContainerId completedContainerId1 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(0, 0), 0), 0);
  ContainerId completedContainerId2 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(1, 1), 1), 1);
      
  RMNodeStatusEvent statusEvent1 = getMockRMNodeStatusEvent();
  RMNodeStatusEvent statusEvent2 = getMockRMNodeStatusEvent();

  ContainerStatus containerStatus1 = mock(ContainerStatus.class);
  ContainerStatus containerStatus2 = mock(ContainerStatus.class);

  doReturn(completedContainerId1).when(containerStatus1).getContainerId();
  doReturn(Collections.singletonList(containerStatus1))
      .when(statusEvent1).getContainers();
   
  doReturn(completedContainerId2).when(containerStatus2).getContainerId();
  doReturn(Collections.singletonList(containerStatus2))
      .when(statusEvent2).getContainers();

  verify(scheduler,times(1)).handle(any(NodeUpdateSchedulerEvent.class)); 
  node.handle(statusEvent1);
  node.handle(statusEvent2);
  verify(scheduler,times(1)).handle(any(NodeUpdateSchedulerEvent.class));
  Assert.assertEquals(2, node.getQueueSize());
  node.handle(new RMNodeEvent(node.getNodeID(), RMNodeEventType.EXPIRE));
  Assert.assertEquals(0, node.getQueueSize());
}
 

@Test(timeout=20000)
public void testUpdateHeartbeatResponseForCleanup() {
  RMNodeImpl node = getRunningNode();
  NodeId nodeId = node.getNodeID();

  // Expire a container
ContainerId completedContainerId = BuilderUtils.newContainerId(
		BuilderUtils.newApplicationAttemptId(
				BuilderUtils.newApplicationId(0, 0), 0), 0);
  node.handle(new RMNodeCleanContainerEvent(nodeId, completedContainerId));
  Assert.assertEquals(1, node.getContainersToCleanUp().size());

  // Finish an application
  ApplicationId finishedAppId = BuilderUtils.newApplicationId(0, 1);
  node.handle(new RMNodeCleanAppEvent(nodeId, finishedAppId));
  Assert.assertEquals(1, node.getAppsToCleanup().size());

  // Verify status update does not clear containers/apps to cleanup
  // but updating heartbeat response for cleanup does
  RMNodeStatusEvent statusEvent = getMockRMNodeStatusEvent();
  node.handle(statusEvent);
  Assert.assertEquals(1, node.getContainersToCleanUp().size());
  Assert.assertEquals(1, node.getAppsToCleanup().size());
  NodeHeartbeatResponse hbrsp = Records.newRecord(NodeHeartbeatResponse.class);
  node.updateNodeHeartbeatResponseForCleanup(hbrsp);
  Assert.assertEquals(0, node.getContainersToCleanUp().size());
  Assert.assertEquals(0, node.getAppsToCleanup().size());
  Assert.assertEquals(1, hbrsp.getContainersToCleanup().size());
  Assert.assertEquals(completedContainerId, hbrsp.getContainersToCleanup().get(0));
  Assert.assertEquals(1, hbrsp.getApplicationsToCleanup().size());
  Assert.assertEquals(finishedAppId, hbrsp.getApplicationsToCleanup().get(0));
}
 

@Test
@SuppressWarnings("unchecked")
public void testStateStoreFailedLocalization() throws Exception {
  final String user = "someuser";
  final ApplicationId appId = ApplicationId.newInstance(1, 1);
  // This is a random path. NO File creation will take place at this place.
  final Path localDir = new Path("/tmp");
  Configuration conf = new YarnConfiguration();
  DrainDispatcher dispatcher = null;
  dispatcher = createDispatcher(conf);
  EventHandler<LocalizerEvent> localizerEventHandler =
      mock(EventHandler.class);
  EventHandler<LocalizerEvent> containerEventHandler =
      mock(EventHandler.class);
  dispatcher.register(LocalizerEventType.class, localizerEventHandler);
  dispatcher.register(ContainerEventType.class, containerEventHandler);
  NMStateStoreService stateStore = mock(NMStateStoreService.class);

  try {
    LocalResourcesTracker tracker = new LocalResourcesTrackerImpl(user,
        appId, dispatcher, false, conf, stateStore);
    // Container 1 needs lr1 resource
    ContainerId cId1 = BuilderUtils.newContainerId(1, 1, 1, 1);
    LocalResourceRequest lr1 = createLocalResourceRequest(user, 1, 1,
        LocalResourceVisibility.APPLICATION);
    LocalizerContext lc1 = new LocalizerContext(user, cId1, null);

    // Container 1 requests lr1 to be localized
    ResourceEvent reqEvent1 = new ResourceRequestEvent(lr1,
        LocalResourceVisibility.APPLICATION, lc1);
    tracker.handle(reqEvent1);
    dispatcher.await();

    // Simulate the process of localization of lr1
    Path hierarchicalPath1 = tracker.getPathForLocalization(lr1, localDir);

    ArgumentCaptor<LocalResourceProto> localResourceCaptor =
        ArgumentCaptor.forClass(LocalResourceProto.class);
    ArgumentCaptor<Path> pathCaptor = ArgumentCaptor.forClass(Path.class);
    verify(stateStore).startResourceLocalization(eq(user), eq(appId),
        localResourceCaptor.capture(), pathCaptor.capture());
    LocalResourceProto lrProto = localResourceCaptor.getValue();
    Path localizedPath1 = pathCaptor.getValue();
    Assert.assertEquals(lr1,
        new LocalResourceRequest(new LocalResourcePBImpl(lrProto)));
    Assert.assertEquals(hierarchicalPath1, localizedPath1.getParent());

    ResourceFailedLocalizationEvent rfe1 =
        new ResourceFailedLocalizationEvent(
            lr1, new Exception("Test").toString());
    tracker.handle(rfe1);
    dispatcher.await();
    verify(stateStore).removeLocalizedResource(eq(user), eq(appId),
        eq(localizedPath1));
  } finally {
    if (dispatcher != null) {
      dispatcher.stop();
    }
  }
}
 

@Test
public void testContainerTokenStorage() throws IOException {
  // test empty when no state
  RecoveredContainerTokensState state =
      stateStore.loadContainerTokensState();
  assertNull(state.getCurrentMasterKey());
  assertNull(state.getPreviousMasterKey());
  assertTrue(state.getActiveTokens().isEmpty());

  // store a master key and verify recovered
  ContainerTokenKeyGeneratorForTest keygen =
      new ContainerTokenKeyGeneratorForTest(new YarnConfiguration());
  MasterKey currentKey = keygen.generateKey();
  stateStore.storeContainerTokenCurrentMasterKey(currentKey);
  restartStateStore();
  state = stateStore.loadContainerTokensState();
  assertEquals(currentKey, state.getCurrentMasterKey());
  assertNull(state.getPreviousMasterKey());
  assertTrue(state.getActiveTokens().isEmpty());

  // store a previous key and verify recovered
  MasterKey prevKey = keygen.generateKey();
  stateStore.storeContainerTokenPreviousMasterKey(prevKey);
  restartStateStore();
  state = stateStore.loadContainerTokensState();
  assertEquals(currentKey, state.getCurrentMasterKey());
  assertEquals(prevKey, state.getPreviousMasterKey());
  assertTrue(state.getActiveTokens().isEmpty());

  // store a few container tokens and verify recovered
  ContainerId cid1 = BuilderUtils.newContainerId(1, 1, 1, 1);
  Long expTime1 = 1234567890L;
  ContainerId cid2 = BuilderUtils.newContainerId(2, 2, 2, 2);
  Long expTime2 = 9876543210L;
  stateStore.storeContainerToken(cid1, expTime1);
  stateStore.storeContainerToken(cid2, expTime2);
  restartStateStore();
  state = stateStore.loadContainerTokensState();
  assertEquals(currentKey, state.getCurrentMasterKey());
  assertEquals(prevKey, state.getPreviousMasterKey());
  Map<ContainerId, Long> loadedActiveTokens =
      state.getActiveTokens();
  assertEquals(2, loadedActiveTokens.size());
  assertEquals(expTime1, loadedActiveTokens.get(cid1));
  assertEquals(expTime2, loadedActiveTokens.get(cid2));

  // add/update/remove tokens and verify recovered
  ContainerId cid3 = BuilderUtils.newContainerId(3, 3, 3, 3);
  Long expTime3 = 135798642L;
  stateStore.storeContainerToken(cid3, expTime3);
  stateStore.removeContainerToken(cid1);
  expTime2 += 246897531L;
  stateStore.storeContainerToken(cid2, expTime2);
  prevKey = currentKey;
  stateStore.storeContainerTokenPreviousMasterKey(prevKey);
  currentKey = keygen.generateKey();
  stateStore.storeContainerTokenCurrentMasterKey(currentKey);
  restartStateStore();
  state = stateStore.loadContainerTokensState();
  assertEquals(currentKey, state.getCurrentMasterKey());
  assertEquals(prevKey, state.getPreviousMasterKey());
  loadedActiveTokens = state.getActiveTokens();
  assertEquals(2, loadedActiveTokens.size());
  assertNull(loadedActiveTokens.get(cid1));
  assertEquals(expTime2, loadedActiveTokens.get(cid2));
  assertEquals(expTime3, loadedActiveTokens.get(cid3));
}
 

@Test (timeout = 5000)
public void testContainerUpdate() throws InterruptedException{
  //Start the node
  node.handle(new RMNodeStartedEvent(null, null, null));
  
  NodeId nodeId = BuilderUtils.newNodeId("localhost:1", 1);
  RMNodeImpl node2 = new RMNodeImpl(nodeId, rmContext, null, 0, 0, null, null, null);
  node2.handle(new RMNodeStartedEvent(null, null, null));
  
  ContainerId completedContainerIdFromNode1 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(0, 0), 0), 0);
  ContainerId completedContainerIdFromNode2_1 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(1, 1), 1), 1);
  ContainerId completedContainerIdFromNode2_2 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(1, 1), 1), 2);
 
  RMNodeStatusEvent statusEventFromNode1 = getMockRMNodeStatusEvent();
  RMNodeStatusEvent statusEventFromNode2_1 = getMockRMNodeStatusEvent();
  RMNodeStatusEvent statusEventFromNode2_2 = getMockRMNodeStatusEvent();
  
  ContainerStatus containerStatusFromNode1 = mock(ContainerStatus.class);
  ContainerStatus containerStatusFromNode2_1 = mock(ContainerStatus.class);
  ContainerStatus containerStatusFromNode2_2 = mock(ContainerStatus.class);

  doReturn(completedContainerIdFromNode1).when(containerStatusFromNode1)
      .getContainerId();
  doReturn(Collections.singletonList(containerStatusFromNode1))
      .when(statusEventFromNode1).getContainers();
  node.handle(statusEventFromNode1);
  Assert.assertEquals(1, completedContainers.size());
  Assert.assertEquals(completedContainerIdFromNode1,
      completedContainers.get(0).getContainerId());

  completedContainers.clear();

  doReturn(completedContainerIdFromNode2_1).when(containerStatusFromNode2_1)
      .getContainerId();
  doReturn(Collections.singletonList(containerStatusFromNode2_1))
      .when(statusEventFromNode2_1).getContainers();

  doReturn(completedContainerIdFromNode2_2).when(containerStatusFromNode2_2)
      .getContainerId();
  doReturn(Collections.singletonList(containerStatusFromNode2_2))
      .when(statusEventFromNode2_2).getContainers();

  node2.setNextHeartBeat(false);
  node2.handle(statusEventFromNode2_1);
  node2.setNextHeartBeat(true);
  node2.handle(statusEventFromNode2_2);

  Assert.assertEquals(2, completedContainers.size());
  Assert.assertEquals(completedContainerIdFromNode2_1,completedContainers.get(0)
      .getContainerId()); 
  Assert.assertEquals(completedContainerIdFromNode2_2,completedContainers.get(1)
      .getContainerId());   
}
 

@Test
@SuppressWarnings("unchecked")
public void testStateStoreSuccessfulLocalization() throws Exception {
  final String user = "someuser";
  final ApplicationId appId = ApplicationId.newInstance(1, 1);
  // This is a random path. NO File creation will take place at this place.
  final Path localDir = new Path("/tmp");
  Configuration conf = new YarnConfiguration();
  DrainDispatcher dispatcher = null;
  dispatcher = createDispatcher(conf);
  EventHandler<LocalizerEvent> localizerEventHandler =
      mock(EventHandler.class);
  EventHandler<LocalizerEvent> containerEventHandler =
      mock(EventHandler.class);
  dispatcher.register(LocalizerEventType.class, localizerEventHandler);
  dispatcher.register(ContainerEventType.class, containerEventHandler);
  DeletionService mockDelService = mock(DeletionService.class);
  NMStateStoreService stateStore = mock(NMStateStoreService.class);

  try {
    LocalResourcesTracker tracker = new LocalResourcesTrackerImpl(user,
        appId, dispatcher, false, conf, stateStore);
    // Container 1 needs lr1 resource
    ContainerId cId1 = BuilderUtils.newContainerId(1, 1, 1, 1);
    LocalResourceRequest lr1 = createLocalResourceRequest(user, 1, 1,
        LocalResourceVisibility.APPLICATION);
    LocalizerContext lc1 = new LocalizerContext(user, cId1, null);

    // Container 1 requests lr1 to be localized
    ResourceEvent reqEvent1 = new ResourceRequestEvent(lr1,
        LocalResourceVisibility.APPLICATION, lc1);
    tracker.handle(reqEvent1);
    dispatcher.await();

    // Simulate the process of localization of lr1
    Path hierarchicalPath1 = tracker.getPathForLocalization(lr1, localDir);

    ArgumentCaptor<LocalResourceProto> localResourceCaptor =
        ArgumentCaptor.forClass(LocalResourceProto.class);
    ArgumentCaptor<Path> pathCaptor = ArgumentCaptor.forClass(Path.class);
    verify(stateStore).startResourceLocalization(eq(user), eq(appId),
        localResourceCaptor.capture(), pathCaptor.capture());
    LocalResourceProto lrProto = localResourceCaptor.getValue();
    Path localizedPath1 = pathCaptor.getValue();
    Assert.assertEquals(lr1,
        new LocalResourceRequest(new LocalResourcePBImpl(lrProto)));
    Assert.assertEquals(hierarchicalPath1, localizedPath1.getParent());

    // Simulate lr1 getting localized
    ResourceLocalizedEvent rle1 =
        new ResourceLocalizedEvent(lr1, pathCaptor.getValue(), 120);
    tracker.handle(rle1);
    dispatcher.await();

    ArgumentCaptor<LocalizedResourceProto> localizedProtoCaptor =
        ArgumentCaptor.forClass(LocalizedResourceProto.class);
    verify(stateStore).finishResourceLocalization(eq(user), eq(appId),
        localizedProtoCaptor.capture());
    LocalizedResourceProto localizedProto = localizedProtoCaptor.getValue();
    Assert.assertEquals(lr1, new LocalResourceRequest(
        new LocalResourcePBImpl(localizedProto.getResource())));
    Assert.assertEquals(localizedPath1.toString(),
        localizedProto.getLocalPath());
    LocalizedResource localizedRsrc1 = tracker.getLocalizedResource(lr1);
    Assert.assertNotNull(localizedRsrc1);

    // simulate release and retention processing
    tracker.handle(new ResourceReleaseEvent(lr1, cId1));
    dispatcher.await();
    boolean removeResult = tracker.remove(localizedRsrc1, mockDelService);

    Assert.assertTrue(removeResult);
    verify(stateStore).removeLocalizedResource(eq(user), eq(appId),
        eq(localizedPath1));
  } finally {
    if (dispatcher != null) {
      dispatcher.stop();
    }
  }
}
 

private int assignContainer(FiCaSchedulerNode node, FiCaSchedulerApp application, 
    Priority priority, int assignableContainers, 
    ResourceRequest request, NodeType type) {
  LOG.debug("assignContainers:" +
      " node=" + node.getRMNode().getNodeAddress() + 
      " application=" + application.getApplicationId().getId() + 
      " priority=" + priority.getPriority() + 
      " assignableContainers=" + assignableContainers +
      " request=" + request + " type=" + type);
  Resource capability = request.getCapability();

  int availableContainers = 
    node.getAvailableResource().getMemory() / capability.getMemory(); // TODO: A buggy
                                                                      // application
                                                                      // with this
                                                                      // zero would
                                                                      // crash the
                                                                      // scheduler.
  int assignedContainers = 
    Math.min(assignableContainers, availableContainers);

  if (assignedContainers > 0) {
    for (int i=0; i < assignedContainers; ++i) {

      NodeId nodeId = node.getRMNode().getNodeID();
      ContainerId containerId = BuilderUtils.newContainerId(application
          .getApplicationAttemptId(), application.getNewContainerId());

      // Create the container
      Container container =
          BuilderUtils.newContainer(containerId, nodeId, node.getRMNode()
            .getHttpAddress(), capability, priority, null);
      
      // Allocate!
      
      // Inform the application
      RMContainer rmContainer =
          application.allocate(type, node, priority, request, container);
      
      // Inform the node
      node.allocateContainer(rmContainer);

      // Update usage for this container
      increaseUsedResources(rmContainer);
    }

  }
  
  return assignedContainers;
}
 

@Test
public void testLogAggregationForRealContainerLaunch() throws IOException,
    InterruptedException, YarnException {

  this.containerManager.start();


  File scriptFile = new File(tmpDir, "scriptFile.sh");
  PrintWriter fileWriter = new PrintWriter(scriptFile);
  fileWriter.write("\necho Hello World! Stdout! > "
      + new File(localLogDir, "stdout"));
  fileWriter.write("\necho Hello World! Stderr! > "
      + new File(localLogDir, "stderr"));
  fileWriter.write("\necho Hello World! Syslog! > "
      + new File(localLogDir, "syslog"));
  fileWriter.close();

  ContainerLaunchContext containerLaunchContext =
      recordFactory.newRecordInstance(ContainerLaunchContext.class);
  // ////// Construct the Container-id
  ApplicationId appId = ApplicationId.newInstance(0, 0);
  ApplicationAttemptId appAttemptId =
      BuilderUtils.newApplicationAttemptId(appId, 1);
  ContainerId cId = BuilderUtils.newContainerId(appAttemptId, 0);

  URL resource_alpha =
      ConverterUtils.getYarnUrlFromPath(localFS
          .makeQualified(new Path(scriptFile.getAbsolutePath())));
  LocalResource rsrc_alpha =
      recordFactory.newRecordInstance(LocalResource.class);
  rsrc_alpha.setResource(resource_alpha);
  rsrc_alpha.setSize(-1);
  rsrc_alpha.setVisibility(LocalResourceVisibility.APPLICATION);
  rsrc_alpha.setType(LocalResourceType.FILE);
  rsrc_alpha.setTimestamp(scriptFile.lastModified());
  String destinationFile = "dest_file";
  Map<String, LocalResource> localResources = 
      new HashMap<String, LocalResource>();
  localResources.put(destinationFile, rsrc_alpha);
  containerLaunchContext.setLocalResources(localResources);
  List<String> commands = new ArrayList<String>();
  commands.add("/bin/bash");
  commands.add(scriptFile.getAbsolutePath());
  containerLaunchContext.setCommands(commands);

  StartContainerRequest scRequest =
      StartContainerRequest.newInstance(containerLaunchContext,
        TestContainerManager.createContainerToken(
          cId, DUMMY_RM_IDENTIFIER, context.getNodeId(), user,
          context.getContainerTokenSecretManager()));
  List<StartContainerRequest> list = new ArrayList<StartContainerRequest>();
  list.add(scRequest);
  StartContainersRequest allRequests =
      StartContainersRequest.newInstance(list);
  this.containerManager.startContainers(allRequests);
  
  BaseContainerManagerTest.waitForContainerState(this.containerManager,
      cId, ContainerState.COMPLETE);

  this.containerManager.handle(new CMgrCompletedAppsEvent(Arrays
      .asList(appId), CMgrCompletedAppsEvent.Reason.ON_SHUTDOWN));
  this.containerManager.stop();
}
 

private int assignContainer(FiCaSchedulerNode node, FiCaSchedulerApp application, 
    Priority priority, int assignableContainers, 
    ResourceRequest request, NodeType type) {
  LOG.debug("assignContainers:" +
      " node=" + node.getRMNode().getNodeAddress() + 
      " application=" + application.getApplicationId().getId() + 
      " priority=" + priority.getPriority() + 
      " assignableContainers=" + assignableContainers +
      " request=" + request + " type=" + type);
  Resource capability = request.getCapability();

  int availableContainers = 
    node.getAvailableResource().getMemory() / capability.getMemory(); // TODO: A buggy
                                                                      // application
                                                                      // with this
                                                                      // zero would
                                                                      // crash the
                                                                      // scheduler.
  int assignedContainers = 
    Math.min(assignableContainers, availableContainers);

  if (assignedContainers > 0) {
    for (int i=0; i < assignedContainers; ++i) {

      NodeId nodeId = node.getRMNode().getNodeID();
      ContainerId containerId = BuilderUtils.newContainerId(application
          .getApplicationAttemptId(), application.getNewContainerId());

      // Create the container
      Container container =
          BuilderUtils.newContainer(containerId, nodeId, node.getRMNode()
            .getHttpAddress(), capability, priority, null);
      
      // Allocate!
      
      // Inform the application
      RMContainer rmContainer =
          application.allocate(type, node, priority, request, container);
      
      // Inform the node
      node.allocateContainer(rmContainer);

      // Update usage for this container
      increaseUsedResources(rmContainer);
    }

  }
  
  return assignedContainers;
}
 

@Test
public void testExpireWhileRunning() {

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

  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.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 EXPIRE and associated actions.
  reset(appAttemptEventHandler);
  ContainerStatus containerStatus = SchedulerUtils
      .createAbnormalContainerStatus(containerId,
          SchedulerUtils.EXPIRED_CONTAINER);
  rmContainer.handle(new RMContainerFinishedEvent(containerId,
      containerStatus, RMContainerEventType.EXPIRE));
  drainDispatcher.await();
  assertEquals(RMContainerState.RUNNING, rmContainer.getState());
  verify(writer, never()).containerFinished(any(RMContainer.class));
  verify(publisher, never()).containerFinished(any(RMContainer.class),
      anyLong());
}
 

@Test
public void testFinishedContainer() {
  Container amContainer = allocateApplicationAttempt();
  launchApplicationAttempt(amContainer);
  runApplicationAttempt(amContainer, "host", 8042, "oldtrackingurl", false);

  // Complete one container
  ContainerId containerId1 = BuilderUtils.newContainerId(applicationAttempt
      .getAppAttemptId(), 2);
  Container container1 = mock(Container.class);
  ContainerStatus containerStatus1 = mock(ContainerStatus.class);
  when(container1.getId()).thenReturn(
      containerId1);
  when(containerStatus1.getContainerId()).thenReturn(containerId1);
  when(container1.getNodeId()).thenReturn(NodeId.newInstance("host", 1234));

  application.handle(new RMAppRunningOnNodeEvent(application
      .getApplicationId(),
      container1.getNodeId()));
  applicationAttempt.handle(new RMAppAttemptContainerFinishedEvent(
      applicationAttempt.getAppAttemptId(), containerStatus1,
      container1.getNodeId()));

  ArgumentCaptor<RMNodeFinishedContainersPulledByAMEvent> captor =
      ArgumentCaptor.forClass(RMNodeFinishedContainersPulledByAMEvent.class);

  // Verify justFinishedContainers
  Assert.assertEquals(1, applicationAttempt.getJustFinishedContainers()
      .size());
  Assert.assertEquals(container1.getId(), applicationAttempt
      .getJustFinishedContainers().get(0).getContainerId());
  Assert.assertEquals(0, getFinishedContainersSentToAM(applicationAttempt)
      .size());

  // Verify finishedContainersSentToAM gets container after pull
  List<ContainerStatus> containerStatuses = applicationAttempt
      .pullJustFinishedContainers();
  Assert.assertEquals(1, containerStatuses.size());
  Mockito.verify(rmnodeEventHandler, never()).handle(Mockito
      .any(RMNodeEvent.class));
  Assert.assertTrue(applicationAttempt.getJustFinishedContainers().isEmpty());
  Assert.assertEquals(1, getFinishedContainersSentToAM(applicationAttempt)
      .size());

  // Verify container is acked to NM via the RMNodeEvent after second pull
  containerStatuses = applicationAttempt.pullJustFinishedContainers();
  Assert.assertEquals(0, containerStatuses.size());
  Mockito.verify(rmnodeEventHandler).handle(captor.capture());
  Assert.assertEquals(container1.getId(), captor.getValue().getContainers()
      .get(0));
  Assert.assertTrue(applicationAttempt.getJustFinishedContainers().isEmpty());
  Assert.assertEquals(0, getFinishedContainersSentToAM(applicationAttempt)
      .size());
}
 

@Test
@SuppressWarnings("unchecked")
public void testRecoveredResource() throws Exception {
  final String user = "someuser";
  final ApplicationId appId = ApplicationId.newInstance(1, 1);
  // This is a random path. NO File creation will take place at this place.
  final Path localDir = new Path("/tmp/localdir");
  Configuration conf = new YarnConfiguration();
  DrainDispatcher dispatcher = null;
  dispatcher = createDispatcher(conf);
  EventHandler<LocalizerEvent> localizerEventHandler =
      mock(EventHandler.class);
  EventHandler<LocalizerEvent> containerEventHandler =
      mock(EventHandler.class);
  dispatcher.register(LocalizerEventType.class, localizerEventHandler);
  dispatcher.register(ContainerEventType.class, containerEventHandler);
  NMStateStoreService stateStore = mock(NMStateStoreService.class);

  try {
    LocalResourcesTracker tracker = new LocalResourcesTrackerImpl(user,
        appId, dispatcher, false, conf, stateStore);
    // Container 1 needs lr1 resource
    ContainerId cId1 = BuilderUtils.newContainerId(1, 1, 1, 1);
    LocalResourceRequest lr1 = createLocalResourceRequest(user, 1, 1,
        LocalResourceVisibility.APPLICATION);
    Assert.assertNull(tracker.getLocalizedResource(lr1));
    final long localizedId1 = 52;
    Path hierarchicalPath1 = new Path(localDir,
        Long.toString(localizedId1));
    Path localizedPath1 = new Path(hierarchicalPath1, "resource.jar");
    tracker.handle(new ResourceRecoveredEvent(lr1, localizedPath1, 120));
    dispatcher.await();
    Assert.assertNotNull(tracker.getLocalizedResource(lr1));

    // verify new paths reflect recovery of previous resources
    LocalResourceRequest lr2 = createLocalResourceRequest(user, 2, 2,
        LocalResourceVisibility.APPLICATION);
    LocalizerContext lc2 = new LocalizerContext(user, cId1, null);
    ResourceEvent reqEvent2 = new ResourceRequestEvent(lr2,
        LocalResourceVisibility.APPLICATION, lc2);
    tracker.handle(reqEvent2);
    dispatcher.await();
    Path hierarchicalPath2 = tracker.getPathForLocalization(lr2, localDir,
        null);
    long localizedId2 = Long.parseLong(hierarchicalPath2.getName());
    Assert.assertEquals(localizedId1 + 1, localizedId2);
  } finally {
    if (dispatcher != null) {
      dispatcher.stop();
    }
  }
}
 

@Test
@SuppressWarnings("unchecked")
public void testStateStoreSuccessfulLocalization() throws Exception {
  final String user = "someuser";
  final ApplicationId appId = ApplicationId.newInstance(1, 1);
  // This is a random path. NO File creation will take place at this place.
  final Path localDir = new Path("/tmp");
  Configuration conf = new YarnConfiguration();
  DrainDispatcher dispatcher = null;
  dispatcher = createDispatcher(conf);
  EventHandler<LocalizerEvent> localizerEventHandler =
      mock(EventHandler.class);
  EventHandler<LocalizerEvent> containerEventHandler =
      mock(EventHandler.class);
  dispatcher.register(LocalizerEventType.class, localizerEventHandler);
  dispatcher.register(ContainerEventType.class, containerEventHandler);
  DeletionService mockDelService = mock(DeletionService.class);
  NMStateStoreService stateStore = mock(NMStateStoreService.class);

  try {
    LocalResourcesTracker tracker = new LocalResourcesTrackerImpl(user,
        appId, dispatcher, false, conf, stateStore);
    // Container 1 needs lr1 resource
    ContainerId cId1 = BuilderUtils.newContainerId(1, 1, 1, 1);
    LocalResourceRequest lr1 = createLocalResourceRequest(user, 1, 1,
        LocalResourceVisibility.APPLICATION);
    LocalizerContext lc1 = new LocalizerContext(user, cId1, null);

    // Container 1 requests lr1 to be localized
    ResourceEvent reqEvent1 = new ResourceRequestEvent(lr1,
        LocalResourceVisibility.APPLICATION, lc1);
    tracker.handle(reqEvent1);
    dispatcher.await();

    // Simulate the process of localization of lr1
    Path hierarchicalPath1 = tracker.getPathForLocalization(lr1, localDir,
        null);

    ArgumentCaptor<LocalResourceProto> localResourceCaptor =
        ArgumentCaptor.forClass(LocalResourceProto.class);
    ArgumentCaptor<Path> pathCaptor = ArgumentCaptor.forClass(Path.class);
    verify(stateStore).startResourceLocalization(eq(user), eq(appId),
        localResourceCaptor.capture(), pathCaptor.capture());
    LocalResourceProto lrProto = localResourceCaptor.getValue();
    Path localizedPath1 = pathCaptor.getValue();
    Assert.assertEquals(lr1,
        new LocalResourceRequest(new LocalResourcePBImpl(lrProto)));
    Assert.assertEquals(hierarchicalPath1, localizedPath1.getParent());

    // Simulate lr1 getting localized
    ResourceLocalizedEvent rle1 =
        new ResourceLocalizedEvent(lr1, pathCaptor.getValue(), 120);
    tracker.handle(rle1);
    dispatcher.await();

    ArgumentCaptor<LocalizedResourceProto> localizedProtoCaptor =
        ArgumentCaptor.forClass(LocalizedResourceProto.class);
    verify(stateStore).finishResourceLocalization(eq(user), eq(appId),
        localizedProtoCaptor.capture());
    LocalizedResourceProto localizedProto = localizedProtoCaptor.getValue();
    Assert.assertEquals(lr1, new LocalResourceRequest(
        new LocalResourcePBImpl(localizedProto.getResource())));
    Assert.assertEquals(localizedPath1.toString(),
        localizedProto.getLocalPath());
    LocalizedResource localizedRsrc1 = tracker.getLocalizedResource(lr1);
    Assert.assertNotNull(localizedRsrc1);

    // simulate release and retention processing
    tracker.handle(new ResourceReleaseEvent(lr1, cId1));
    dispatcher.await();
    boolean removeResult = tracker.remove(localizedRsrc1, mockDelService);

    Assert.assertTrue(removeResult);
    verify(stateStore).removeLocalizedResource(eq(user), eq(appId),
        eq(localizedPath1));
  } finally {
    if (dispatcher != null) {
      dispatcher.stop();
    }
  }
}
 

@Test(timeout=10000)
@SuppressWarnings("unchecked")
public void testConsistency() {
  String user = "testuser";
  DrainDispatcher dispatcher = null;
  try {
    Configuration conf = new Configuration();
    dispatcher = createDispatcher(conf);
    EventHandler<LocalizerEvent> localizerEventHandler = mock(EventHandler.class);
    EventHandler<LocalizerEvent> containerEventHandler = mock(EventHandler.class);
    dispatcher.register(LocalizerEventType.class, localizerEventHandler);
    dispatcher.register(ContainerEventType.class, containerEventHandler);

    ContainerId cId1 = BuilderUtils.newContainerId(1, 1, 1, 1);
    LocalizerContext lc1 = new LocalizerContext(user, cId1, null);
    LocalResourceRequest req1 = createLocalResourceRequest(user, 1, 1,
        LocalResourceVisibility.PUBLIC);
    LocalizedResource lr1 = createLocalizedResource(req1, dispatcher);
    ConcurrentMap<LocalResourceRequest, LocalizedResource> localrsrc = new ConcurrentHashMap<LocalResourceRequest, LocalizedResource>();
    localrsrc.put(req1, lr1);
    LocalResourcesTracker tracker = new LocalResourcesTrackerImpl(user,
        null, dispatcher, localrsrc, false, conf,
        new NMNullStateStoreService(), null);

    ResourceEvent req11Event = new ResourceRequestEvent(req1,
        LocalResourceVisibility.PUBLIC, lc1);

    ResourceEvent rel11Event = new ResourceReleaseEvent(req1, cId1);

    // Localize R1 for C1
    tracker.handle(req11Event);

    dispatcher.await();

    // Verify refCount for R1 is 1
    Assert.assertEquals(1, lr1.getRefCount());

    dispatcher.await();
    verifyTrackedResourceCount(tracker, 1);

    // Localize resource1
    ResourceLocalizedEvent rle = new ResourceLocalizedEvent(req1, new Path(
        "file:///tmp/r1"), 1);
    lr1.handle(rle);
    Assert.assertTrue(lr1.getState().equals(ResourceState.LOCALIZED));
    Assert.assertTrue(createdummylocalizefile(new Path("file:///tmp/r1")));
    LocalizedResource rsrcbefore = tracker.iterator().next();
    File resFile = new File(lr1.getLocalPath().toUri().getRawPath()
        .toString());
    Assert.assertTrue(resFile.exists());
    Assert.assertTrue(resFile.delete());

    // Localize R1 for C1
    tracker.handle(req11Event);

    dispatcher.await();
    lr1.handle(rle);
    Assert.assertTrue(lr1.getState().equals(ResourceState.LOCALIZED));
    LocalizedResource rsrcafter = tracker.iterator().next();
    if (rsrcbefore == rsrcafter) {
      Assert.fail("Localized resource should not be equal");
    }
    // Release resource1
    tracker.handle(rel11Event);
  } finally {
    if (dispatcher != null) {
      dispatcher.stop();
    }
  }
}
 

@Test
  public void testNMWebApp() throws IOException, YarnException {
    Context nmContext = new NodeManager.NMContext(null, null, null, null,
        null);
    ResourceView resourceView = new ResourceView() {
      @Override
      public long getVmemAllocatedForContainers() {
        return 0;
      }
      @Override
      public long getPmemAllocatedForContainers() {
        return 0;
      }
      @Override
      public long getVCoresAllocatedForContainers() {
        return 0;
      }
      @Override
      public long getGCoresAllocatedForContainers() {
        return 0;
      }
      @Override
      public boolean isVmemCheckEnabled() {
        return true;
      }
      @Override
      public boolean isPmemCheckEnabled() {
        return true;
      }
    };
    Configuration conf = new Configuration();
    conf.set(YarnConfiguration.NM_LOCAL_DIRS, testRootDir.getAbsolutePath());
    conf.set(YarnConfiguration.NM_LOG_DIRS, testLogDir.getAbsolutePath());
    NodeHealthCheckerService healthChecker = createNodeHealthCheckerService(conf);
    healthChecker.init(conf);
    LocalDirsHandlerService dirsHandler = healthChecker.getDiskHandler();

    WebServer server = new WebServer(nmContext, resourceView,
        new ApplicationACLsManager(conf), dirsHandler);
    server.init(conf);
    server.start();

    // Add an application and the corresponding containers
    RecordFactory recordFactory =
        RecordFactoryProvider.getRecordFactory(conf);
    Dispatcher dispatcher = new AsyncDispatcher();
    String user = "nobody";
    long clusterTimeStamp = 1234;
    ApplicationId appId =
        BuilderUtils.newApplicationId(recordFactory, clusterTimeStamp, 1);
    Application app = mock(Application.class);
    when(app.getUser()).thenReturn(user);
    when(app.getAppId()).thenReturn(appId);
    nmContext.getApplications().put(appId, app);
    ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
        appId, 1);
    ContainerId container1 =
        BuilderUtils.newContainerId(recordFactory, appId, appAttemptId, 0);
    ContainerId container2 =
        BuilderUtils.newContainerId(recordFactory, appId, appAttemptId, 1);
    NodeManagerMetrics metrics = mock(NodeManagerMetrics.class);
    NMStateStoreService stateStore = new NMNullStateStoreService();
    for (ContainerId containerId : new ContainerId[] { container1,
        container2}) {
      // TODO: Use builder utils
      ContainerLaunchContext launchContext =
          recordFactory.newRecordInstance(ContainerLaunchContext.class);
      long currentTime = System.currentTimeMillis();
      Token containerToken =
          BuilderUtils.newContainerToken(containerId, "127.0.0.1", 1234, user,
            BuilderUtils.newResource(1024, 1), currentTime + 10000L, 123,
            "password".getBytes(), currentTime);
      Container container =
          new ContainerImpl(conf, dispatcher, stateStore, launchContext,
            null, metrics,
            BuilderUtils.newContainerTokenIdentifier(containerToken)) {

            @Override
            public ContainerState getContainerState() {
              return ContainerState.RUNNING;
            };
          };
      nmContext.getContainers().put(containerId, container);
      //TODO: Gross hack. Fix in code.
      ApplicationId applicationId = 
          containerId.getApplicationAttemptId().getApplicationId();
      nmContext.getApplications().get(applicationId).getContainers()
          .put(containerId, container);
      writeContainerLogs(nmContext, containerId, dirsHandler);

    }
    // TODO: Pull logs and test contents.
//    Thread.sleep(1000000);
  }
 

private ContainerId newContainerId(int appId, int appAttemptId, int cId) {
  return BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(System.currentTimeMillis(), appId),
          appAttemptId), cId);
}
 

@Test(timeout=30000)
public void testContainerLogDirs() throws IOException, YarnException {
  File absLogDir = new File("target",
    TestNMWebServer.class.getSimpleName() + "LogDir").getAbsoluteFile();
  String logdirwithFile = absLogDir.toURI().toString();
  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.NM_LOG_DIRS, logdirwithFile);
  NodeHealthCheckerService healthChecker = createNodeHealthCheckerService(conf);
  healthChecker.init(conf);
  LocalDirsHandlerService dirsHandler = healthChecker.getDiskHandler();
  NMContext nmContext = new NodeManager.NMContext(null, null, dirsHandler,
      new ApplicationACLsManager(conf), new NMNullStateStoreService());
  // Add an application and the corresponding containers
  RecordFactory recordFactory = RecordFactoryProvider.getRecordFactory(conf);
  String user = "nobody";
  long clusterTimeStamp = 1234;
  ApplicationId appId = BuilderUtils.newApplicationId(recordFactory,
      clusterTimeStamp, 1);
  Application app = mock(Application.class);
  when(app.getUser()).thenReturn(user);
  when(app.getAppId()).thenReturn(appId);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      appId, 1);
  ContainerId container1 = BuilderUtils.newContainerId(recordFactory, appId,
      appAttemptId, 0);
  nmContext.getApplications().put(appId, app);

  MockContainer container =
      new MockContainer(appAttemptId, new AsyncDispatcher(), conf, user,
          appId, 1);
  container.setState(ContainerState.RUNNING);
  nmContext.getContainers().put(container1, container);   
  List<File> files = null;
  files = ContainerLogsUtils.getContainerLogDirs(container1, user, nmContext);
  Assert.assertTrue(!(files.get(0).toString().contains("file:")));
  
  // After container is completed, it is removed from nmContext
  nmContext.getContainers().remove(container1);
  Assert.assertNull(nmContext.getContainers().get(container1));
  files = ContainerLogsUtils.getContainerLogDirs(container1, user, nmContext);
  Assert.assertTrue(!(files.get(0).toString().contains("file:")));

  // Create a new context to check if correct container log dirs are fetched
  // on full disk.
  LocalDirsHandlerService dirsHandlerForFullDisk = spy(dirsHandler);
  // good log dirs are empty and nm log dir is in the full log dir list.
  when(dirsHandlerForFullDisk.getLogDirs()).
      thenReturn(new ArrayList<String>());
  when(dirsHandlerForFullDisk.getLogDirsForRead()).
      thenReturn(Arrays.asList(new String[] {absLogDir.getAbsolutePath()}));
  nmContext = new NodeManager.NMContext(null, null, dirsHandlerForFullDisk,
      new ApplicationACLsManager(conf), new NMNullStateStoreService());
  nmContext.getApplications().put(appId, app);
  container.setState(ContainerState.RUNNING);
  nmContext.getContainers().put(container1, container);
  List<File> dirs =
      ContainerLogsUtils.getContainerLogDirs(container1, user, nmContext);
  File containerLogDir = new File(absLogDir, appId + "/" + container1);
  Assert.assertTrue(dirs.contains(containerLogDir));
}