From c8f4673f2b6c98c22534a6cebb19009497f318b2 Mon Sep 17 00:00:00 2001
From: Junfan Zhang <zuston@apache.org>
Date: Fri, 12 Jul 2024 11:33:09 +0800
Subject: [PATCH] YARN-11704. Avoid nested 'AND' placement constraint for non
 tags in scheduling request

---
 .../MemoryPlacementConstraintManager.java      |  4 ++++
 .../TestPlacementConstraintManagerService.java | 18 ++++--------------
 2 files changed, 8 insertions(+), 14 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/constraint/MemoryPlacementConstraintManager.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/MemoryPlacementConstraintManager.java
index 6308b41058a75..c0339b41cfeda 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/MemoryPlacementConstraintManager.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/MemoryPlacementConstraintManager.java
@@ -266,6 +266,10 @@ public PlacementConstraint getMultilevelConstraint(ApplicationId appId,
             .distinct()
             .collect(Collectors.toList());
 
+    if (allConstraints != null && allConstraints.size() == 1) {
+      return allConstraints.get(0).build();
+    }
+
     // Compose an AND constraint
     // When merge request(RC), app(AC) and global constraint(GC),
     // we do a merge on them with CC=AND(GC, AC, RC) and returns a
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/constraint/TestPlacementConstraintManagerService.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/TestPlacementConstraintManagerService.java
index d9376037de67d..3b3c73963c861 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/TestPlacementConstraintManagerService.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/constraint/TestPlacementConstraintManagerService.java
@@ -194,10 +194,8 @@ public void testGetRequestConstraint() {
     // AC = null
     // GC = null
     constraint = pcm.getMultilevelConstraint(appId1, null, c1);
-    Assert.assertTrue(constraint.getConstraintExpr() instanceof And);
-    mergedConstraint = (And) constraint.getConstraintExpr();
-    Assert.assertEquals(1, mergedConstraint.getChildren().size());
-    Assert.assertEquals(c1, mergedConstraint.getChildren().get(0).build());
+    Assert.assertTrue(constraint.getConstraintExpr().getClass() == c1.getConstraintExpr().getClass());
+    Assert.assertEquals(c1, constraint);
 
     // RC = null
     // AC = tag1->c1, tag2->c2
@@ -213,11 +211,7 @@ public void testGetRequestConstraint() {
     Assert.assertEquals(0, mergedConstraint.getChildren().size());
     // if a mapping is found for a given source tag
     constraint = pcm.getMultilevelConstraint(appId1, sourceTag1, null);
-    Assert.assertTrue(constraint.getConstraintExpr() instanceof And);
-    mergedConstraint = (And) constraint.getConstraintExpr();
-    // AND(c1)
-    Assert.assertEquals(1, mergedConstraint.getChildren().size());
-    Assert.assertEquals(c1, mergedConstraint.getChildren().get(0).build());
+    Assert.assertEquals(c1, constraint);
     pcm.unregisterApplication(appId1);
 
     // RC = null
@@ -226,11 +220,7 @@ public void testGetRequestConstraint() {
     pcm.addGlobalConstraint(sourceTag1, c1, true);
     constraint = pcm.getMultilevelConstraint(appId1,
         Sets.newHashSet(sourceTag1), null);
-    Assert.assertTrue(constraint.getConstraintExpr() instanceof And);
-    mergedConstraint = (And) constraint.getConstraintExpr();
-    // AND(c1)
-    Assert.assertEquals(1, mergedConstraint.getChildren().size());
-    Assert.assertEquals(c1, mergedConstraint.getChildren().get(0).build());
+    Assert.assertEquals(c1, constraint);
     pcm.removeGlobalConstraint(sourceTag1);
 
     // RC = c2