From 0de10680b7e5a9dfc85173bcfd338fd3656aa57f Mon Sep 17 00:00:00 2001
From: Daniel Templeton <templedf@apache.org>
Date: Wed, 8 Nov 2017 16:43:49 -0800
Subject: [PATCH] YARN-7166. Container REST endpoints should report resource
 types

Change-Id: If9c2fe58d4cf758bb6b6cf363dc01f35f8720987
---
 .../yarn/server/webapp/dao/ContainerInfo.java | 39 ++++++++++++++++---
 1 file changed, 34 insertions(+), 5 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/webapp/dao/ContainerInfo.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/webapp/dao/ContainerInfo.java
index 1a5ee85cf8..26a822c13d 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/webapp/dao/ContainerInfo.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/webapp/dao/ContainerInfo.java
@@ -18,6 +18,9 @@
 
 package org.apache.hadoop.yarn.server.webapp.dao;
 
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import javax.xml.bind.annotation.XmlAccessType;
 import javax.xml.bind.annotation.XmlAccessorType;
 import javax.xml.bind.annotation.XmlRootElement;
@@ -27,6 +30,8 @@
 
 import org.apache.hadoop.yarn.api.records.ContainerReport;
 import org.apache.hadoop.yarn.api.records.ContainerState;
+import org.apache.hadoop.yarn.api.records.Resource;
+import org.apache.hadoop.yarn.api.records.ResourceInformation;
 import org.apache.hadoop.yarn.util.Times;
 
 @Public
@@ -49,20 +54,18 @@ public class ContainerInfo {
   protected ContainerState containerState;
   protected String nodeHttpAddress;
   protected String nodeId;
+  protected Map<String, Long> allocatedResources;
 
   public ContainerInfo() {
     // JAXB needs this
   }
 
   public ContainerInfo(ContainerReport container) {
-    containerId = container.getContainerId().toString();
-    if (container.getAllocatedResource() != null) {
-      allocatedMB = container.getAllocatedResource().getMemorySize();
-      allocatedVCores = container.getAllocatedResource().getVirtualCores();
-    }
     if (container.getAssignedNode() != null) {
       assignedNodeId = container.getAssignedNode().toString();
     }
+
+    containerId = container.getContainerId().toString();
     priority = container.getPriority().getPriority();
     startedTime = container.getCreationTime();
     finishedTime = container.getFinishTime();
@@ -73,6 +76,22 @@ public ContainerInfo(ContainerReport container) {
     containerState = container.getContainerState();
     nodeHttpAddress = container.getNodeHttpAddress();
     nodeId = container.getAssignedNode().toString();
+
+    Resource allocated = container.getAllocatedResource();
+
+    if (allocated != null) {
+      allocatedMB = allocated.getMemorySize();
+      allocatedVCores = allocated.getVirtualCores();
+
+      // Now populate the allocated resources. This map will include memory
+      // and CPU, because it's where they belong. We still keep allocatedMB
+      // and allocatedVCores so that we don't break the API.
+      allocatedResources = new HashMap<>();
+
+      for (ResourceInformation info : allocated.getResources()) {
+        allocatedResources.put(info.getName(), info.getValue());
+      }
+    }
   }
 
   public String getContainerId() {
@@ -130,4 +149,14 @@ public String getNodeHttpAddress() {
   public String getNodeId() {
     return nodeId;
   }
+
+  /**
+   * Return a map of the allocated resources. The map key is the resource name,
+   * and the value is the resource value.
+   *
+   * @return the allocated resources map
+   */
+  public Map<String, Long> getAllocatedResources() {
+    return Collections.unmodifiableMap(allocatedResources);
+  }
 }