From 83dcb9d87ec75f2be0acb8972f5f0faefe6ffbcd Mon Sep 17 00:00:00 2001
From: Weiwei Yang <wwei@apache.org>
Date: Fri, 21 Jun 2019 17:40:55 +0800
Subject: [PATCH] YARN-9209. When nodePartition is not set in Placement
 Constraints, containers are allocated only in default partition. Contributed
 by Tarun Parimi.

---
 .../scheduler/AppSchedulingInfo.java               | 14 ++++++++++++++
 .../SingleConstraintAppPlacementAllocator.java     |  5 ++++-
 .../TestSingleConstraintAppPlacementAllocator.java |  7 ++++---
 3 files changed, 22 insertions(+), 4 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/AppSchedulingInfo.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AppSchedulingInfo.java
index c3269fbf7d..4abb7129cc 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AppSchedulingInfo.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/AppSchedulingInfo.java
@@ -795,4 +795,18 @@ public boolean precheckNode(SchedulerRequestKey schedulerKey,
   public Map<String, String> getApplicationSchedulingEnvs() {
     return applicationSchedulingEnvs;
   }
+
+  /**
+   * Get the defaultNodeLabelExpression for the application's current queue.
+   *
+   * @return defaultNodeLabelExpression
+   */
+  public String getDefaultNodeLabelExpression() {
+    try {
+      this.readLock.lock();
+      return queue.getDefaultNodeLabelExpression();
+    } finally {
+      this.readLock.unlock();
+    }
+  }
 }
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/placement/SingleConstraintAppPlacementAllocator.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/placement/SingleConstraintAppPlacementAllocator.java
index e0307b8982..3fc3afbd09 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/placement/SingleConstraintAppPlacementAllocator.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/placement/SingleConstraintAppPlacementAllocator.java
@@ -255,7 +255,10 @@ private void validateAndSetSchedulingRequest(SchedulingRequest
   // Currently only single constraint is handled.
   private String validateAndGetTargetNodePartition(
       PlacementConstraint placementConstraint) {
-    String nodePartition = RMNodeLabelsManager.NO_LABEL;
+    String defaultNodeLabelExpression =
+        appSchedulingInfo.getDefaultNodeLabelExpression();
+    String nodePartition = defaultNodeLabelExpression == null ?
+        RMNodeLabelsManager.NO_LABEL : defaultNodeLabelExpression;
     if (placementConstraint != null &&
         placementConstraint.getConstraintExpr() != null) {
       PlacementConstraint.AbstractConstraint ac =
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/placement/TestSingleConstraintAppPlacementAllocator.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/placement/TestSingleConstraintAppPlacementAllocator.java
index d9b6c20844..e651831266 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/placement/TestSingleConstraintAppPlacementAllocator.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/placement/TestSingleConstraintAppPlacementAllocator.java
@@ -67,7 +67,7 @@ public void setup() throws Exception {
         TestUtils.getMockApplicationId(1));
     when(appSchedulingInfo.getApplicationAttemptId()).thenReturn(
         TestUtils.getMockApplicationAttemptId(1, 1));
-
+    when(appSchedulingInfo.getDefaultNodeLabelExpression()).thenReturn("y");
     // stub RMContext
     rmContext = TestUtils.getMockRMContext();
 
@@ -153,7 +153,8 @@ public void testSchedulingRequestValidation() {
         .resourceSizing(
             ResourceSizing.newInstance(1, Resource.newInstance(1024, 1)))
         .build());
-    Assert.assertEquals("", allocator.getTargetNodePartition());
+    // Node partition is unspecified, use the default node label expression y
+    Assert.assertEquals("y", allocator.getTargetNodePartition());
 
     // Valid (with application Id target)
     assertValidSchedulingRequest(SchedulingRequest.newBuilder().executionType(
@@ -167,7 +168,7 @@ public void testSchedulingRequestValidation() {
             ResourceSizing.newInstance(1, Resource.newInstance(1024, 1)))
         .build());
     // Allocation tags should not include application Id
-    Assert.assertEquals("", allocator.getTargetNodePartition());
+    Assert.assertEquals("y", allocator.getTargetNodePartition());
 
     // Invalid (without sizing)
     assertInvalidSchedulingRequest(SchedulingRequest.newBuilder().executionType(