YARN-3103. AMRMClientImpl does not update AMRM token properly. Contributed by Jason Lowe

(cherry picked from commit 6d2bdbd7dab179dfb4f19bb41809e97f1db88c6b)

(cherry picked from commit 12522fd9cbd8da8c040a5b7bb71fcdaa256daf89)
(cherry picked from commit f50f5ad49d3b70448647384fc5f020214cb58f10)

hadoop-yarn-project/CHANGES.txt

@@ -69,6 +69,9 @@ Release 2.6.1 - UNRELEASED
     YARN-3011. Possible IllegalArgumentException in ResourceLocalizationService
     might lead NM to crash. (Varun Saxena via jianhe)
 
+    YARN-3103. AMRMClientImpl does not update AMRM token properly. (Jason Lowe
+    via jianhe)
+
 Release 2.6.0 - 2014-11-18
 
   INCOMPATIBLE CHANGES

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/AMRMClientImpl.java

@@ -756,11 +756,11 @@ public class AMRMClientImpl<T extends ContainerRequest> extends AMRMClient<T> {
         new org.apache.hadoop.security.token.Token<AMRMTokenIdentifier>(token
           .getIdentifier().array(), token.getPassword().array(), new Text(
           token.getKind()), new Text(token.getService()));
-    amrmToken.setService(ClientRMProxy.getAMRMTokenService(getConfig()));
+    // Preserve the token service sent by the RM when adding the token
+    // to ensure we replace the previous token setup by the RM.
+    // Afterwards we can update the service address for the RPC layer.
     UserGroupInformation currentUGI = UserGroupInformation.getCurrentUser();
-    if (UserGroupInformation.isSecurityEnabled()) {
-      currentUGI = UserGroupInformation.getLoginUser();
-    }
     currentUGI.addToken(amrmToken);
+    amrmToken.setService(ClientRMProxy.getAMRMTokenService(getConfig()));
   }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestAMRMClient.java

@@ -200,8 +200,11 @@ public class TestAMRMClient {
     // of testing.
     UserGroupInformation.setLoginUser(UserGroupInformation
       .createRemoteUser(UserGroupInformation.getCurrentUser().getUserName()));
-    appAttempt.getAMRMToken().setService(ClientRMProxy.getAMRMTokenService(conf));
+
+    // emulate RM setup of AMRM token in credentials by adding the token
+    // *before* setting the token service
     UserGroupInformation.getCurrentUser().addToken(appAttempt.getAMRMToken());
+    appAttempt.getAMRMToken().setService(ClientRMProxy.getAMRMTokenService(conf));
   }
   
   @After
@@ -1026,13 +1029,18 @@ public class TestAMRMClient {
         UserGroupInformation.getCurrentUser().getCredentials();
     Iterator<org.apache.hadoop.security.token.Token<?>> iter =
         credentials.getAllTokens().iterator();
+    org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> result = null;
     while (iter.hasNext()) {
       org.apache.hadoop.security.token.Token<?> token = iter.next();
       if (token.getKind().equals(AMRMTokenIdentifier.KIND_NAME)) {
-        return (org.apache.hadoop.security.token.Token<AMRMTokenIdentifier>)
+        if (result != null) {
+          Assert.fail("credentials has more than one AMRM token."
+              + " token1: " + result + " token2: " + token);
+        }
+        result = (org.apache.hadoop.security.token.Token<AMRMTokenIdentifier>)
             token;
       }
     }
-    return null;
+    return result;
   }
 }