MockRM rm = new MockRM();
try {
rm.start();
MockNM nm1 = rm.registerNode("h1:1234", 10000);
NMTokenSecretManagerInRM nmTokenSecretManager =
rm.getRMContext().getNMTokenSecretManager();
// submitting new application
RMApp app = rm.submitApp(1000);
// start scheduling.
nm1.nodeHeartbeat(true);
// Starting application attempt and launching
// It should get registered with NMTokenSecretManager.
RMAppAttempt attempt = app.getCurrentAppAttempt();
MockAM am = rm.sendAMLaunched(attempt.getAppAttemptId());
Assert.assertTrue(nmTokenSecretManager
.isApplicationAttemptRegistered(attempt.getAppAttemptId()));
// This will register application master.
am.registerAppAttempt();
ArrayList<Container> containersReceivedForNM1 =
new ArrayList<Container>();
List<ContainerId> releaseContainerList =
new ArrayList<ContainerId>();
HashMap<String, Token> nmTokens = new HashMap<String, Token>();
// initially requesting 2 containers.
AllocateResponse response =
am.allocate("h1", 1000, 2, releaseContainerList);
Assert.assertEquals(0, response.getAllocatedContainers().size());
allocateContainersAndValidateNMTokens(am, containersReceivedForNM1, 2,
nmTokens, nm1);
Assert.assertEquals(1, nmTokens.size());
// requesting 2 more containers.
response = am.allocate("h1", 1000, 2, releaseContainerList);
Assert.assertEquals(0, response.getAllocatedContainers().size());
allocateContainersAndValidateNMTokens(am, containersReceivedForNM1, 4,
nmTokens, nm1);
Assert.assertEquals(1, nmTokens.size());
// We will be simulating NM restart so restarting newly added h2:1234
// NM 2 now registers.
MockNM nm2 = rm.registerNode("h2:1234", 10000);
nm2.nodeHeartbeat(true);
ArrayList<Container> containersReceivedForNM2 =
new ArrayList<Container>();
response = am.allocate("h2", 1000, 2, releaseContainerList);
Assert.assertEquals(0, response.getAllocatedContainers().size());
allocateContainersAndValidateNMTokens(am, containersReceivedForNM2, 2,
nmTokens, nm2);
Assert.assertEquals(2, nmTokens.size());
// Simulating NM-2 restart.
nm2 = rm.registerNode("h2:1234", 10000);
// Wait for reconnect to make it through the RM and create a new RMNode
Map<NodeId, RMNode> nodes = rm.getRMContext().getRMNodes();
while (nodes.get(nm2.getNodeId()).getLastNodeHeartBeatResponse()
.getResponseId() > 0) {
Thread.sleep(WAIT_SLEEP_MS);
}
int interval = 40;
// Wait for nm Token to be cleared.
while (nmTokenSecretManager
.isApplicationAttemptNMTokenPresent(attempt.getAppAttemptId(),
nm2.getNodeId()) && interval-- > 0) {
LOG.info("waiting for nmToken to be cleared for : " + nm2.getNodeId());
Thread.sleep(WAIT_SLEEP_MS);
}
Assert.assertTrue(nmTokenSecretManager
.isApplicationAttemptRegistered(attempt.getAppAttemptId()));
// removing NMToken for h2:1234
nmTokens.remove(nm2.getNodeId().toString());
Assert.assertEquals(1, nmTokens.size());
// We should again receive the NMToken.
response = am.allocate("h2", 1000, 2, releaseContainerList);
Assert.assertEquals(0, response.getAllocatedContainers().size());
allocateContainersAndValidateNMTokens(am, containersReceivedForNM2, 4,
nmTokens, nm2);
Assert.assertEquals(2, nmTokens.size());
// Now rolling over NMToken masterKey. it should resend the NMToken in
// next allocate call.
Assert.assertTrue(nmTokenSecretManager
.isApplicationAttemptNMTokenPresent(attempt.getAppAttemptId(),
nm1.getNodeId()));
Assert.assertTrue(nmTokenSecretManager
.isApplicationAttemptNMTokenPresent(attempt.getAppAttemptId(),
nm2.getNodeId()));
nmTokenSecretManager.rollMasterKey();
nmTokenSecretManager.activateNextMasterKey();
Assert.assertFalse(nmTokenSecretManager
.isApplicationAttemptNMTokenPresent(attempt.getAppAttemptId(),
nm1.getNodeId()));
Assert.assertFalse(nmTokenSecretManager
.isApplicationAttemptNMTokenPresent(attempt.getAppAttemptId(),
nm2.getNodeId()));
// It should not remove application attempt entry.
Assert.assertTrue(nmTokenSecretManager
.isApplicationAttemptRegistered(attempt.getAppAttemptId()));
nmTokens.clear();
Assert.assertEquals(0, nmTokens.size());
// We should again receive the NMToken.
response = am.allocate("h2", 1000, 1, releaseContainerList);
Assert.assertEquals(0, response.getAllocatedContainers().size());
allocateContainersAndValidateNMTokens(am, containersReceivedForNM2, 5,
nmTokens, nm2);
Assert.assertEquals(1, nmTokens.size());
Assert.assertTrue(nmTokenSecretManager
.isApplicationAttemptNMTokenPresent(attempt.getAppAttemptId(),
nm2.getNodeId()));
// After AM is finished making sure that nmtoken entry for app
Assert.assertTrue(nmTokenSecretManager
.isApplicationAttemptRegistered(attempt.getAppAttemptId()));
am.unregisterAppAttempt();
// marking all the containers as finished.
for (Container container : containersReceivedForNM1) {
nm1.nodeHeartbeat(attempt.getAppAttemptId(), container.getId().getId(),
ContainerState.COMPLETE);
}
for (Container container : containersReceivedForNM2) {
nm2.nodeHeartbeat(attempt.getAppAttemptId(), container.getId().getId(),
ContainerState.COMPLETE);
}
nm1.nodeHeartbeat(am.getApplicationAttemptId(), 1,
ContainerState.COMPLETE);
am.waitForState(RMAppAttemptState.FINISHED);
Assert.assertFalse(nmTokenSecretManager
.isApplicationAttemptRegistered(attempt.getAppAttemptId()));
} finally {
rm.stop();
}
}