Java Code Examples for org.apache.hadoop.yarn.event.DrainDispatcher#await()

protected void waitForContainerCleanup(DrainDispatcher dispatcher, MockNM nm,
    NodeHeartbeatResponse resp) throws Exception {
  int waitCount = 0, cleanedConts = 0;
  List<ContainerId> contsToClean;
  do {
    dispatcher.await();
    contsToClean = resp.getContainersToCleanup();
    cleanedConts += contsToClean.size();
    if (cleanedConts >= 1) {
      break;
    }
    Thread.sleep(100);
    resp = nm.nodeHeartbeat(true);
  } while(waitCount++ < 200);

  if (contsToClean.isEmpty()) {
    LOG.error("Failed to get any containers to cleanup");
  } else {
    LOG.info("Got cleanup for " + contsToClean.get(0));
  }
  Assert.assertEquals(1, cleanedConts);
}
 

protected void waitForContainerCleanup(DrainDispatcher dispatcher, MockNM nm,
    NodeHeartbeatResponse resp) throws Exception {
  int waitCount = 0, cleanedConts = 0;
  List<ContainerId> contsToClean;
  do {
    dispatcher.await();
    contsToClean = resp.getContainersToCleanup();
    cleanedConts += contsToClean.size();
    if (cleanedConts >= 1) {
      break;
    }
    Thread.sleep(100);
    resp = nm.nodeHeartbeat(true);
  } while(waitCount++ < 200);

  if (contsToClean.isEmpty()) {
    LOG.error("Failed to get any containers to cleanup");
  } else {
    LOG.info("Got cleanup for " + contsToClean.get(0));
  }
  Assert.assertEquals(1, cleanedConts);
}
 

private void finishTask(DrainDispatcher rmDispatcher, MockNM node,
    MRApp mrApp, Task task) throws Exception {
  TaskAttempt attempt = task.getAttempts().values().iterator().next();
  List<ContainerStatus> contStatus = new ArrayList<ContainerStatus>(1);
  contStatus.add(ContainerStatus.newInstance(attempt.getAssignedContainerID(),
      ContainerState.COMPLETE, "", 0));
  Map<ApplicationId,List<ContainerStatus>> statusUpdate =
      new HashMap<ApplicationId,List<ContainerStatus>>(1);
  statusUpdate.put(mrApp.getAppID(), contStatus);
  node.nodeHeartbeat(statusUpdate, true);
  rmDispatcher.await();
  mrApp.getContext().getEventHandler().handle(
        new TaskAttemptEvent(attempt.getID(), TaskAttemptEventType.TA_DONE));
  mrApp.waitForState(task, TaskState.SUCCEEDED);
}
 

private
    List<TaskAttemptContainerAssignedEvent> getContainerOnHost(JobId jobId,
        int taskAttemptId, int memory, String[] hosts, MockNM mockNM,
        DrainDispatcher dispatcher, MyContainerAllocator allocator,
        int expectedAdditions1, int expectedRemovals1,
        int expectedAdditions2, int expectedRemovals2, MyResourceManager rm)
        throws Exception {
  ContainerRequestEvent reqEvent =
      createReq(jobId, taskAttemptId, memory, hosts);
  allocator.sendRequest(reqEvent);

  // Send the request to the RM
  List<TaskAttemptContainerAssignedEvent> assigned = allocator.schedule();
  dispatcher.await();
  assertBlacklistAdditionsAndRemovals(
      expectedAdditions1, expectedRemovals1, rm);
  Assert.assertEquals("No of assignments must be 0", 0, assigned.size());

  // Heartbeat from the required nodeManager
  mockNM.nodeHeartbeat(true);
  dispatcher.await();

  assigned = allocator.schedule();
  dispatcher.await();
  assertBlacklistAdditionsAndRemovals(
      expectedAdditions2, expectedRemovals2, rm);
  return assigned;
}
 

@Test
public void testUnregistrationOnlyIfRegistered() throws Exception {
  Configuration conf = new Configuration();
  final MyResourceManager rm = new MyResourceManager(conf);
  rm.start();
  DrainDispatcher rmDispatcher =
      (DrainDispatcher) rm.getRMContext().getDispatcher();

  // Submit the application
  RMApp rmApp = rm.submitApp(1024);
  rmDispatcher.await();

  MockNM amNodeManager = rm.registerNode("127.0.0.1:1234", 11264);
  amNodeManager.nodeHeartbeat(true);
  rmDispatcher.await();

  final ApplicationAttemptId appAttemptId =
      rmApp.getCurrentAppAttempt().getAppAttemptId();
  rm.sendAMLaunched(appAttemptId);
  rmDispatcher.await();

  MRApp mrApp =
      new MRApp(appAttemptId, ContainerId.newContainerId(appAttemptId, 0), 10,
          0, false, this.getClass().getName(), true, 1) {
        @Override
        protected Dispatcher createDispatcher() {
          return new DrainDispatcher();
        }

        protected ContainerAllocator createContainerAllocator(
            ClientService clientService, AppContext context) {
          return new MyContainerAllocator(rm, appAttemptId, context);
        };
      };

  mrApp.submit(conf);
  DrainDispatcher amDispatcher = (DrainDispatcher) mrApp.getDispatcher();
  MyContainerAllocator allocator =
      (MyContainerAllocator) mrApp.getContainerAllocator();
  amDispatcher.await();
  Assert.assertTrue(allocator.isApplicationMasterRegistered());
  mrApp.stop();
  Assert.assertTrue(allocator.isUnregistered());
}
 

@Test(timeout = 30000)
public void testReducerRampdownDiagnostics() throws Exception {
  LOG.info("Running tesReducerRampdownDiagnostics");

  final Configuration conf = new Configuration();
  conf.setFloat(MRJobConfig.COMPLETED_MAPS_FOR_REDUCE_SLOWSTART, 0.0f);
  final MyResourceManager rm = new MyResourceManager(conf);
  rm.start();
  final DrainDispatcher dispatcher = (DrainDispatcher) rm.getRMContext()
      .getDispatcher();
  final RMApp app = rm.submitApp(1024);
  dispatcher.await();

  final String host = "host1";
  final MockNM nm = rm.registerNode(String.format("%s:1234", host), 2048);
  nm.nodeHeartbeat(true);
  dispatcher.await();
  final ApplicationAttemptId appAttemptId = app.getCurrentAppAttempt()
      .getAppAttemptId();
  rm.sendAMLaunched(appAttemptId);
  dispatcher.await();
  final JobId jobId = MRBuilderUtils
      .newJobId(appAttemptId.getApplicationId(), 0);
  final Job mockJob = mock(Job.class);
  when(mockJob.getReport()).thenReturn(
      MRBuilderUtils.newJobReport(jobId, "job", "user", JobState.RUNNING, 0,
          0, 0, 0, 0, 0, 0, "jobfile", null, false, ""));
  final MyContainerAllocator allocator = new MyContainerAllocator(rm, conf,
      appAttemptId, mockJob);
  // add resources to scheduler
  dispatcher.await();

  // create the container request
  final String[] locations = new String[] { host };
  allocator.sendRequest(createReq(jobId, 0, 1024, locations, false, true));
  for (int i = 0; i < 1;) {
    dispatcher.await();
    i += allocator.schedule().size();
    nm.nodeHeartbeat(true);
  }

  allocator.sendRequest(createReq(jobId, 0, 1024, locations, true, false));
  while (allocator.getTaskAttemptKillEvents().size() == 0) {
    dispatcher.await();
    allocator.schedule().size();
    nm.nodeHeartbeat(true);
  }
  final String killEventMessage = allocator.getTaskAttemptKillEvents().get(0)
      .getMessage();
  Assert.assertTrue("No reducer rampDown preemption message",
      killEventMessage.contains(RMContainerAllocator.RAMPDOWN_DIAGNOSTIC));
}
 

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

@Test 
public void testMapNodeLocality() throws Exception {
  // test checks that ordering of allocated containers list from the RM does 
  // not affect the map->container assignment done by the AM. If there is a 
  // node local container available for a map then it should be assigned to 
  // that container and not a rack-local container that happened to be seen 
  // earlier in the allocated containers list from the RM.
  // Regression test for MAPREDUCE-4893
  LOG.info("Running testMapNodeLocality");

  Configuration conf = new Configuration();
  MyResourceManager rm = new MyResourceManager(conf);
  rm.start();
  DrainDispatcher dispatcher = (DrainDispatcher) rm.getRMContext()
      .getDispatcher();

  // Submit the application
  RMApp app = rm.submitApp(1024);
  dispatcher.await();

  MockNM amNodeManager = rm.registerNode("amNM:1234", 2048);
  amNodeManager.nodeHeartbeat(true);
  dispatcher.await();

  ApplicationAttemptId appAttemptId = app.getCurrentAppAttempt()
      .getAppAttemptId();
  rm.sendAMLaunched(appAttemptId);
  dispatcher.await();

  JobId jobId = MRBuilderUtils.newJobId(appAttemptId.getApplicationId(), 0);
  Job mockJob = mock(Job.class);
  when(mockJob.getReport()).thenReturn(
      MRBuilderUtils.newJobReport(jobId, "job", "user", JobState.RUNNING, 0, 
          0, 0, 0, 0, 0, 0, "jobfile", null, false, ""));
  MyContainerAllocator allocator = new MyContainerAllocator(rm, conf,
      appAttemptId, mockJob);

  // add resources to scheduler
  MockNM nodeManager1 = rm.registerNode("h1:1234", 3072); // can assign 2 maps 
  rm.registerNode("h2:1234", 10240); // wont heartbeat on node local node
  MockNM nodeManager3 = rm.registerNode("h3:1234", 1536); // assign 1 map
  dispatcher.await();

  // create the container requests for maps
  ContainerRequestEvent event1 = createReq(jobId, 1, 1024,
      new String[] { "h1" });
  allocator.sendRequest(event1);
  ContainerRequestEvent event2 = createReq(jobId, 2, 1024,
      new String[] { "h1" });
  allocator.sendRequest(event2);
  ContainerRequestEvent event3 = createReq(jobId, 3, 1024,
      new String[] { "h2" });
  allocator.sendRequest(event3);

  // this tells the scheduler about the requests
  // as nodes are not added, no allocations
  List<TaskAttemptContainerAssignedEvent> assigned = allocator.schedule();
  dispatcher.await();
  Assert.assertEquals("No of assignments must be 0", 0, assigned.size());

  // update resources in scheduler
  // Node heartbeat from rack-local first. This makes node h3 the first in the
  // list of allocated containers but it should not be assigned to task1.
  nodeManager3.nodeHeartbeat(true);
  // Node heartbeat from node-local next. This allocates 2 node local 
  // containers for task1 and task2. These should be matched with those tasks.
  nodeManager1.nodeHeartbeat(true);
  dispatcher.await();

  assigned = allocator.schedule();
  dispatcher.await();
  checkAssignments(new ContainerRequestEvent[] { event1, event2, event3 },
      assigned, false);
  // remove the rack-local assignment that should have happened for task3
  for(TaskAttemptContainerAssignedEvent event : assigned) {
    if(event.getTaskAttemptID().equals(event3.getAttemptID())) {
      assigned.remove(event);
      Assert.assertTrue(
                  event.getContainer().getNodeId().getHost().equals("h3"));
      break;
    }
  }
  checkAssignments(new ContainerRequestEvent[] { event1, event2},
      assigned, true);
}
 

@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);
    long localizedId2 = Long.parseLong(hierarchicalPath2.getName());
    Assert.assertEquals(localizedId1 + 1, localizedId2);
  } finally {
    if (dispatcher != null) {
      dispatcher.stop();
    }
  }
}
 

/**
 * Decommissioning using a pre-configured exclude hosts file
 */
@Test
public void testDecommissionWithExcludeHosts() throws Exception {
  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.RM_NODES_EXCLUDE_FILE_PATH, hostFile
      .getAbsolutePath());

  writeToHostsFile("");
  final DrainDispatcher dispatcher = new DrainDispatcher();
  rm = new MockRM(conf) {
    @Override
    protected Dispatcher createDispatcher() {
      return dispatcher;
    }
  };
  rm.start();

  MockNM nm1 = rm.registerNode("host1:1234", 5120);
  MockNM nm2 = rm.registerNode("host2:5678", 10240);
  MockNM nm3 = rm.registerNode("localhost:4433", 1024);

  dispatcher.await();

  int metricCount = ClusterMetrics.getMetrics().getNumDecommisionedNMs();
  NodeHeartbeatResponse nodeHeartbeat = nm1.nodeHeartbeat(true);
  Assert.assertTrue(NodeAction.NORMAL.equals(nodeHeartbeat.getNodeAction()));
  nodeHeartbeat = nm2.nodeHeartbeat(true);
  Assert.assertTrue(NodeAction.NORMAL.equals(nodeHeartbeat.getNodeAction()));
  dispatcher.await();

  // To test that IPs also work
  String ip = NetUtils.normalizeHostName("localhost");
  writeToHostsFile("host2", ip);

  rm.getNodesListManager().refreshNodes(conf);

  checkDecommissionedNMCount(rm, metricCount + 2);

  nodeHeartbeat = nm1.nodeHeartbeat(true);
  Assert.assertTrue(NodeAction.NORMAL.equals(nodeHeartbeat.getNodeAction()));
  nodeHeartbeat = nm2.nodeHeartbeat(true);
  Assert.assertTrue("The decommisioned metrics are not updated",
      NodeAction.SHUTDOWN.equals(nodeHeartbeat.getNodeAction()));

  nodeHeartbeat = nm3.nodeHeartbeat(true);
  Assert.assertTrue("The decommisioned metrics are not updated",
      NodeAction.SHUTDOWN.equals(nodeHeartbeat.getNodeAction()));
  dispatcher.await();

  writeToHostsFile("");
  rm.getNodesListManager().refreshNodes(conf);

  nm3 = rm.registerNode("localhost:4433", 1024);
  dispatcher.await();
  nodeHeartbeat = nm3.nodeHeartbeat(true);
  dispatcher.await();
  Assert.assertTrue(NodeAction.NORMAL.equals(nodeHeartbeat.getNodeAction()));
  // decommissined node is 1 since 1 node is rejoined after updating exclude
  // file
  checkDecommissionedNMCount(rm, metricCount + 1);
}
 

@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 testResource() throws Exception {

  LOG.info("Running testResource");

  Configuration conf = new Configuration();
  MyResourceManager rm = new MyResourceManager(conf);
  rm.start();
  DrainDispatcher dispatcher = (DrainDispatcher) rm.getRMContext()
      .getDispatcher();

  // Submit the application
  RMApp app = rm.submitApp(1024);
  dispatcher.await();

  MockNM amNodeManager = rm.registerNode("amNM:1234", 2048);
  amNodeManager.nodeHeartbeat(true);
  dispatcher.await();

  ApplicationAttemptId appAttemptId = app.getCurrentAppAttempt()
      .getAppAttemptId();
  rm.sendAMLaunched(appAttemptId);
  dispatcher.await();

  JobId jobId = MRBuilderUtils.newJobId(appAttemptId.getApplicationId(), 0);
  Job mockJob = mock(Job.class);
  when(mockJob.getReport()).thenReturn(
      MRBuilderUtils.newJobReport(jobId, "job", "user", JobState.RUNNING, 0,
          0, 0, 0, 0, 0, 0, "jobfile", null, false, ""));
  MyContainerAllocator allocator = new MyContainerAllocator(rm, conf,
      appAttemptId, mockJob);

  // add resources to scheduler
  MockNM nodeManager1 = rm.registerNode("h1:1234", 10240);
  MockNM nodeManager2 = rm.registerNode("h2:1234", 10240);
  MockNM nodeManager3 = rm.registerNode("h3:1234", 10240);
  dispatcher.await();

  // create the container request
  ContainerRequestEvent event1 = createReq(jobId, 1, 1024,
      new String[] { "h1" });
  allocator.sendRequest(event1);

  // send 1 more request with different resource req
  ContainerRequestEvent event2 = createReq(jobId, 2, 2048,
      new String[] { "h2" });
  allocator.sendRequest(event2);

  // this tells the scheduler about the requests
  // as nodes are not added, no allocations
  List<TaskAttemptContainerAssignedEvent> assigned = allocator.schedule();
  dispatcher.await();
  Assert.assertEquals("No of assignments must be 0", 0, assigned.size());

  // update resources in scheduler
  nodeManager1.nodeHeartbeat(true); // Node heartbeat
  nodeManager2.nodeHeartbeat(true); // Node heartbeat
  nodeManager3.nodeHeartbeat(true); // Node heartbeat
  dispatcher.await();

  assigned = allocator.schedule();
  dispatcher.await();
  checkAssignments(new ContainerRequestEvent[] { event1, event2 },
      assigned, false);
}
 

private MockNM registerNodeManager(int i, MyResourceManager rm,
    DrainDispatcher dispatcher) throws Exception {
  MockNM nm = rm.registerNode("h" + (i + 1) + ":1234", 10240);
  dispatcher.await();
  return nm;
}
 

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

@SuppressWarnings("unchecked")
@Test
public void testResourcePresentInGoodDir() throws IOException {
  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);
    LocalResourceRequest req2 =
        createLocalResourceRequest(user, 2, 1, LocalResourceVisibility.PUBLIC);
    LocalizedResource lr1 = createLocalizedResource(req1, dispatcher);
    LocalizedResource lr2 = createLocalizedResource(req2, dispatcher);
    ConcurrentMap<LocalResourceRequest, LocalizedResource> localrsrc =
        new ConcurrentHashMap<LocalResourceRequest, LocalizedResource>();
    localrsrc.put(req1, lr1);
    localrsrc.put(req2, lr2);
    LocalDirsHandlerService dirsHandler = mock(LocalDirsHandlerService.class);
    List<String> goodDirs = new ArrayList<String>();
    // /tmp/somedir2 is bad
    goodDirs.add("/tmp/somedir1/");
    goodDirs.add("/tmp/somedir2");
    Mockito.when(dirsHandler.getLocalDirs()).thenReturn(goodDirs);
    Mockito.when(dirsHandler.getLocalDirsForRead()).thenReturn(goodDirs);
    LocalResourcesTrackerImpl tracker =
        new LocalResourcesTrackerImpl(user, null, dispatcher, localrsrc,
            true , conf, new NMNullStateStoreService(), dirsHandler);
    ResourceEvent req11Event =
        new ResourceRequestEvent(req1, LocalResourceVisibility.PUBLIC, lc1);
    ResourceEvent req21Event =
        new ResourceRequestEvent(req2, LocalResourceVisibility.PUBLIC, lc1);
    // Localize R1 for C1
    tracker.handle(req11Event);
    // Localize R2 for C1
    tracker.handle(req21Event);
    dispatcher.await();
    // Localize resource1
    Path p1 = tracker.getPathForLocalization(req1,
        new Path("/tmp/somedir1"), null);
    Path p2 = tracker.getPathForLocalization(req2,
        new Path("/tmp/somedir2"), null);
    ResourceLocalizedEvent rle1 = new ResourceLocalizedEvent(req1, p1, 1);
    tracker.handle(rle1);
    ResourceLocalizedEvent rle2 = new ResourceLocalizedEvent(req2, p2, 1);
    tracker.handle(rle2);
    dispatcher.await();
    // Remove somedir2 from gooddirs
    Assert.assertTrue(tracker.checkLocalResource(lr2));
    goodDirs.remove(1);
    Assert.assertFalse(tracker.checkLocalResource(lr2));
  } finally {
    if (dispatcher != null) {
      dispatcher.stop();
    }
  }
}
 

@Test
@SuppressWarnings("unchecked")
public void testRecoveredResourceWithDirCacheMgr() 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 localDirRoot = 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 {
    LocalResourcesTrackerImpl tracker = new LocalResourcesTrackerImpl(user,
        appId, dispatcher, true, conf, stateStore);
    LocalResourceRequest lr1 = createLocalResourceRequest(user, 1, 1,
        LocalResourceVisibility.PUBLIC);
    Assert.assertNull(tracker.getLocalizedResource(lr1));
    final long localizedId1 = 52;
    Path hierarchicalPath1 = new Path(localDirRoot + "/4/2",
        Long.toString(localizedId1));
    Path localizedPath1 = new Path(hierarchicalPath1, "resource.jar");
    tracker.handle(new ResourceRecoveredEvent(lr1, localizedPath1, 120));
    dispatcher.await();
    Assert.assertNotNull(tracker.getLocalizedResource(lr1));
    LocalCacheDirectoryManager dirMgrRoot =
        tracker.getDirectoryManager(localDirRoot);
    Assert.assertEquals(0, dirMgrRoot.getDirectory("").getCount());
    Assert.assertEquals(1, dirMgrRoot.getDirectory("4/2").getCount());

    LocalResourceRequest lr2 = createLocalResourceRequest(user, 2, 2,
        LocalResourceVisibility.PUBLIC);
    Assert.assertNull(tracker.getLocalizedResource(lr2));
    final long localizedId2 = localizedId1 + 1;
    Path hierarchicalPath2 = new Path(localDirRoot + "/4/2",
        Long.toString(localizedId2));
    Path localizedPath2 = new Path(hierarchicalPath2, "resource.jar");
    tracker.handle(new ResourceRecoveredEvent(lr2, localizedPath2, 120));
    dispatcher.await();
    Assert.assertNotNull(tracker.getLocalizedResource(lr2));
    Assert.assertEquals(0, dirMgrRoot.getDirectory("").getCount());
    Assert.assertEquals(2, dirMgrRoot.getDirectory("4/2").getCount());

    LocalResourceRequest lr3 = createLocalResourceRequest(user, 3, 3,
        LocalResourceVisibility.PUBLIC);
    Assert.assertNull(tracker.getLocalizedResource(lr3));
    final long localizedId3 = 128;
    Path hierarchicalPath3 = new Path(localDirRoot + "/4/3",
        Long.toString(localizedId3));
    Path localizedPath3 = new Path(hierarchicalPath3, "resource.jar");
    tracker.handle(new ResourceRecoveredEvent(lr3, localizedPath3, 120));
    dispatcher.await();
    Assert.assertNotNull(tracker.getLocalizedResource(lr3));
    Assert.assertEquals(0, dirMgrRoot.getDirectory("").getCount());
    Assert.assertEquals(2, dirMgrRoot.getDirectory("4/2").getCount());
    Assert.assertEquals(1, dirMgrRoot.getDirectory("4/3").getCount());

    LocalResourceRequest lr4 = createLocalResourceRequest(user, 4, 4,
        LocalResourceVisibility.PUBLIC);
    Assert.assertNull(tracker.getLocalizedResource(lr4));
    final long localizedId4 = 256;
    Path hierarchicalPath4 = new Path(localDirRoot + "/4",
        Long.toString(localizedId4));
    Path localizedPath4 = new Path(hierarchicalPath4, "resource.jar");
    tracker.handle(new ResourceRecoveredEvent(lr4, localizedPath4, 120));
    dispatcher.await();
    Assert.assertNotNull(tracker.getLocalizedResource(lr4));
    Assert.assertEquals(0, dirMgrRoot.getDirectory("").getCount());
    Assert.assertEquals(1, dirMgrRoot.getDirectory("4").getCount());
    Assert.assertEquals(2, dirMgrRoot.getDirectory("4/2").getCount());
    Assert.assertEquals(1, dirMgrRoot.getDirectory("4/3").getCount());
  } finally {
    if (dispatcher != null) {
      dispatcher.stop();
    }
  }
}
 

@Test (timeout = 60000)
public void testDecomissionedNMsMetricsOnRMRestart() throws Exception {
  YarnConfiguration conf = new YarnConfiguration();
  conf.set(YarnConfiguration.RM_NODES_EXCLUDE_FILE_PATH,
    hostFile.getAbsolutePath());
  writeToHostsFile("");
  final DrainDispatcher dispatcher = new DrainDispatcher();
  MockRM rm1 = null, rm2 = null;
  try {
    rm1 = new MockRM(conf) {
      @Override
      protected Dispatcher createDispatcher() {
        return dispatcher;
      }
    };
    rm1.start();
    MockNM nm1 = rm1.registerNode("localhost:1234", 8000);
    MockNM nm2 = rm1.registerNode("host2:1234", 8000);
    Assert
        .assertEquals(0,
            ClusterMetrics.getMetrics().getNumDecommisionedNMs());
    String ip = NetUtils.normalizeHostName("localhost");
    // Add 2 hosts to exclude list.
    writeToHostsFile("host2", ip);

    // refresh nodes
    rm1.getNodesListManager().refreshNodes(conf);
    NodeHeartbeatResponse nodeHeartbeat = nm1.nodeHeartbeat(true);
    Assert
        .assertTrue(
            NodeAction.SHUTDOWN.equals(nodeHeartbeat.getNodeAction()));
    nodeHeartbeat = nm2.nodeHeartbeat(true);
    Assert.assertTrue("The decommisioned metrics are not updated",
        NodeAction.SHUTDOWN.equals(nodeHeartbeat.getNodeAction()));

    dispatcher.await();
    Assert
        .assertEquals(2,
            ClusterMetrics.getMetrics().getNumDecommisionedNMs());
    rm1.stop();
    rm1 = null;
    Assert
        .assertEquals(0,
            ClusterMetrics.getMetrics().getNumDecommisionedNMs());

    // restart RM.
    rm2 = new MockRM(conf);
    rm2.start();
    Assert
        .assertEquals(2,
            ClusterMetrics.getMetrics().getNumDecommisionedNMs());
  } finally {
    if (rm1 != null) {
      rm1.stop();
    }
    if (rm2 != null) {
      rm2.stop();
    }
  }
}
 

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

    DeletionService mockDelService = mock(DeletionService.class);

    ContainerId cId1 = BuilderUtils.newContainerId(1, 1, 1, 1);
    LocalizerContext lc1 = new LocalizerContext(user, cId1, null);
    ContainerId cId2 = BuilderUtils.newContainerId(1, 1, 1, 2);
    LocalizerContext lc2 = new LocalizerContext(user, cId2, null);

    LocalResourceRequest req1 =
        createLocalResourceRequest(user, 1, 1, LocalResourceVisibility.PUBLIC);
    LocalResourceRequest req2 =
        createLocalResourceRequest(user, 2, 1, LocalResourceVisibility.PUBLIC);
    LocalizedResource lr1 = createLocalizedResource(req1, dispatcher);
    LocalizedResource lr2 = createLocalizedResource(req2, dispatcher);
    ConcurrentMap<LocalResourceRequest, LocalizedResource> localrsrc =
        new ConcurrentHashMap<LocalResourceRequest, LocalizedResource>();
    localrsrc.put(req1, lr1);
    localrsrc.put(req2, lr2);
    LocalResourcesTracker tracker =
        new LocalResourcesTrackerImpl(user, null, dispatcher, localrsrc,
            false, conf, new NMNullStateStoreService());

    ResourceEvent req11Event =
        new ResourceRequestEvent(req1, LocalResourceVisibility.PUBLIC, lc1);
    ResourceEvent req12Event =
        new ResourceRequestEvent(req1, LocalResourceVisibility.PUBLIC, lc2);
    ResourceEvent req21Event =
        new ResourceRequestEvent(req2, LocalResourceVisibility.PUBLIC, lc1);

    ResourceEvent rel11Event = new ResourceReleaseEvent(req1, cId1);
    ResourceEvent rel12Event = new ResourceReleaseEvent(req1, cId2);
    ResourceEvent rel21Event = new ResourceReleaseEvent(req2, cId1);

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

    // Localize R1 for C2
    tracker.handle(req12Event);

    // Localize R2 for C1
    tracker.handle(req21Event);

    dispatcher.await();
    verify(localizerEventHandler, times(3)).handle(
        any(LocalizerResourceRequestEvent.class));
    // Verify refCount for R1 is 2
    Assert.assertEquals(2, lr1.getRefCount());
    // Verify refCount for R2 is 1
    Assert.assertEquals(1, lr2.getRefCount());

    // Release R2 for C1
    tracker.handle(rel21Event);

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

    // Verify resource with non zero ref count is not removed.
    Assert.assertEquals(2, lr1.getRefCount());
    Assert.assertFalse(tracker.remove(lr1, mockDelService));
    verifyTrackedResourceCount(tracker, 2);

    // Localize resource1
    ResourceLocalizedEvent rle =
        new ResourceLocalizedEvent(req1, new Path("file:///tmp/r1"), 1);
    lr1.handle(rle);
    Assert.assertTrue(lr1.getState().equals(ResourceState.LOCALIZED));

    // Release resource1
    tracker.handle(rel11Event);
    tracker.handle(rel12Event);
    Assert.assertEquals(0, lr1.getRefCount());

    // Verify resources in state LOCALIZED with ref-count=0 is removed.
    Assert.assertTrue(tracker.remove(lr1, mockDelService));
    verifyTrackedResourceCount(tracker, 1);
  } finally {
    if (dispatcher != null) {
      dispatcher.stop();
    }
  }
}