From 29f2230cb615b975950691044a621ff008088cf5 Mon Sep 17 00:00:00 2001
From: Szilard Nemeth <snemeth@apache.org>
Date: Wed, 25 Jan 2023 15:13:24 +0100
Subject: [PATCH] YARN-5607. Document
 TestContainerResourceUsage#waitForContainerCompletion. Contributed by Susheel
 Gupta

---
 .../yarn/server/resourcemanager/MockRM.java   | 25 +++++++++++++++++++
 .../TestContainerResourceUsage.java           | 18 ++-----------
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockRM.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockRM.java
index 95db1171b3..316f8e06cb 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockRM.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockRM.java
@@ -1011,6 +1011,31 @@ public void waitForAppRemovedFromScheduler(ApplicationId appId)
     LOG.info("app is removed from scheduler, " + appId);
   }
 
+  /**
+   * Wait until a container has reached a completion state.
+   * The timeout is 20 seconds.
+   * @param nm A mock nodemanager
+   * @param rm A mock resourcemanager
+   * @param amContainerId The id of an am container
+   * @param container A container
+   * @throws Exception
+   *         if interrupted while waiting for the completion transition
+   *         or an unexpected error while MockNM is hearbeating.
+   */
+  public static void waitForContainerCompletion(MockRM rm, MockNM nm,
+    ContainerId amContainerId, RMContainer container) throws Exception {
+    ContainerId containerId = container.getContainerId();
+    if (null != rm.scheduler.getRMContainer(containerId)) {
+      if (containerId.equals(amContainerId)) {
+        rm.waitForState(nm, containerId, RMContainerState.COMPLETED);
+      } else {
+        rm.waitForState(nm, containerId, RMContainerState.KILLED);
+      }
+    } else {
+      rm.drainEvents();
+    }
+  }
+
   private void drainEventsImplicitly() {
     if (!disableDrainEventsImplicitly) {
       drainEvents();
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestContainerResourceUsage.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestContainerResourceUsage.java
index 9a01087ce3..54fd650964 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestContainerResourceUsage.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestContainerResourceUsage.java
@@ -348,7 +348,7 @@ private void amRestartTests(boolean keepRunningContainers)
       // If keepRunningContainers is false, all live containers should now
       // be completed. Calculate the resource usage metrics for all of them.
       for (RMContainer c : rmContainers) {
-        waitforContainerCompletion(rm, nm, amContainerId, c);
+        MockRM.waitForContainerCompletion(rm, nm, amContainerId, c);
         AggregateAppResourceUsage ru = calculateContainerResourceMetrics(c);
         memorySeconds += ru.getMemorySeconds();
         vcoreSeconds += ru.getVcoreSeconds();
@@ -400,7 +400,7 @@ private void amRestartTests(boolean keepRunningContainers)
 
     // Calculate container usage metrics for second attempt.
     for (RMContainer c : rmContainers) {
-      waitforContainerCompletion(rm, nm, amContainerId, c);
+      MockRM.waitForContainerCompletion(rm, nm, amContainerId, c);
       AggregateAppResourceUsage ru = calculateContainerResourceMetrics(c);
       memorySeconds += ru.getMemorySeconds();
       vcoreSeconds += ru.getVcoreSeconds();
@@ -417,20 +417,6 @@ private void amRestartTests(boolean keepRunningContainers)
     return;
   }
 
-  private void waitforContainerCompletion(MockRM rm, MockNM nm,
-      ContainerId amContainerId, RMContainer container) throws Exception {
-    ContainerId containerId = container.getContainerId();
-    if (null != rm.scheduler.getRMContainer(containerId)) {
-      if (containerId.equals(amContainerId)) {
-        rm.waitForState(nm, containerId, RMContainerState.COMPLETED);
-      } else {
-        rm.waitForState(nm, containerId, RMContainerState.KILLED);
-      }
-    } else {
-      rm.drainEvents();
-    }
-  }
-
   private AggregateAppResourceUsage calculateContainerResourceMetrics(
       RMContainer rmContainer) {
     Resource resource = rmContainer.getContainer().getResource();