From 9b81cb0537e5b731581e6a375bf0a59abf61c359 Mon Sep 17 00:00:00 2001
From: Sunil G <sunilg@apache.org>
Date: Tue, 23 Jan 2018 17:09:58 +0530
Subject: [PATCH] YARN-7779. Display allocation tags in RM web UI and expose
 same through REST API. Contributed by Weiwei Yang.

---
 .../hadoop/yarn/sls/nodemanager/NodeInfo.java |  6 ++
 .../yarn/sls/scheduler/RMNodeWrapper.java     |  6 ++
 .../server/resourcemanager/rmnode/RMNode.java |  7 ++
 .../resourcemanager/rmnode/RMNodeImpl.java    |  6 ++
 .../constraint/AllocationTagsManager.java     | 11 +++
 .../resourcemanager/webapp/NodesPage.java     |  3 +
 .../webapp/dao/AllocationTagInfo.java         | 56 ++++++++++++++
 .../webapp/dao/AllocationTagsInfo.java        | 59 ++++++++++++++
 .../resourcemanager/webapp/dao/NodeInfo.java  | 15 ++++
 .../server/resourcemanager/MockNodes.java     |  6 ++
 .../resourcemanager/webapp/TestNodesPage.java |  4 +-
 .../webapp/TestRMWebServicesNodes.java        | 77 ++++++++++++++++++-
 12 files changed, 253 insertions(+), 3 deletions(-)
 create mode 100644 hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/AllocationTagInfo.java
 create mode 100644 hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/AllocationTagsInfo.java

diff --git a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/nodemanager/NodeInfo.java b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/nodemanager/NodeInfo.java
index 1016ce1679..0c99139d54 100644
--- a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/nodemanager/NodeInfo.java
+++ b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/nodemanager/NodeInfo.java
@@ -20,6 +20,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import org.apache.hadoop.classification.InterfaceAudience.Private;
@@ -212,6 +213,11 @@ public Integer getDecommissioningTimeout() {
       return null;
     }
 
+    @Override
+    public Map<String, Long> getAllocationTagsWithCount() {
+      return null;
+    }
+
     @Override
     public Resource getPhysicalResource() {
       return null;
diff --git a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/scheduler/RMNodeWrapper.java b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/scheduler/RMNodeWrapper.java
index fdad826994..92f9b0fbeb 100644
--- a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/scheduler/RMNodeWrapper.java
+++ b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/scheduler/RMNodeWrapper.java
@@ -37,6 +37,7 @@
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 @Private
@@ -202,6 +203,11 @@ public Integer getDecommissioningTimeout() {
     return null;
   }
 
+  @Override
+  public Map<String, Long> getAllocationTagsWithCount() {
+    return node.getAllocationTagsWithCount();
+  }
+
   @Override
   public Resource getPhysicalResource() {
     return null;
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNode.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNode.java
index a5615ef06e..872f2a6ef8 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNode.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNode.java
@@ -20,6 +20,7 @@
 
 
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import org.apache.hadoop.net.Node;
@@ -182,4 +183,10 @@ public interface RMNode {
    * @return the decommissioning timeout in second.
    */
   Integer getDecommissioningTimeout();
+
+  /**
+   * Get the allocation tags and their counts associated with this node.
+   * @return a map of each allocation tag and its count.
+   */
+  Map<String, Long> getAllocationTagsWithCount();
 }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNodeImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNodeImpl.java
index da54eb93a2..4fc2d8a4a8 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNodeImpl.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNodeImpl.java
@@ -1529,4 +1529,10 @@ public void setUntrackedTimeStamp(long ts) {
   public Integer getDecommissioningTimeout() {
     return decommissioningTimeout;
   }
+
+  @Override
+  public Map<String, Long> getAllocationTagsWithCount() {
+    return context.getAllocationTagsManager()
+        .getAllocationTagsWithCount(getNodeID());
+  }
 }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/AllocationTagsManager.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/AllocationTagsManager.java
index 7ad5e8c372..42a78c98ed 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/AllocationTagsManager.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/AllocationTagsManager.java
@@ -548,4 +548,15 @@ public long getRackCardinalityByOp(String rack, ApplicationId applicationId,
       readLock.unlock();
     }
   }
+
+  /**
+   * Returns a map whose key is the allocation tag and value is the
+   * count of allocations with this tag.
+   *
+   * @param nodeId
+   * @return allocation tag to count mapping
+   */
+  public Map<String, Long> getAllocationTagsWithCount(NodeId nodeId) {
+    return globalNodeMapping.getTypeToTagsWithCount().get(nodeId);
+  }
 }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodesPage.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodesPage.java
index d0e384da34..3e78cf4f40 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodesPage.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodesPage.java
@@ -81,12 +81,14 @@ protected void render(Block html) {
 
       if (!this.opportunisticContainersEnabled) {
         trbody.th(".containers", "Containers")
+            .th(".allocationTags", "Allocation Tags")
             .th(".mem", "Mem Used")
             .th(".mem", "Mem Avail")
             .th(".vcores", "VCores Used")
             .th(".vcores", "VCores Avail");
       } else {
         trbody.th(".containers", "Running Containers (G)")
+            .th(".allocationTags", "Allocation Tags")
             .th(".mem", "Mem Used (G)")
             .th(".mem", "Mem Avail (G)")
             .th(".vcores", "VCores Used (G)")
@@ -167,6 +169,7 @@ protected void render(Block html) {
             .append(Times.format(info.getLastHealthUpdate())).append("\",\"")
             .append(info.getHealthReport()).append("\",\"")
             .append(String.valueOf(info.getNumContainers())).append("\",\"")
+            .append(info.getAllocationTagsSummary()).append("\",\"")
             .append("<br title='").append(String.valueOf(usedMemory))
             .append("'>").append(StringUtils.byteDesc(usedMemory * BYTES_IN_MB))
             .append("\",\"").append("<br title='")
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/AllocationTagInfo.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/AllocationTagInfo.java
new file mode 100644
index 0000000000..97f9e90022
--- /dev/null
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/AllocationTagInfo.java
@@ -0,0 +1,56 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.yarn.server.resourcemanager.webapp.dao;
+
+import javax.xml.bind.annotation.XmlAccessType;
+import javax.xml.bind.annotation.XmlAccessorType;
+import javax.xml.bind.annotation.XmlRootElement;
+
+/**
+ * DAO object to display node allocation tag.
+ */
+@XmlRootElement(name = "allocationTagInfo")
+@XmlAccessorType(XmlAccessType.FIELD)
+public class AllocationTagInfo {
+
+  private String allocationTag;
+  private long allocationsCount;
+
+  public AllocationTagInfo() {
+    // JAXB needs this
+  }
+
+  public AllocationTagInfo(String tag, long count) {
+    this.allocationTag = tag;
+    this.allocationsCount = count;
+  }
+
+  public String getAllocationTag() {
+    return this.allocationTag;
+  }
+
+  public long getAllocationsCount() {
+    return this.allocationsCount;
+  }
+
+  @Override
+  public String toString() {
+    return allocationTag + "(" + allocationsCount + ")";
+  }
+}
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/AllocationTagsInfo.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/AllocationTagsInfo.java
new file mode 100644
index 0000000000..ee09aa2f03
--- /dev/null
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/AllocationTagsInfo.java
@@ -0,0 +1,59 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.yarn.server.resourcemanager.webapp.dao;
+
+import org.apache.hadoop.util.StringUtils;
+
+import javax.xml.bind.annotation.XmlAccessType;
+import javax.xml.bind.annotation.XmlAccessorType;
+import javax.xml.bind.annotation.XmlRootElement;
+import java.util.ArrayList;
+import java.util.Iterator;
+
+/**
+ * DAO object to display node allocation tags.
+ */
+@XmlRootElement(name = "allocationTagsInfo")
+@XmlAccessorType(XmlAccessType.FIELD)
+public class AllocationTagsInfo {
+
+  private ArrayList<AllocationTagInfo> allocationTagInfo;
+
+  public AllocationTagsInfo() {
+    allocationTagInfo = new ArrayList<>();
+  }
+
+  public void addAllocationTag(AllocationTagInfo info) {
+    allocationTagInfo.add(info);
+  }
+
+  @Override
+  public String toString() {
+    StringBuffer sb = new StringBuffer();
+    Iterator<AllocationTagInfo> it = allocationTagInfo.iterator();
+    while (it.hasNext()) {
+      AllocationTagInfo current = it.next();
+      sb.append(current.toString());
+      if (it.hasNext()) {
+        sb.append(StringUtils.COMMA);
+      }
+    }
+    return sb.toString();
+  }
+}
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/NodeInfo.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/NodeInfo.java
index 3cec215f0e..46a6e60ce3 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/NodeInfo.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/NodeInfo.java
@@ -20,6 +20,7 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Map;
 import java.util.Set;
 
 import javax.xml.bind.annotation.XmlAccessType;
@@ -57,6 +58,7 @@ public class NodeInfo {
   private long usedVirtualCoresOpport;
   private int numQueuedContainers;
   protected ArrayList<String> nodeLabels = new ArrayList<String>();
+  private AllocationTagsInfo allocationTags;
   protected ResourceUtilizationInfo resourceUtilization;
   protected ResourceInfo usedResource;
   protected ResourceInfo availableResource;
@@ -111,6 +113,14 @@ public NodeInfo(RMNode ni, ResourceScheduler sched) {
       Collections.sort(nodeLabels);
     }
 
+    // add allocation tags
+    allocationTags = new AllocationTagsInfo();
+    Map<String, Long> allocationTagsInfo = ni.getAllocationTagsWithCount();
+    if (allocationTagsInfo != null) {
+      allocationTagsInfo.forEach((tag, count) ->
+          allocationTags.addAllocationTag(new AllocationTagInfo(tag, count)));
+    }
+
     // update node and containers resource utilization
     this.resourceUtilization = new ResourceUtilizationInfo(ni);
   }
@@ -207,6 +217,11 @@ public ResourceUtilizationInfo getResourceUtilization() {
     return this.resourceUtilization;
   }
 
+  public String getAllocationTagsSummary() {
+    return this.allocationTags == null ? "" :
+        this.allocationTags.toString();
+  }
+
   @VisibleForTesting
   public void setId(String id) {
     this.id = id;
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockNodes.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockNodes.java
index d6549b99c0..84105d9303 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockNodes.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockNodes.java
@@ -22,6 +22,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
+import java.util.Map;
 
 import org.apache.hadoop.net.Node;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
@@ -279,6 +280,11 @@ public Integer getDecommissioningTimeout() {
       return null;
     }
 
+    @Override
+    public Map<String, Long> getAllocationTagsWithCount() {
+      return null;
+    }
+
     @Override
     public Resource getPhysicalResource() {
       return this.physicalResource;
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestNodesPage.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestNodesPage.java
index cc97674696..26e8c2ab66 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestNodesPage.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestNodesPage.java
@@ -48,8 +48,8 @@ public class TestNodesPage {
 
   // Number of Actual Table Headers for NodesPage.NodesBlock might change in
   // future. In that case this value should be adjusted to the new value.
-  final int numberOfThInMetricsTable = 23;
-  final int numberOfActualTableHeaders = 13;
+  private final int numberOfThInMetricsTable = 23;
+  private final int numberOfActualTableHeaders = 14;
   private final int numberOfThForOpportunisticContainers = 4;
 
   private Injector injector;
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestRMWebServicesNodes.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestRMWebServicesNodes.java
index fb597fc7b7..7ea7e811f3 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestRMWebServicesNodes.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestRMWebServicesNodes.java
@@ -22,10 +22,15 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import java.io.StringReader;
 import java.util.ArrayList;
 import java.util.EnumSet;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.Iterator;
 
 import javax.ws.rs.core.MediaType;
 import javax.xml.parsers.DocumentBuilder;
@@ -51,6 +56,7 @@
 import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNodeStatusEvent;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerNodeReport;
+import org.apache.hadoop.yarn.server.resourcemanager.scheduler.constraint.AllocationTagsManager;
 import org.apache.hadoop.yarn.util.RackResolver;
 import org.apache.hadoop.yarn.util.YarnVersionInfo;
 import org.apache.hadoop.yarn.webapp.GenericExceptionHandler;
@@ -734,7 +740,7 @@ public void verifyNodesXML(NodeList nodes, RMNode nm)
 
   public void verifyNodeInfo(JSONObject nodeInfo, RMNode nm)
       throws JSONException, Exception {
-    assertEquals("incorrect number of elements", 18, nodeInfo.length());
+    assertEquals("incorrect number of elements", 19, nodeInfo.length());
 
     JSONObject resourceInfo = nodeInfo.getJSONObject("resourceUtilization");
     verifyNodeInfoGeneric(nm, nodeInfo.getString("state"),
@@ -837,4 +843,73 @@ public void verifyNodeInfoGeneric(RMNode node, String state, String rack,
     }
   }
 
+  @Test
+  public void testNodesAllocationTags() throws Exception {
+    NodeId nm1 = NodeId.newInstance("host1", 1234);
+    NodeId nm2 = NodeId.newInstance("host2", 2345);
+    AllocationTagsManager atm = mock(AllocationTagsManager.class);
+
+    Map<String, Map<String, Long>> expectedAllocationTags = new TreeMap<>();
+    Map<String, Long> nm1Tags = new TreeMap<>();
+    nm1Tags.put("A", 1L);
+    nm1Tags.put("B", 2L);
+    Map<String, Long> nm2Tags = new TreeMap<>();
+    nm2Tags.put("C", 1L);
+    nm2Tags.put("D", 2L);
+    expectedAllocationTags.put(nm1.toString(), nm1Tags);
+    expectedAllocationTags.put(nm2.toString(), nm2Tags);
+
+    when(atm.getAllocationTagsWithCount(nm1)).thenReturn(nm1Tags);
+    when(atm.getAllocationTagsWithCount(nm2)).thenReturn(nm2Tags);
+    rm.getRMContext().setAllocationTagsManager(atm);
+
+    rm.start();
+
+    rm.registerNode(nm1.toString(), 1024);
+    rm.registerNode(nm2.toString(), 1024);
+
+    WebResource r = resource();
+    ClientResponse response = r.path("ws").path("v1").path("cluster")
+        .path("nodes").accept("application/json").get(ClientResponse.class);
+    assertEquals(MediaType.APPLICATION_JSON + "; " + JettyUtils.UTF_8,
+        response.getType().toString());
+    JSONObject nodesInfoJson = response.getEntity(JSONObject.class);
+    verifyNodeAllocationTag(nodesInfoJson, expectedAllocationTags);
+
+    rm.stop();
+  }
+
+  private void verifyNodeAllocationTag(JSONObject json,
+      Map<String, Map<String, Long>> expectedAllocationTags)
+      throws JSONException {
+    JSONArray nodes = json.getJSONObject("nodes").getJSONArray("node");
+    assertEquals(expectedAllocationTags.size(), nodes.length());
+    for (int i=0; i<nodes.length(); i++) {
+      JSONObject nodeJson = nodes.getJSONObject(i);
+      String nodeId = nodeJson.getString("id");
+
+      // Ensure the response contains all nodes info
+      assertTrue("Nodes info should have expected node IDs",
+          expectedAllocationTags.containsKey(nodeId));
+
+      Map<String, Long> expectedTags = expectedAllocationTags.get(nodeId);
+      JSONArray tagsInfo = nodeJson.getJSONObject("allocationTags")
+          .getJSONArray("allocationTagInfo");
+
+      // Ensure number of tags are expected.
+      assertEquals(expectedTags.size(), tagsInfo.length());
+
+      // Iterate expected tags and make sure the actual
+      // tags/counts are matched.
+      Iterator<String> it = expectedTags.keySet().iterator();
+      for (int j=0; j<tagsInfo.length(); j++) {
+        JSONObject tagInfo = tagsInfo.getJSONObject(j);
+        String expectedTag = it.next();
+        assertEquals(tagInfo.getString("allocationTag"), expectedTag);
+        assertEquals(tagInfo.getLong("allocationsCount"),
+            expectedTags.get(expectedTag).longValue());
+      }
+    }
+  }
+
 }