Assert.assertEquals(0, NMTokenCache.numberOfNMTokensInCache());
HashMap<String, Token> receivedNMTokens = new HashMap<String, Token>();
while (allocatedContainerCount < containersRequestedAny
&& iterationsLeft-- > 0) {
AllocateResponse allocResponse = amClient.allocate(0.1f);
assertTrue(amClient.ask.size() == 0);
assertTrue(amClient.release.size() == 0);
assertTrue(nodeCount == amClient.getClusterNodeCount());
allocatedContainerCount += allocResponse.getAllocatedContainers().size();
for(Container container : allocResponse.getAllocatedContainers()) {
ContainerId rejectContainerId = container.getId();
releases.add(rejectContainerId);
amClient.releaseAssignedContainer(rejectContainerId);
}
for (NMToken token : allocResponse.getNMTokens()) {
String nodeID = token.getNodeId().toString();
if (receivedNMTokens.containsKey(nodeID)) {
Assert.fail("Received token again for : " + nodeID);
}
receivedNMTokens.put(nodeID, token.getToken());
}
if(allocatedContainerCount < containersRequestedAny) {
// sleep to let NM's heartbeat to RM and trigger allocations
sleep(100);
}
}
// Should receive atleast 1 token
Assert.assertTrue(receivedNMTokens.size() > 0
&& receivedNMTokens.size() <= nodeCount);
assertTrue(allocatedContainerCount == containersRequestedAny);
assertTrue(amClient.release.size() == 2);
assertTrue(amClient.ask.size() == 0);
// need to tell the AMRMClient that we dont need these resources anymore
amClient.removeContainerRequest(
new ContainerRequest(capability, nodes, racks, priority));
amClient.removeContainerRequest(
new ContainerRequest(capability, nodes, racks, priority));
assertTrue(amClient.ask.size() == 3);
// send 0 container count request for resources that are no longer needed
ResourceRequest snoopRequest = amClient.ask.iterator().next();
assertTrue(snoopRequest.getNumContainers() == 0);
// test RPC exception handling
amClient.addContainerRequest(new ContainerRequest(capability, nodes,
racks, priority));
amClient.addContainerRequest(new ContainerRequest(capability, nodes,
racks, priority));
snoopRequest = amClient.ask.iterator().next();
assertTrue(snoopRequest.getNumContainers() == 2);
ApplicationMasterProtocol realRM = amClient.rmClient;
try {
ApplicationMasterProtocol mockRM = mock(ApplicationMasterProtocol.class);
when(mockRM.allocate(any(AllocateRequest.class))).thenAnswer(
new Answer<AllocateResponse>() {
public AllocateResponse answer(InvocationOnMock invocation)
throws Exception {
amClient.removeContainerRequest(
new ContainerRequest(capability, nodes,
racks, priority));
amClient.removeContainerRequest(
new ContainerRequest(capability, nodes, racks, priority));
throw new Exception();
}
});
amClient.rmClient = mockRM;
amClient.allocate(0.1f);
}catch (Exception ioe) {}
finally {
amClient.rmClient = realRM;
}
assertTrue(amClient.release.size() == 2);
assertTrue(amClient.ask.size() == 3);
snoopRequest = amClient.ask.iterator().next();
// verify that the remove request made in between makeRequest and allocate
// has not been lost
assertTrue(snoopRequest.getNumContainers() == 0);
iterationsLeft = 3;
// do a few iterations to ensure RM is not going send new containers
while(!releases.isEmpty() || iterationsLeft-- > 0) {
// inform RM of rejection
AllocateResponse allocResponse = amClient.allocate(0.1f);
// RM did not send new containers because AM does not need any
assertTrue(allocResponse.getAllocatedContainers().size() == 0);
if(allocResponse.getCompletedContainersStatuses().size() > 0) {
for(ContainerStatus cStatus :allocResponse
.getCompletedContainersStatuses()) {
if(releases.contains(cStatus.getContainerId())) {
assertTrue(cStatus.getState() == ContainerState.COMPLETE);
assertTrue(cStatus.getExitStatus() == -100);
releases.remove(cStatus.getContainerId());