diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/protocolrecords/RegisterApplicationMasterResponse.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/protocolrecords/RegisterApplicationMasterResponse.java index 363d872..79f9f3a 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/protocolrecords/RegisterApplicationMasterResponse.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/protocolrecords/RegisterApplicationMasterResponse.java @@ -29,6 +29,7 @@ import org.apache.hadoop.yarn.api.ApplicationMasterProtocol; import org.apache.hadoop.yarn.api.records.ApplicationAccessType; import org.apache.hadoop.yarn.api.records.Container; +import org.apache.hadoop.yarn.api.records.NMToken; import org.apache.hadoop.yarn.api.records.Resource; import org.apache.hadoop.yarn.util.Records; @@ -55,13 +56,15 @@ public static RegisterApplicationMasterResponse newInstance( Resource minCapability, Resource maxCapability, Map acls, ByteBuffer key, - List containersFromPreviousAttempt, String queue) { + List containersFromPreviousAttempt, String queue, + List nmTokensFromPreviousAttempts) { RegisterApplicationMasterResponse response = Records.newRecord(RegisterApplicationMasterResponse.class); response.setMaximumResourceCapability(maxCapability); response.setApplicationACLs(acls); response.setClientToAMTokenMasterKey(key); - response.setContainersFromPreviousAttempt(containersFromPreviousAttempt); + response.setContainersFromPreviousAttempts(containersFromPreviousAttempt); + response.setNMTokensFromPreviousAttempts(nmTokensFromPreviousAttempts); response.setQueue(queue); return response; } @@ -129,26 +132,52 @@ public static RegisterApplicationMasterResponse newInstance( /** *

* Get the list of running containers as viewed by - * ResourceManager from previous application attempt. + * ResourceManager from previous application attempts. *

* * @return the list of running containers as viewed by - * ResourceManager from previous application attempt + * ResourceManager from previous application attempts + * @see RegisterApplicationMasterResponse#getNMTokensFromPreviousAttempts() */ @Public @Unstable - public abstract List getContainersFromPreviousAttempt(); + public abstract List getContainersFromPreviousAttempts(); /** * Set the list of running containers as viewed by - * ResourceManager from previous application attempt. + * ResourceManager from previous application attempts. * * @param containersFromPreviousAttempt * the list of running containers as viewed by - * ResourceManager from previous application attempt. + * ResourceManager from previous application attempts. */ @Private @Unstable - public abstract void setContainersFromPreviousAttempt( + public abstract void setContainersFromPreviousAttempts( List containersFromPreviousAttempt); + + /** + * Get the list of NMTokens for communicating with the NMs where the + * containers of previous application attempts are running. + * + * @return the list of NMTokens for communicating with the NMs where the + * containers of previous application attempts are running. + * + * @see RegisterApplicationMasterResponse#getContainersFromPreviousAttempts() + */ + @Public + @Stable + public abstract List getNMTokensFromPreviousAttempts(); + + /** + * Set the list of NMTokens for communicating with the NMs where the the + * containers of previous application attempts are running. + * + * @param nmTokens + * the list of NMTokens for communicating with the NMs where the + * containers of previous application attempts are running. + */ + @Private + @Unstable + public abstract void setNMTokensFromPreviousAttempts(List nmTokens); } diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_service_protos.proto hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_service_protos.proto index eff5cd7..a1f6d2e 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_service_protos.proto +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_service_protos.proto @@ -44,8 +44,9 @@ message RegisterApplicationMasterResponseProto { optional ResourceProto maximumCapability = 1; optional bytes client_to_am_token_master_key = 2; repeated ApplicationACLMapProto application_ACLs = 3; - repeated ContainerProto containers_from_previous_attempt = 4; + repeated ContainerProto containers_from_previous_attempts = 4; optional string queue = 5; + repeated NMTokenProto nm_tokens_from_previous_attempts = 6; } message FinishApplicationMasterRequestProto { diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-applications-distributedshell/src/main/java/org/apache/hadoop/yarn/applications/distributedshell/ApplicationMaster.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-applications-distributedshell/src/main/java/org/apache/hadoop/yarn/applications/distributedshell/ApplicationMaster.java index 8d869a2..775df36 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-applications-distributedshell/src/main/java/org/apache/hadoop/yarn/applications/distributedshell/ApplicationMaster.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-applications-distributedshell/src/main/java/org/apache/hadoop/yarn/applications/distributedshell/ApplicationMaster.java @@ -542,7 +542,7 @@ public void run() throws YarnException, IOException { } List previousAMRunningContainers = - response.getContainersFromPreviousAttempt(); + response.getContainersFromPreviousAttempts(); LOG.info("Received " + previousAMRunningContainers.size() + " previous AM's running containers on AM registration."); numAllocatedContainers.addAndGet(previousAMRunningContainers.size()); diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/protocolrecords/impl/pb/RegisterApplicationMasterResponsePBImpl.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/protocolrecords/impl/pb/RegisterApplicationMasterResponsePBImpl.java index ae488c4..06a637a 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/protocolrecords/impl/pb/RegisterApplicationMasterResponsePBImpl.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/protocolrecords/impl/pb/RegisterApplicationMasterResponsePBImpl.java @@ -31,13 +31,16 @@ import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse; import org.apache.hadoop.yarn.api.records.ApplicationAccessType; import org.apache.hadoop.yarn.api.records.Container; +import org.apache.hadoop.yarn.api.records.NMToken; import org.apache.hadoop.yarn.api.records.Resource; import org.apache.hadoop.yarn.api.records.impl.pb.ContainerPBImpl; +import org.apache.hadoop.yarn.api.records.impl.pb.NMTokenPBImpl; import org.apache.hadoop.yarn.api.records.impl.pb.ProtoUtils; import org.apache.hadoop.yarn.api.records.impl.pb.ResourcePBImpl; import org.apache.hadoop.yarn.proto.YarnProtos.ApplicationACLMapProto; import org.apache.hadoop.yarn.proto.YarnProtos.ContainerProto; import org.apache.hadoop.yarn.proto.YarnProtos.ResourceProto; +import org.apache.hadoop.yarn.proto.YarnServiceProtos.NMTokenProto; import org.apache.hadoop.yarn.proto.YarnServiceProtos.RegisterApplicationMasterResponseProto; import org.apache.hadoop.yarn.proto.YarnServiceProtos.RegisterApplicationMasterResponseProtoOrBuilder; @@ -56,7 +59,8 @@ private Resource maximumResourceCapability; private Map applicationACLS = null; - private List containersFromPreviousAttempt = null; + private List containersFromPreviousAttempts = null; + private List nmTokens = null; public RegisterApplicationMasterResponsePBImpl() { builder = RegisterApplicationMasterResponseProto.newBuilder(); @@ -110,8 +114,13 @@ private void mergeLocalToBuilder() { if (this.applicationACLS != null) { addApplicationACLs(); } - if (this.containersFromPreviousAttempt != null) { - addRunningContainersToProto(); + if (this.containersFromPreviousAttempts != null) { + addContainersFromPreviousAttemptToProto(); + } + if (nmTokens != null) { + builder.clearNmTokensFromPreviousAttempts(); + Iterable iterable = getTokenProtoIterable(nmTokens); + builder.addAllNmTokensFromPreviousAttempts(iterable); } } @@ -236,21 +245,22 @@ public ByteBuffer getClientToAMTokenMasterKey() { } @Override - public List getContainersFromPreviousAttempt() { - if (this.containersFromPreviousAttempt != null) { - return this.containersFromPreviousAttempt; + public List getContainersFromPreviousAttempts() { + if (this.containersFromPreviousAttempts != null) { + return this.containersFromPreviousAttempts; } - initRunningContainersList(); - return this.containersFromPreviousAttempt; + initContainersPreviousAttemptList(); + return this.containersFromPreviousAttempts; } @Override - public void setContainersFromPreviousAttempt(final List containers) { + public void + setContainersFromPreviousAttempts(final List containers) { if (containers == null) { return; } - this.containersFromPreviousAttempt = new ArrayList(); - this.containersFromPreviousAttempt.addAll(containers); + this.containersFromPreviousAttempts = new ArrayList(); + this.containersFromPreviousAttempts.addAll(containers); } @Override @@ -272,25 +282,88 @@ public void setQueue(String queue) { } } - private void initRunningContainersList() { - RegisterApplicationMasterResponseProtoOrBuilder p = viaProto ? proto : builder; - List list = p.getContainersFromPreviousAttemptList(); - containersFromPreviousAttempt = new ArrayList(); + + private void initContainersPreviousAttemptList() { + RegisterApplicationMasterResponseProtoOrBuilder p = + viaProto ? proto : builder; + List list = p.getContainersFromPreviousAttemptsList(); + containersFromPreviousAttempts = new ArrayList(); for (ContainerProto c : list) { - containersFromPreviousAttempt.add(convertFromProtoFormat(c)); + containersFromPreviousAttempts.add(convertFromProtoFormat(c)); } } - private void addRunningContainersToProto() { + private void addContainersFromPreviousAttemptToProto() { maybeInitBuilder(); - builder.clearContainersFromPreviousAttempt(); + builder.clearContainersFromPreviousAttempts(); List list = new ArrayList(); - for (Container c : containersFromPreviousAttempt) { + for (Container c : containersFromPreviousAttempts) { list.add(convertToProtoFormat(c)); } - builder.addAllContainersFromPreviousAttempt(list); + builder.addAllContainersFromPreviousAttempts(list); + } + + + @Override + public List getNMTokensFromPreviousAttempts() { + if (nmTokens != null) { + return nmTokens; + } + initLocalNewNMTokenList(); + return nmTokens; } + @Override + public void setNMTokensFromPreviousAttempts(final List nmTokens) { + if (nmTokens == null || nmTokens.isEmpty()) { + if (this.nmTokens != null) { + this.nmTokens.clear(); + } + builder.clearNmTokensFromPreviousAttempts(); + return; + } + this.nmTokens = new ArrayList(); + this.nmTokens.addAll(nmTokens); + } + + private synchronized void initLocalNewNMTokenList() { + RegisterApplicationMasterResponseProtoOrBuilder p = viaProto ? proto : builder; + List list = p.getNmTokensFromPreviousAttemptsList(); + nmTokens = new ArrayList(); + for (NMTokenProto t : list) { + nmTokens.add(convertFromProtoFormat(t)); + } + } + + private synchronized Iterable getTokenProtoIterable( + final List nmTokenList) { + maybeInitBuilder(); + return new Iterable() { + @Override + public synchronized Iterator iterator() { + return new Iterator() { + + Iterator iter = nmTokenList.iterator(); + + @Override + public boolean hasNext() { + return iter.hasNext(); + } + + @Override + public NMTokenProto next() { + return convertToProtoFormat(iter.next()); + } + + @Override + public void remove() { + throw new UnsupportedOperationException(); + } + }; + } + }; + } + private Resource convertFromProtoFormat(ResourceProto resource) { return new ResourcePBImpl(resource); } @@ -306,4 +379,12 @@ private ContainerPBImpl convertFromProtoFormat(ContainerProto p) { private ContainerProto convertToProtoFormat(Container t) { return ((ContainerPBImpl) t).getProto(); } + + private NMTokenProto convertToProtoFormat(NMToken token) { + return ((NMTokenPBImpl) token).getProto(); + } + + private NMToken convertFromProtoFormat(NMTokenProto proto) { + return new NMTokenPBImpl(proto); + } } diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NMTokenPBImpl.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NMTokenPBImpl.java index bcbf0a3..eec3343 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NMTokenPBImpl.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NMTokenPBImpl.java @@ -47,9 +47,9 @@ public NMTokenPBImpl(NMTokenProto proto) { this.proto = proto; viaProto = true; } - + @Override - public synchronized NodeId getNodeId() { + public NodeId getNodeId() { NMTokenProtoOrBuilder p = viaProto ? proto : builder; if (this.nodeId != null) { return nodeId; @@ -62,7 +62,7 @@ public synchronized NodeId getNodeId() { } @Override - public synchronized void setNodeId(NodeId nodeId) { + public void setNodeId(NodeId nodeId) { maybeInitBuilder(); if (nodeId == null) { builder.clearNodeId(); @@ -71,7 +71,7 @@ public synchronized void setNodeId(NodeId nodeId) { } @Override - public synchronized Token getToken() { + public Token getToken() { NMTokenProtoOrBuilder p = viaProto ? proto : builder; if (this.token != null) { return this.token; @@ -84,7 +84,7 @@ public synchronized Token getToken() { } @Override - public synchronized void setToken(Token token) { + public void setToken(Token token) { maybeInitBuilder(); if (token == null) { builder.clearToken(); @@ -92,14 +92,14 @@ public synchronized void setToken(Token token) { this.token = token; } - public synchronized NMTokenProto getProto() { + public NMTokenProto getProto() { mergeLocalToProto(); proto = viaProto ? proto : builder.build(); viaProto = true; return proto; } - private synchronized void mergeLocalToProto() { + private void mergeLocalToProto() { if (viaProto) { maybeInitBuilder(); } @@ -108,7 +108,7 @@ private synchronized void mergeLocalToProto() { viaProto = true; } - private synchronized void mergeLocalToBuilder() { + private void mergeLocalToBuilder() { if (this.nodeId != null) { builder.setNodeId(convertToProtoFormat(nodeId)); } @@ -117,26 +117,57 @@ private synchronized void mergeLocalToBuilder() { } } - private synchronized void maybeInitBuilder() { + private void maybeInitBuilder() { if(viaProto || builder == null) { builder = NMTokenProto.newBuilder(proto); } viaProto = false; } - private synchronized NodeId convertFromProtoFormat(NodeIdProto p) { + private NodeId convertFromProtoFormat(NodeIdProto p) { return new NodeIdPBImpl(p); } - private synchronized NodeIdProto convertToProtoFormat(NodeId nodeId) { + private NodeIdProto convertToProtoFormat(NodeId nodeId) { return ((NodeIdPBImpl)nodeId).getProto(); } - private synchronized TokenProto convertToProtoFormat(Token token) { + private TokenProto convertToProtoFormat(Token token) { return ((TokenPBImpl)token).getProto(); } - private synchronized Token convertFromProtoFormat(TokenProto proto) { + private Token convertFromProtoFormat(TokenProto proto) { return new TokenPBImpl(proto); } + + @Override + public int hashCode() { + final int prime = 31; + int result = 1; + result = prime * result + ((nodeId == null) ? 0 : nodeId.hashCode()); + result = prime * result + ((token == null) ? 0 : token.hashCode()); + return result; + } + + @Override + public boolean equals(Object obj) { + if (this == obj) + return true; + if (obj == null) + return false; + if (getClass() != obj.getClass()) + return false; + NMTokenPBImpl other = (NMTokenPBImpl) obj; + if (nodeId == null) { + if (other.nodeId != null) + return false; + } else if (!nodeId.equals(other.nodeId)) + return false; + if (token == null) { + if (other.token != null) + return false; + } else if (!token.equals(other.token)) + return false; + return true; + } } diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/ApplicationMasterService.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/ApplicationMasterService.java index db81dd8..2e1df5b 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/ApplicationMasterService.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/ApplicationMasterService.java @@ -51,6 +51,7 @@ import org.apache.hadoop.yarn.api.records.ApplicationId; import org.apache.hadoop.yarn.api.records.Container; import org.apache.hadoop.yarn.api.records.ContainerId; +import org.apache.hadoop.yarn.api.records.NMToken; import org.apache.hadoop.yarn.api.records.NodeReport; import org.apache.hadoop.yarn.api.records.PreemptionContainer; import org.apache.hadoop.yarn.api.records.PreemptionContract; @@ -280,10 +281,21 @@ public RegisterApplicationMasterResponse registerApplicationMaster( .getMasterKey(applicationAttemptId).getEncoded())); } - List containerList = + List transferredContainers = ((AbstractYarnScheduler) rScheduler) .getTransferredContainers(applicationAttemptId); - response.setContainersFromPreviousAttempt(containerList); + response.setContainersFromPreviousAttempts(transferredContainers); + if (!transferredContainers.isEmpty()) { + List nmTokens = new ArrayList(); + for (Container container : transferredContainers) { + nmTokens.add(rmContext.getNMTokenSecretManager().createAndGetNMToken( + app.getUser(), applicationAttemptId, container)); + } + response.setNMTokensFromPreviousAttempts(nmTokens); + LOG.info("Application " + appID + " retrieved " + + transferredContainers.size() + " containers from previous" + + " attempts and " + nmTokens.size() + " NM tokens."); + } return response; } } diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockRM.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockRM.java index 63efe8f..c91ce35 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockRM.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockRM.java @@ -486,6 +486,7 @@ public static void finishApplicationMaster(RMApp rmApp, MockRM rm, MockNM nm, public static MockAM launchAM(RMApp app, MockRM rm, MockNM nm) throws Exception { + rm.waitForState(app.getApplicationId(), RMAppState.ACCEPTED); RMAppAttempt attempt = app.getCurrentAppAttempt(); nm.nodeHeartbeat(true); MockAM am = rm.sendAMLaunched(attempt.getAppAttemptId()); diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/applicationsmanager/TestAMRestart.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/applicationsmanager/TestAMRestart.java index ca9befd..66d7acb 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/applicationsmanager/TestAMRestart.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/applicationsmanager/TestAMRestart.java @@ -24,6 +24,7 @@ import junit.framework.Assert; +import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse; import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse; import org.apache.hadoop.yarn.api.records.ApplicationAccessType; import org.apache.hadoop.yarn.api.records.ApplicationAttemptId; @@ -31,6 +32,7 @@ import org.apache.hadoop.yarn.api.records.ContainerId; import org.apache.hadoop.yarn.api.records.ContainerState; import org.apache.hadoop.yarn.api.records.ContainerStatus; +import org.apache.hadoop.yarn.api.records.NMToken; import org.apache.hadoop.yarn.api.records.ResourceRequest; import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.server.resourcemanager.MockAM; @@ -160,11 +162,11 @@ public void testAMRestartWithExistingContainers() throws Exception { am2.registerAppAttempt(); // Assert two containers are running: container2 and container3; - Assert.assertEquals(2, registerResponse.getContainersFromPreviousAttempt() + Assert.assertEquals(2, registerResponse.getContainersFromPreviousAttempts() .size()); boolean containerId2Exists = false, containerId3Exists = false; for (Container container : registerResponse - .getContainersFromPreviousAttempt()) { + .getContainersFromPreviousAttempts()) { if (container.getId().equals(containerId2)) { containerId2Exists = true; } @@ -232,4 +234,100 @@ public void testAMRestartWithExistingContainers() throws Exception { rm1.stop(); } + + @Test + public void testNMTokensRebindOnAMRestart() throws Exception { + YarnConfiguration conf = new YarnConfiguration(); + conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 3); + + MockRM rm1 = new MockRM(conf); + rm1.start(); + RMApp app1 = + rm1.submitApp(200, "myname", "myuser", + new HashMap(), false, "default", -1, + null, "MAPREDUCE", false, true); + MockNM nm1 = + new MockNM("127.0.0.1:1234", 8000, rm1.getResourceTrackerService()); + nm1.registerNode(); + MockNM nm2 = + new MockNM("127.1.1.1:4321", 8000, rm1.getResourceTrackerService()); + nm2.registerNode(); + MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1); + + int NUM_CONTAINERS = 1; + List containers = new ArrayList(); + // nmTokens keeps track of all the nmTokens issued in the allocate call. + List expectedNMTokens = new ArrayList(); + + // am1 allocate 1 container on nm1. + while (true) { + AllocateResponse response = + am1.allocate("127.0.0.1", 2000, NUM_CONTAINERS, + new ArrayList()); + nm1.nodeHeartbeat(true); + containers.addAll(response.getAllocatedContainers()); + expectedNMTokens.addAll(response.getNMTokens()); + if (containers.size() == NUM_CONTAINERS) { + break; + } + Thread.sleep(200); + System.out.println("Waiting for container to be allocated."); + } + // launch the container + nm1.nodeHeartbeat(am1.getApplicationAttemptId(), 2, ContainerState.RUNNING); + ContainerId containerId2 = + ContainerId.newInstance(am1.getApplicationAttemptId(), 2); + rm1.waitForState(nm1, containerId2, RMContainerState.RUNNING); + + // fail am1 + nm1.nodeHeartbeat(am1.getApplicationAttemptId(), 1, ContainerState.COMPLETE); + am1.waitForState(RMAppAttemptState.FAILED); + rm1.waitForState(app1.getApplicationId(), RMAppState.ACCEPTED); + + // restart the am + MockAM am2 = MockRM.launchAM(app1, rm1, nm1); + RegisterApplicationMasterResponse registerResponse = + am2.registerAppAttempt(); + rm1.waitForState(app1.getApplicationId(), RMAppState.RUNNING); + + // check am2 get the nm token from am1. + Assert.assertEquals(expectedNMTokens, + registerResponse.getNMTokensFromPreviousAttempts()); + + // am2 allocate 1 container on nm2 + containers = new ArrayList(); + while (true) { + AllocateResponse allocateResponse = + am2.allocate("127.1.1.1", 4000, NUM_CONTAINERS, + new ArrayList()); + nm2.nodeHeartbeat(true); + containers.addAll(allocateResponse.getAllocatedContainers()); + expectedNMTokens.addAll(allocateResponse.getNMTokens()); + if (containers.size() == NUM_CONTAINERS) { + break; + } + Thread.sleep(200); + System.out.println("Waiting for container to be allocated."); + } + nm1.nodeHeartbeat(am2.getApplicationAttemptId(), 2, ContainerState.RUNNING); + ContainerId am2ContainerId2 = + ContainerId.newInstance(am2.getApplicationAttemptId(), 2); + rm1.waitForState(nm1, am2ContainerId2, RMContainerState.RUNNING); + + // fail am2. + nm1.nodeHeartbeat(am2.getApplicationAttemptId(), 1, ContainerState.COMPLETE); + am2.waitForState(RMAppAttemptState.FAILED); + rm1.waitForState(app1.getApplicationId(), RMAppState.ACCEPTED); + + // restart am + MockAM am3 = MockRM.launchAM(app1, rm1, nm1); + registerResponse = am3.registerAppAttempt(); + rm1.waitForState(app1.getApplicationId(), RMAppState.RUNNING); + + // check am3 get the NM token from both am1 and am2; + List transferredTokens = registerResponse.getNMTokensFromPreviousAttempts(); + Assert.assertEquals(2, transferredTokens.size()); + Assert.assertTrue(transferredTokens.containsAll(expectedNMTokens)); + rm1.stop(); + } }