From aa602381c595db4b958709c03874ad54597ba197 Mon Sep 17 00:00:00 2001
From: slfan1989 <55643692+slfan1989@users.noreply.github.com>
Date: Tue, 28 Mar 2023 22:48:46 +0800
Subject: [PATCH] YARN-11426. Improve YARN NodeLabel Memory Display. (#5335)

YARN-11426. Improve YARN NodeLabel Memory Display.
Co-authored-by: slfan1989 <louj1988@@>
Reviewed-by: Inigo Goiri <inigoiri@apache.org>
Reviewed-by: Chris Nauroth <cnauroth@apache.org>
Signed-off-by: Shilun Fan <slfan1989@apache.org>
---
 .../hadoop-yarn/hadoop-yarn-api/pom.xml       |  5 ++++
 .../hadoop/yarn/api/records/Resource.java     | 11 ++++++++
 .../hadoop/yarn/api/records/TestResource.java | 25 +++++++++++++++++++
 .../webapp/NodeLabelsPage.java                |  2 +-
 4 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/pom.xml b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/pom.xml
index 61747c2cd80..e4b8ee28227 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/pom.xml
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/pom.xml
@@ -115,6 +115,11 @@
       <groupId>com.fasterxml.jackson.core</groupId>
       <artifactId>jackson-annotations</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/Resource.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/Resource.java
index 0c10e017685..80e569d5a9e 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/Resource.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/Resource.java
@@ -29,6 +29,7 @@ import org.apache.hadoop.classification.InterfaceAudience.Public;
 import org.apache.hadoop.classification.InterfaceStability;
 import org.apache.hadoop.classification.InterfaceStability.Evolving;
 import org.apache.hadoop.classification.InterfaceStability.Stable;
+import org.apache.hadoop.classification.VisibleForTesting;
 import org.apache.hadoop.util.StringUtils;
 import org.apache.hadoop.yarn.api.ApplicationMasterProtocol;
 import org.apache.hadoop.yarn.api.protocolrecords.ResourceTypes;
@@ -543,4 +544,14 @@ public abstract class Resource implements Comparable<Resource> {
     ri.setMaximumAllocation(Long.MAX_VALUE);
     return ri;
   }
+
+  @VisibleForTesting
+  protected void setResources(ResourceInformation[] resources) {
+    this.resources = resources;
+  }
+
+  public String getFormattedString(long memory) {
+    return getFormattedString(
+        StringUtils.byteDesc(memory * 1024 * 1024));
+  }
 }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/api/records/TestResource.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/api/records/TestResource.java
index 638ecf9d322..060ca39c697 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/api/records/TestResource.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/api/records/TestResource.java
@@ -20,6 +20,8 @@ package org.apache.hadoop.yarn.api.records;
 import org.junit.jupiter.api.Test;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
 
 /**
  * The class to test {@link Resource}.
@@ -42,4 +44,27 @@ class TestResource {
         "Cast to Integer.MAX_VALUE if the long is greater than "
             + "Integer.MAX_VALUE");
   }
+
+  @Test
+  public void testResourceFormatted() {
+    Resource resource = spy(Resource.class);
+    resource.setResources(new ResourceInformation[0]);
+    when(resource.getVirtualCores()).thenReturn(1);
+
+    // We set 10MB
+    String expectedResult1 = "<memory:10 MB, vCores:1>";
+    assertEquals(expectedResult1, resource.getFormattedString(10));
+
+    // We set 1024 MB = 1GB
+    String expectedResult2 = "<memory:1 GB, vCores:1>";
+    assertEquals(expectedResult2, resource.getFormattedString(1024));
+
+    // We set 1024 * 1024 MB = 1024 GB = 1TB
+    String expectedResult3 = "<memory:1 TB, vCores:1>";
+    assertEquals(expectedResult3, resource.getFormattedString(1024 * 1024));
+
+    // We set 1024 * 1024 * 1024 MB = 1024 * 1024 GB = 1 * 1024 TB = 1 PB
+    String expectedResult4 = "<memory:1 PB, vCores:1>";
+    assertEquals(expectedResult4, resource.getFormattedString(1024 * 1024 * 1024));
+  }
 }
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/NodeLabelsPage.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodeLabelsPage.java
index 6ff76281007..c4df6aa0e27 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodeLabelsPage.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodeLabelsPage.java
@@ -75,7 +75,7 @@ public class NodeLabelsPage extends RmView {
         } else {
           row = row.td(String.valueOf(nActiveNMs));
         }
-        row.td(info.getResource().toString()).__();
+        row.td(info.getResource().toFormattedString()).__();
       }
       tbody.__().__();
     }