From 20416bc62dce85c2c1aecbc36b8373993a3c55c3 Mon Sep 17 00:00:00 2001
From: Peter Bacsko <pbacsko@cloudera.com>
Date: Tue, 9 Mar 2021 11:34:37 +0100
Subject: [PATCH] YARN-10675. Consolidate YARN-10672 and YARN-10447.
 Contributed by Szilard Nemeth.

---
 .../scheduler/capacity/CapacityScheduler.java | 20 +++----------------
 .../scheduler/capacity/TestLeafQueue.java     |  7 +++++--
 2 files changed, 8 insertions(+), 19 deletions(-)

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/capacity/CapacityScheduler.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacityScheduler.java
index 467dacbfbb..d259c93a1f 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacityScheduler.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacityScheduler.java
@@ -249,8 +249,6 @@ public Configuration getConf() {
 
   private CSMaxRunningAppsEnforcer maxRunningEnforcer;
 
-  private boolean activitiesManagerEnabled = true;
-
   public CapacityScheduler() {
     super(CapacityScheduler.class.getName());
     this.maxRunningEnforcer = new CSMaxRunningAppsEnforcer(this);
@@ -351,9 +349,7 @@ void initScheduler(Configuration configuration) throws
       this.workflowPriorityMappingsMgr = new WorkflowPriorityMappingsManager();
 
       this.activitiesManager = new ActivitiesManager(rmContext);
-      if (activitiesManagerEnabled) {
-        activitiesManager.init(conf);
-      }
+      activitiesManager.init(conf);
       initializeQueues(this.conf);
       this.isLazyPreemptionEnabled = conf.getLazyPreemptionEnabled();
 
@@ -411,9 +407,7 @@ void initScheduler(Configuration configuration) throws
   private void startSchedulerThreads() {
     writeLock.lock();
     try {
-      if (activitiesManagerEnabled) {
-        activitiesManager.start();
-      }
+      activitiesManager.start();
       if (scheduleAsynchronously) {
         Preconditions.checkNotNull(asyncSchedulerThreads,
             "asyncSchedulerThreads is null");
@@ -447,9 +441,7 @@ public void serviceStart() throws Exception {
   public void serviceStop() throws Exception {
     writeLock.lock();
     try {
-      if (activitiesManagerEnabled) {
-        this.activitiesManager.stop();
-      }
+      this.activitiesManager.stop();
       if (scheduleAsynchronously && asyncSchedulerThreads != null) {
         for (Thread t : asyncSchedulerThreads) {
           t.interrupt();
@@ -3480,7 +3472,6 @@ public void setMaxRunningAppsEnforcer(CSMaxRunningAppsEnforcer enforcer) {
     this.maxRunningEnforcer = enforcer;
   }
 
-
   /**
    * Returning true as capacity scheduler supports placement constraints.
    */
@@ -3489,11 +3480,6 @@ public boolean placementConstraintEnabled() {
     return true;
   }
 
-  @VisibleForTesting
-  public void setActivitiesManagerEnabled(boolean enabled) {
-    this.activitiesManagerEnabled = enabled;
-  }
-
   @VisibleForTesting
   public void setQueueManager(CapacitySchedulerQueueManager qm) {
     this.queueManager = qm;
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestLeafQueue.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestLeafQueue.java
index 6262d85598..4b023107b3 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestLeafQueue.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestLeafQueue.java
@@ -162,9 +162,13 @@ private void setUpWithNodeLabels() throws Exception {
   private void setUpInternal(ResourceCalculator rC, boolean withNodeLabels)
       throws Exception {
     CapacityScheduler spyCs = new CapacityScheduler();
-    spyCs.setActivitiesManagerEnabled(false);
     queues = new CSQueueStore();
     cs = spy(spyCs);
+
+    //All stub calls on the spy object of the 'cs' field should happen
+    //before cs.start() is invoked. See YARN-10672 for more details.
+    when(cs.getNumClusterNodes()).thenReturn(3);
+
     rmContext = TestUtils.getMockRMContext();
     spyRMContext = spy(rmContext);
 
@@ -231,7 +235,6 @@ private void setUpInternal(ResourceCalculator rC, boolean withNodeLabels)
     when(spyRMContext.getScheduler()).thenReturn(cs);
     when(spyRMContext.getYarnConfiguration())
         .thenReturn(new YarnConfiguration());
-    when(cs.getNumClusterNodes()).thenReturn(3);
     cs.start();
   }