YARN-1321. Changed NMTokenCache to support both singleton and an instance usage. Contributed by Alejandro Abdelnur.

git-svn-id: https://svn.apache.org/repos/asf/hadoop/common/trunk@1537334 13f79535-47bb-0310-9956-ffa450edef68

hadoop-yarn-project/CHANGES.txt

@@ -123,6 +123,9 @@ Release 2.2.1 - UNRELEASED
     YARN-1109. Demote NodeManager "Sending out status for container" logs to
     debug (haosdent via Sandy Ryza)
 
+    YARN-1321. Changed NMTokenCache to support both singleton and an instance
+    usage. (Alejandro Abdelnur via vinodkv) 
+
   OPTIMIZATIONS
 
   BUG FIXES

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

@@ -59,9 +59,12 @@ public abstract class AMRMClient<T extends AMRMClient.ContainerRequest> extends
     return client;
   }
 
+  private NMTokenCache nmTokenCache;
+
   @Private
   protected AMRMClient(String name) {
     super(name);
+    nmTokenCache = NMTokenCache.getSingleton();
   }
 
   /**
@@ -297,4 +300,33 @@ public abstract class AMRMClient<T extends AMRMClient.ContainerRequest> extends
    */
   public abstract void updateBlacklist(List<String> blacklistAdditions,
       List<String> blacklistRemovals);
+
+  /**
+   * Set the NM token cache for the <code>AMRMClient</code>. This cache must
+   * be shared with the {@link NMClient} used to manage containers for the
+   * <code>AMRMClient</code>
+   * <p/>
+   * If a NM token cache is not set, the {@link NMTokenCache#getSingleton()}
+   * singleton instance will be used.
+   *
+   * @param nmTokenCache the NM token cache to use.
+   */
+  public void setNMTokenCache(NMTokenCache nmTokenCache) {
+    this.nmTokenCache = nmTokenCache;
+  }
+
+  /**
+   * Get the NM token cache of the <code>AMRMClient</code>. This cache must be
+   * shared with the {@link NMClient} used to manage containers for the
+   * <code>AMRMClient</code>.
+   * <p/>
+   * If a NM token cache is not set, the {@link NMTokenCache#getSingleton()}
+   * singleton instance will be used.
+   *
+   * @return the NM token cache.
+   */
+  public NMTokenCache getNMTokenCache() {
+    return nmTokenCache;
+  }
+
 }

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

@@ -58,6 +58,8 @@ public abstract class NMClient extends AbstractService {
     return client;
   }
 
+  private NMTokenCache nmTokenCache = NMTokenCache.getSingleton();
+
   @Private
   protected NMClient(String name) {
     super(name);
@@ -118,4 +120,33 @@ public abstract class NMClient extends AbstractService {
    * @param enabled whether the feature is enabled or not
    */
   public abstract void cleanupRunningContainersOnStop(boolean enabled);
+
+  /**
+   * Set the NM Token cache of the <code>NMClient</code>. This cache must be
+   * shared with the {@link AMRMClient} that requested the containers managed
+   * by this <code>NMClient</code>
+   * <p/>
+   * If a NM token cache is not set, the {@link NMTokenCache#getSingleton()}
+   * singleton instance will be used.
+   *
+   * @param nmTokenCache the NM token cache to use.
+   */
+  public void setNMTokenCache(NMTokenCache nmTokenCache) {
+    this.nmTokenCache = nmTokenCache;
+  }
+
+  /**
+   * Get the NM token cache of the <code>NMClient</code>. This cache must be
+   * shared with the {@link AMRMClient} that requested the containers managed
+   * by this <code>NMClient</code>
+   * <p/>
+   * If a NM token cache is not set, the {@link NMTokenCache#getSingleton()}
+   * singleton instance will be used.
+   *
+   * @return the NM token cache
+   */
+  public NMTokenCache getNMTokenCache() {
+    return nmTokenCache;
+  }
+
 }

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

@@ -23,21 +23,139 @@ import java.util.concurrent.ConcurrentHashMap;
 import org.apache.hadoop.classification.InterfaceAudience.Private;
 import org.apache.hadoop.classification.InterfaceAudience.Public;
 import org.apache.hadoop.classification.InterfaceStability.Evolving;
+import org.apache.hadoop.yarn.api.ApplicationMasterProtocol;
+import org.apache.hadoop.yarn.api.ContainerManagementProtocol;
 import org.apache.hadoop.yarn.api.records.Token;
+import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync;
+import org.apache.hadoop.yarn.client.api.async.NMClientAsync;
 
 import com.google.common.annotations.VisibleForTesting;
 
 /**
- * It manages NMTokens required for communicating with Node manager. Its a
- * static token cache.
+ * NMTokenCache manages NMTokens required for an Application Master
+ * communicating with individual NodeManagers.
+ * <p/>
+ * By default Yarn client libraries {@link AMRMClient} and {@link NMClient} use
+ * {@link #getSingleton()} instance of the cache.
+ * <ul>
+ * <li>Using the singleton instance of the cache is appropriate when running a
+ * single ApplicationMaster in the same JVM.</li>
+ * <li>When using the singleton, users don't need to do anything special,
+ * {@link AMRMClient} and {@link NMClient} are already set up to use the default
+ * singleton {@link NMTokenCache}</li>
+ * </ul>
+ * <p/>
+ * If running multiple Application Masters in the same JVM, a different cache
+ * instance should be used for each Application Master.
+ * <p/>
+ * <ul>
+ * <li>
+ * If using the {@link AMRMClient} and the {@link NMClient}, setting up and using
+ * an instance cache is as follows:
+ * <p/>
+ * 
+ * <pre>
+ *   NMTokenCache nmTokenCache = new NMTokenCache();
+ *   AMRMClient rmClient = AMRMClient.createAMRMClient();
+ *   NMClient nmClient = NMClient.createNMClient();
+ *   nmClient.setNMTokenCache(nmTokenCache);
+ *   ...
+ * </pre>
+ * </li>
+ * <li>
+ * If using the {@link AMRMClientAsync} and the {@link NMClientAsync}, setting up
+ * and using an instance cache is as follows:
+ * <p/>
+ * 
+ * <pre>
+ *   NMTokenCache nmTokenCache = new NMTokenCache();
+ *   AMRMClient rmClient = AMRMClient.createAMRMClient();
+ *   NMClient nmClient = NMClient.createNMClient();
+ *   nmClient.setNMTokenCache(nmTokenCache);
+ *   AMRMClientAsync rmClientAsync = new AMRMClientAsync(rmClient, 1000, [AMRM_CALLBACK]);
+ *   NMClientAsync nmClientAsync = new NMClientAsync("nmClient", nmClient, [NM_CALLBACK]);
+ *   ...
+ * </pre>
+ * </li>
+ * <li>
+ * If using {@link ApplicationMasterProtocol} and
+ * {@link ContainerManagementProtocol} directly, setting up and using an
+ * instance cache is as follows:
+ * <p/>
+ * 
+ * <pre>
+ *   NMTokenCache nmTokenCache = new NMTokenCache();
+ *   ...
+ *   ApplicationMasterProtocol amPro = ClientRMProxy.createRMProxy(conf, ApplicationMasterProtocol.class);
+ *   ...
+ *   AllocateRequest allocateRequest = ...
+ *   ...
+ *   AllocateResponse allocateResponse = rmClient.allocate(allocateRequest);
+ *   for (NMToken token : allocateResponse.getNMTokens()) {
+ *     nmTokenCache.setToken(token.getNodeId().toString(), token.getToken());
+ *   }
+ *   ...
+ *   ContainerManagementProtocolProxy nmPro = ContainerManagementProtocolProxy(conf, nmTokenCache);
+ *   ...
+ *   nmPro.startContainer(container, containerContext);
+ *   ...
+ * </pre>
+ * </li>
+ * </ul>
+ * It is also possible to mix the usage of a client (<code>AMRMClient</code> or
+ * <code>NMClient</code>, or the async versions of them) with a protocol proxy (
+ * <code>ContainerManagementProtocolProxy</code> or
+ * <code>ApplicationMasterProtocol</code>).
  */
 @Public
 @Evolving
 public class NMTokenCache {
-  private static ConcurrentHashMap<String, Token> nmTokens;
+  private static final NMTokenCache NM_TOKEN_CACHE = new NMTokenCache();
   
+  /**
+   * Returns the singleton NM token cache.
+   *
+   * @return the singleton NM token cache.
+   */
+  public static NMTokenCache getSingleton() {
+    return NM_TOKEN_CACHE;
+  }
   
-  static {
+  /**
+   * Returns NMToken, null if absent. Only the singleton obtained from
+   * {@link #getSingleton()} is looked at for the tokens. If you are using your
+   * own NMTokenCache that is different from the singleton, use
+   * {@link #getToken(String) }
+   * 
+   * @param nodeAddr
+   * @return {@link Token} NMToken required for communicating with node manager
+   */
+  @Public
+  public static Token getNMToken(String nodeAddr) {
+    return NM_TOKEN_CACHE.getToken(nodeAddr);
+  }
+  
+  /**
+   * Sets the NMToken for node address only in the singleton obtained from
+   * {@link #getSingleton()}. If you are using your own NMTokenCache that is
+   * different from the singleton, use {@link #setToken(String, Token) }
+   * 
+   * @param nodeAddr
+   *          node address (host:port)
+   * @param token
+   *          NMToken
+   */
+  @Public
+  public static void setNMToken(String nodeAddr, Token token) {
+    NM_TOKEN_CACHE.setToken(nodeAddr, token);
+  }
+
+  private ConcurrentHashMap<String, Token> nmTokens;
+
+  /**
+   * Creates a NM token cache instance.
+   */
+  public NMTokenCache() {
     nmTokens = new ConcurrentHashMap<String, Token>();
   }
   
@@ -45,11 +163,11 @@ public class NMTokenCache {
    * Returns NMToken, null if absent
    * @param nodeAddr
    * @return {@link Token} NMToken required for communicating with node
-   * manager
+   *         manager
    */
   @Public
   @Evolving
-  public static Token getNMToken(String nodeAddr) {
+  public Token getToken(String nodeAddr) {
     return nmTokens.get(nodeAddr);
   }
   
@@ -60,7 +178,7 @@ public class NMTokenCache {
    */
   @Public
   @Evolving
-  public static void setNMToken(String nodeAddr, Token token) {
+  public void setToken(String nodeAddr, Token token) {
     nmTokens.put(nodeAddr, token);
   }
   
@@ -69,7 +187,7 @@ public class NMTokenCache {
    */
   @Private
   @VisibleForTesting
-  public static boolean containsNMToken(String nodeAddr) {
+  public boolean containsToken(String nodeAddr) {
     return nmTokens.containsKey(nodeAddr);
   }
   
@@ -78,7 +196,7 @@ public class NMTokenCache {
    */
   @Private
   @VisibleForTesting
-  public static int numberOfNMTokensInCache() {
+  public int numberOfTokensInCache() {
     return nmTokens.size();
   }
   
@@ -88,7 +206,7 @@ public class NMTokenCache {
    */
   @Private
   @VisibleForTesting
-  public static void removeNMToken(String nodeAddr) {
+  public void removeToken(String nodeAddr) {
     nmTokens.remove(nodeAddr);
   }
   
@@ -97,7 +215,7 @@ public class NMTokenCache {
    */
   @Private
   @VisibleForTesting
-  public static void clearCache() {
+  public void clearCache() {
     nmTokens.clear();
   }
 }

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

@@ -58,7 +58,6 @@ import org.apache.hadoop.yarn.client.ClientRMProxy;
 import org.apache.hadoop.yarn.client.api.AMRMClient;
 import org.apache.hadoop.yarn.client.api.AMRMClient.ContainerRequest;
 import org.apache.hadoop.yarn.client.api.InvalidContainerRequestException;
-import org.apache.hadoop.yarn.client.api.NMTokenCache;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.exceptions.YarnRuntimeException;
@@ -288,12 +287,12 @@ public class AMRMClientImpl<T extends ContainerRequest> extends AMRMClient<T> {
   protected void populateNMTokens(AllocateResponse allocateResponse) {
     for (NMToken token : allocateResponse.getNMTokens()) {
       String nodeId = token.getNodeId().toString();
-      if (NMTokenCache.containsNMToken(nodeId)) {
+      if (getNMTokenCache().containsToken(nodeId)) {
         LOG.debug("Replacing token for : " + nodeId);
       } else {
         LOG.debug("Received new token for : " + nodeId);
       }
-      NMTokenCache.setNMToken(nodeId, token.getToken());
+      getNMTokenCache().setToken(nodeId, token.getToken());
     }
   }
 

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

@@ -56,9 +56,16 @@ public class ContainerManagementProtocolProxy {
   private final LinkedHashMap<String, ContainerManagementProtocolProxyData> cmProxy;
   private final Configuration conf;
   private final YarnRPC rpc;
+  private NMTokenCache nmTokenCache;
   
   public ContainerManagementProtocolProxy(Configuration conf) {
+    this(conf, NMTokenCache.getSingleton());
+  }
+
+  public ContainerManagementProtocolProxy(Configuration conf,
+      NMTokenCache nmTokenCache) {
     this.conf = conf;
+    this.nmTokenCache = nmTokenCache;
 
     maxConnectedNMs =
         conf.getInt(YarnConfiguration.NM_CLIENT_MAX_NM_PROXIES,
@@ -86,7 +93,7 @@ public class ContainerManagementProtocolProxy {
 
     while (proxy != null
         && !proxy.token.getIdentifier().equals(
-            NMTokenCache.getNMToken(containerManagerBindAddr).getIdentifier())) {
+            nmTokenCache.getToken(containerManagerBindAddr).getIdentifier())) {
       LOG.info("Refreshing proxy as NMToken got updated for node : "
           + containerManagerBindAddr);
       // Token is updated. check if anyone has already tried closing it.
@@ -109,7 +116,7 @@ public class ContainerManagementProtocolProxy {
     if (proxy == null) {
       proxy =
           new ContainerManagementProtocolProxyData(rpc, containerManagerBindAddr,
-              containerId, NMTokenCache.getNMToken(containerManagerBindAddr));
+              containerId, nmTokenCache.getToken(containerManagerBindAddr));
       if (cmProxy.size() > maxConnectedNMs) {
         // Number of existing proxy exceed the limit.
         String cmAddr = cmProxy.keySet().iterator().next();

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

@@ -130,7 +130,10 @@ public class NMClientImpl extends NMClient {
   @Override
   protected void serviceInit(Configuration conf) throws Exception {
     super.serviceInit(conf);
-    cmProxy = new ContainerManagementProtocolProxy(conf);
+    if (getNMTokenCache() == null) {
+      throw new IllegalStateException("NMTokenCache has not been set");
+    }
+    cmProxy = new ContainerManagementProtocolProxy(conf, getNMTokenCache());
   }
   
   @Override

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

@@ -626,6 +626,13 @@ public class TestAMRMClient {
     try {
       // start am rm client
       amClient = AMRMClient.<ContainerRequest>createAMRMClient();
+
+      //setting an instance NMTokenCache
+      amClient.setNMTokenCache(new NMTokenCache());
+      //asserting we are not using the singleton instance cache
+      Assert.assertNotSame(NMTokenCache.getSingleton(), 
+          amClient.getNMTokenCache());
+
       amClient.init(conf);
       amClient.start();
 
@@ -681,8 +688,8 @@ public class TestAMRMClient {
     int iterationsLeft = 3;
     Set<ContainerId> releases = new TreeSet<ContainerId>();
     
-    NMTokenCache.clearCache();
-    Assert.assertEquals(0, NMTokenCache.numberOfNMTokensInCache());
+    amClient.getNMTokenCache().clearCache();
+    Assert.assertEquals(0, amClient.getNMTokenCache().numberOfTokensInCache());
     HashMap<String, Token> receivedNMTokens = new HashMap<String, Token>();
     
     while (allocatedContainerCount < containersRequestedAny

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

@@ -78,6 +78,7 @@ public class TestNMClient {
   List<NodeReport> nodeReports = null;
   ApplicationAttemptId attemptId = null;
   int nodeCount = 3;
+  NMTokenCache nmTokenCache = null;
   
   @Before
   public void setup() throws YarnException, IOException {
@@ -155,10 +156,16 @@ public class TestNMClient {
       .createRemoteUser(UserGroupInformation.getCurrentUser().getUserName()));
     UserGroupInformation.getCurrentUser().addToken(appAttempt.getAMRMToken());
 
+    //creating an instance NMTokenCase
+    nmTokenCache = new NMTokenCache();
+    
     // start am rm client
     rmClient =
         (AMRMClientImpl<ContainerRequest>) AMRMClient
           .<ContainerRequest> createAMRMClient();
+
+    //setting an instance NMTokenCase
+    rmClient.setNMTokenCache(nmTokenCache);
     rmClient.init(conf);
     rmClient.start();
     assertNotNull(rmClient);
@@ -166,6 +173,9 @@ public class TestNMClient {
 
     // start am nm client
     nmClient = (NMClientImpl) NMClient.createNMClient();
+    
+    //propagating the AMRMClient NMTokenCache instance
+    nmClient.setNMTokenCache(rmClient.getNMTokenCache());
     nmClient.init(conf);
     nmClient.start();
     assertNotNull(nmClient);
@@ -258,7 +268,7 @@ public class TestNMClient {
       }
       if (!allocResponse.getNMTokens().isEmpty()) {
         for (NMToken token : allocResponse.getNMTokens()) {
-          NMTokenCache.setNMToken(token.getNodeId().toString(),
+          rmClient.getNMTokenCache().setToken(token.getNodeId().toString(),
               token.getToken());
         }
       }