@Test
public void testRMContainerAllocatorResendsRequestsOnRMRestart()
    throws Exception {

  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.RECOVERY_ENABLED, "true");
  conf.set(YarnConfiguration.RM_STORE, MemoryRMStateStore.class.getName());
  conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS,
      YarnConfiguration.DEFAULT_RM_AM_MAX_ATTEMPTS);
  conf.setBoolean(YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_ENABLED, true);
  conf.setLong(YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_SCHEDULING_WAIT_MS, 0);
  conf.setBoolean(MRJobConfig.MR_AM_JOB_NODE_BLACKLISTING_ENABLE, true);
  conf.setInt(MRJobConfig.MAX_TASK_FAILURES_PER_TRACKER, 1);
  conf.setInt(
      MRJobConfig.MR_AM_IGNORE_BLACKLISTING_BLACKLISTED_NODE_PERECENT, -1);
  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);
  MyResourceManager rm1 = new MyResourceManager(conf, memStore);
  rm1.start();
  DrainDispatcher dispatcher =
      (DrainDispatcher) rm1.getRMContext().getDispatcher();
  // Submit the application
  RMApp app = rm1.submitApp(1024);
  dispatcher.await();
  MockNM nm1 = new MockNM("h1:1234", 15120, rm1.getResourceTrackerService());
  nm1.registerNode();
  nm1.nodeHeartbeat(true); // Node heartbeat
  dispatcher.await();
  ApplicationAttemptId appAttemptId =
      app.getCurrentAppAttempt().getAppAttemptId();
  rm1.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(rm1, conf, appAttemptId, mockJob);
  // Step-1 : AM send allocate request for 2 ContainerRequests and 1
  // blackListeNode
  // create the container request
  // send MAP request
  ContainerRequestEvent event1 =
      createReq(jobId, 1, 1024, new String[] { "h1" });
  allocator.sendRequest(event1);
  ContainerRequestEvent event2 =
      createReq(jobId, 2, 2048, new String[] { "h1", "h2" });
  allocator.sendRequest(event2);
  // Send events to blacklist h2
  ContainerFailedEvent f1 = createFailEvent(jobId, 1, "h2", false);
  allocator.sendFailure(f1);
  // send allocate request and 1 blacklisted nodes
  List<TaskAttemptContainerAssignedEvent> assignedContainers =
      allocator.schedule();
  dispatcher.await();
  Assert.assertEquals("No of assignments must be 0", 0,
      assignedContainers.size());
  // Why ask is 3, not 4? --> ask from blacklisted node h2 is removed
  assertAsksAndReleases(3, 0, rm1);
  assertBlacklistAdditionsAndRemovals(1, 0, rm1);
  nm1.nodeHeartbeat(true); // Node heartbeat
  dispatcher.await();
  // Step-2 : 2 containers are allocated by RM.
  assignedContainers = allocator.schedule();
  dispatcher.await();
  Assert.assertEquals("No of assignments must be 2", 2,
      assignedContainers.size());
  assertAsksAndReleases(0, 0, rm1);
  assertBlacklistAdditionsAndRemovals(0, 0, rm1);
  assignedContainers = allocator.schedule();
  Assert.assertEquals("No of assignments must be 0", 0,
      assignedContainers.size());
  assertAsksAndReleases(3, 0, rm1);
  assertBlacklistAdditionsAndRemovals(0, 0, rm1);
  // Step-3 : AM Send 1 containerRequest(event3) and 1 releaseRequests to
  // RM
  // send container request
  ContainerRequestEvent event3 =
      createReq(jobId, 3, 1000, new String[] { "h1" });
  allocator.sendRequest(event3);
  // send deallocate request
  ContainerAllocatorEvent deallocate1 =
      createDeallocateEvent(jobId, 1, false);
  allocator.sendDeallocate(deallocate1);
  assignedContainers = allocator.schedule();
  Assert.assertEquals("No of assignments must be 0", 0,
      assignedContainers.size());
  assertAsksAndReleases(3, 1, rm1);
  assertBlacklistAdditionsAndRemovals(0, 0, rm1);
  // Phase-2 start 2nd RM is up
  MyResourceManager rm2 = new MyResourceManager(conf, memStore);
  rm2.start();
  nm1.setResourceTrackerService(rm2.getResourceTrackerService());
  allocator.updateSchedulerProxy(rm2);
  dispatcher = (DrainDispatcher) rm2.getRMContext().getDispatcher();
  // NM should be rebooted on heartbeat, even first heartbeat for nm2
  NodeHeartbeatResponse hbResponse = nm1.nodeHeartbeat(true);
  Assert.assertEquals(NodeAction.RESYNC, hbResponse.getNodeAction());
  // new NM to represent NM re-register
  nm1 = new MockNM("h1:1234", 10240, rm2.getResourceTrackerService());
  nm1.registerNode();
  nm1.nodeHeartbeat(true);
  dispatcher.await();
  // Step-4 : On RM restart, AM(does not know RM is restarted) sends
  // additional containerRequest(event4) and blacklisted nodes.
  // Intern RM send resync command
  // send deallocate request, release=1
  ContainerAllocatorEvent deallocate2 =
      createDeallocateEvent(jobId, 2, false);
  allocator.sendDeallocate(deallocate2);
  // Send events to blacklist nodes h3
  ContainerFailedEvent f2 = createFailEvent(jobId, 1, "h3", false);
  allocator.sendFailure(f2);
  ContainerRequestEvent event4 =
      createReq(jobId, 4, 2000, new String[] { "h1", "h2" });
  allocator.sendRequest(event4);
  // send allocate request to 2nd RM and get resync command
  allocator.schedule();
  dispatcher.await();
  // Step-5 : On Resync,AM sends all outstanding
  // asks,release,blacklistAaddition
  // and another containerRequest(event5)
  ContainerRequestEvent event5 =
      createReq(jobId, 5, 3000, new String[] { "h1", "h2", "h3" });
  allocator.sendRequest(event5);
  // send all outstanding request again.
  assignedContainers = allocator.schedule();
  dispatcher.await();
  assertAsksAndReleases(3, 2, rm2);
  assertBlacklistAdditionsAndRemovals(2, 0, rm2);
  nm1.nodeHeartbeat(true);
  dispatcher.await();
  // Step-6 : RM allocates containers i.e event3,event4 and cRequest5
  assignedContainers = allocator.schedule();
  dispatcher.await();
  Assert.assertEquals("Number of container should be 3", 3,
      assignedContainers.size());
  for (TaskAttemptContainerAssignedEvent assig : assignedContainers) {
    Assert.assertTrue("Assigned count not correct",
        "h1".equals(assig.getContainer().getNodeId().getHost()));
  }
  rm1.stop();
  rm2.stop();
}