/// <exception cref="System.Exception"/> protected internal virtual void AllocateContainersAndValidateNMTokens(MockAM am, AList <Container> containersReceived, int totalContainerRequested, Dictionary <string , Token> nmTokens, MockNM nm) { AList <ContainerId> releaseContainerList = new AList <ContainerId>(); AllocateResponse response; AList <ResourceRequest> resourceRequest = new AList <ResourceRequest>(); while (containersReceived.Count < totalContainerRequested) { nm.NodeHeartbeat(true); Log.Info("requesting containers.."); response = am.Allocate(resourceRequest, releaseContainerList); Sharpen.Collections.AddAll(containersReceived, response.GetAllocatedContainers()); if (!response.GetNMTokens().IsEmpty()) { foreach (NMToken nmToken in response.GetNMTokens()) { string nodeId = nmToken.GetNodeId().ToString(); if (nmTokens.Contains(nodeId)) { NUnit.Framework.Assert.Fail("Duplicate NMToken received for : " + nodeId); } nmTokens[nodeId] = nmToken.GetToken(); } } Log.Info("Got " + containersReceived.Count + " containers. Waiting to get " + totalContainerRequested ); Sharpen.Thread.Sleep(WaitSleepMs); } }
/// <exception cref="System.Exception"/> public virtual void TestNMToken() { MockRM rm = new MockRM(conf); 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()); NUnit.Framework.Assert.IsTrue(nmTokenSecretManager.IsApplicationAttemptRegistered (attempt.GetAppAttemptId())); // This will register application master. am.RegisterAppAttempt(); AList <Container> containersReceivedForNM1 = new AList <Container>(); IList <ContainerId> releaseContainerList = new AList <ContainerId>(); Dictionary <string, Token> nmTokens = new Dictionary <string, Token>(); // initially requesting 2 containers. AllocateResponse response = am.Allocate("h1", 1000, 2, releaseContainerList); NUnit.Framework.Assert.AreEqual(0, response.GetAllocatedContainers().Count); AllocateContainersAndValidateNMTokens(am, containersReceivedForNM1, 2, nmTokens, nm1); NUnit.Framework.Assert.AreEqual(1, nmTokens.Count); // requesting 2 more containers. response = am.Allocate("h1", 1000, 2, releaseContainerList); NUnit.Framework.Assert.AreEqual(0, response.GetAllocatedContainers().Count); AllocateContainersAndValidateNMTokens(am, containersReceivedForNM1, 4, nmTokens, nm1); NUnit.Framework.Assert.AreEqual(1, nmTokens.Count); // 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); AList <Container> containersReceivedForNM2 = new AList <Container>(); response = am.Allocate("h2", 1000, 2, releaseContainerList); NUnit.Framework.Assert.AreEqual(0, response.GetAllocatedContainers().Count); AllocateContainersAndValidateNMTokens(am, containersReceivedForNM2, 2, nmTokens, nm2); NUnit.Framework.Assert.AreEqual(2, nmTokens.Count); // Simulating NM-2 restart. nm2 = rm.RegisterNode("h2:1234", 10000); // Wait for reconnect to make it through the RM and create a new RMNode IDictionary <NodeId, RMNode> nodes = rm.GetRMContext().GetRMNodes(); while (nodes[nm2.GetNodeId()].GetLastNodeHeartBeatResponse().GetResponseId() > 0) { Sharpen.Thread.Sleep(WaitSleepMs); } 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()); Sharpen.Thread.Sleep(WaitSleepMs); } NUnit.Framework.Assert.IsTrue(nmTokenSecretManager.IsApplicationAttemptRegistered (attempt.GetAppAttemptId())); // removing NMToken for h2:1234 Sharpen.Collections.Remove(nmTokens, nm2.GetNodeId().ToString()); NUnit.Framework.Assert.AreEqual(1, nmTokens.Count); // We should again receive the NMToken. response = am.Allocate("h2", 1000, 2, releaseContainerList); NUnit.Framework.Assert.AreEqual(0, response.GetAllocatedContainers().Count); AllocateContainersAndValidateNMTokens(am, containersReceivedForNM2, 4, nmTokens, nm2); NUnit.Framework.Assert.AreEqual(2, nmTokens.Count); // Now rolling over NMToken masterKey. it should resend the NMToken in // next allocate call. NUnit.Framework.Assert.IsTrue(nmTokenSecretManager.IsApplicationAttemptNMTokenPresent (attempt.GetAppAttemptId(), nm1.GetNodeId())); NUnit.Framework.Assert.IsTrue(nmTokenSecretManager.IsApplicationAttemptNMTokenPresent (attempt.GetAppAttemptId(), nm2.GetNodeId())); nmTokenSecretManager.RollMasterKey(); nmTokenSecretManager.ActivateNextMasterKey(); NUnit.Framework.Assert.IsFalse(nmTokenSecretManager.IsApplicationAttemptNMTokenPresent (attempt.GetAppAttemptId(), nm1.GetNodeId())); NUnit.Framework.Assert.IsFalse(nmTokenSecretManager.IsApplicationAttemptNMTokenPresent (attempt.GetAppAttemptId(), nm2.GetNodeId())); // It should not remove application attempt entry. NUnit.Framework.Assert.IsTrue(nmTokenSecretManager.IsApplicationAttemptRegistered (attempt.GetAppAttemptId())); nmTokens.Clear(); NUnit.Framework.Assert.AreEqual(0, nmTokens.Count); // We should again receive the NMToken. response = am.Allocate("h2", 1000, 1, releaseContainerList); NUnit.Framework.Assert.AreEqual(0, response.GetAllocatedContainers().Count); AllocateContainersAndValidateNMTokens(am, containersReceivedForNM2, 5, nmTokens, nm2); NUnit.Framework.Assert.AreEqual(1, nmTokens.Count); NUnit.Framework.Assert.IsTrue(nmTokenSecretManager.IsApplicationAttemptNMTokenPresent (attempt.GetAppAttemptId(), nm2.GetNodeId())); // After AM is finished making sure that nmtoken entry for app NUnit.Framework.Assert.IsTrue(nmTokenSecretManager.IsApplicationAttemptRegistered (attempt.GetAppAttemptId())); am.UnregisterAppAttempt(); // marking all the containers as finished. foreach (Container container in containersReceivedForNM1) { nm1.NodeHeartbeat(attempt.GetAppAttemptId(), container.GetId().GetContainerId(), ContainerState.Complete); } foreach (Container container_1 in containersReceivedForNM2) { nm2.NodeHeartbeat(attempt.GetAppAttemptId(), container_1.GetId().GetContainerId() , ContainerState.Complete); } nm1.NodeHeartbeat(am.GetApplicationAttemptId(), 1, ContainerState.Complete); am.WaitForState(RMAppAttemptState.Finished); NUnit.Framework.Assert.IsFalse(nmTokenSecretManager.IsApplicationAttemptRegistered (attempt.GetAppAttemptId())); } finally { rm.Stop(); } }