@Test(timeout = 30000)
public void testRecoverRequestAfterPreemption() throws Exception {
Configuration conf = new Configuration();
conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
ResourceScheduler.class);
MockRM rm1 = new MockRM(conf);
rm1.start();
MockNM nm1 = rm1.registerNode("127.0.0.1:1234", 8000);
RMApp app1 = rm1.submitApp(1024);
MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);
CapacityScheduler cs = (CapacityScheduler) rm1.getResourceScheduler();
// request a container.
am1.allocate("127.0.0.1", 1024, 1, new ArrayList<ContainerId>());
ContainerId containerId1 = ContainerId.newInstance(am1
.getApplicationAttemptId(), 2);
rm1.waitForState(nm1, containerId1, RMContainerState.ALLOCATED);
RMContainer rmContainer = cs.getRMContainer(containerId1);
List<ResourceRequest> requests = rmContainer.getResourceRequests();
FiCaSchedulerApp app = cs.getApplicationAttempt(am1
.getApplicationAttemptId());
FiCaSchedulerNode node = cs.getNode(rmContainer.getAllocatedNode());
for (ResourceRequest request : requests) {
// Skip the OffRack and RackLocal resource requests.
if (request.getResourceName().equals(node.getRackName())
|| request.getResourceName().equals(ResourceRequest.ANY)) {
continue;
}
// Already the node local resource request is cleared from RM after
// allocation.
Assert.assertNull(app.getResourceRequest(request.getPriority(), request
.getResourceName()));
}
// Call killContainer to preempt the container
cs.killContainer(rmContainer);
Assert.assertEquals(3, requests.size());
for (ResourceRequest request : requests) {
// Resource request must have added back in RM after preempt event
// handling.
Assert.assertEquals(1, app.getResourceRequest(request.getPriority(),
request.getResourceName()).getNumContainers());
}
// New container will be allocated and will move to ALLOCATED state
ContainerId containerId2 = ContainerId.newInstance(am1
.getApplicationAttemptId(), 3);
rm1.waitForState(nm1, containerId2, RMContainerState.ALLOCATED);
// allocate container
List<Container> containers = am1.allocate(new ArrayList<ResourceRequest>(),
new ArrayList<ContainerId>()).getAllocatedContainers();