From fd84ca5161d171f7e754b9b06623c6118e048066 Mon Sep 17 00:00:00 2001
From: Zhankun Tang <ztang@apache.org>
Date: Thu, 24 Oct 2019 10:05:00 +0800
Subject: [PATCH] YARN-9921. Issue in PlacementConstraint when YARN Service AM
 retries allocation on component failure. Contributed by Tarun Parimi

---
 .../impl/pb/SchedulingRequestPBImpl.java      | 23 +++++++++++++++++--
 ...SingleConstraintAppPlacementAllocator.java | 16 +++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/SchedulingRequestPBImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/SchedulingRequestPBImpl.java
index a53dca1be1..222ea63203 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/SchedulingRequestPBImpl.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/SchedulingRequestPBImpl.java
@@ -277,8 +277,27 @@ public boolean equals(Object other) {
     if (other == null) {
       return false;
     }
-    if (other.getClass().isAssignableFrom(this.getClass())) {
-      return this.getProto().equals(this.getClass().cast(other).getProto());
+    if (other instanceof SchedulingRequest) {
+      if (this == other) {
+        return true;
+      }
+      SchedulingRequest that = (SchedulingRequest) other;
+      if (getAllocationRequestId() != that.getAllocationRequestId()) {
+        return false;
+      }
+      if (!getAllocationTags().equals(that.getAllocationTags())) {
+        return false;
+      }
+      if (!getPriority().equals(that.getPriority())) {
+        return false;
+      }
+      if(!getExecutionType().equals(that.getExecutionType())) {
+        return false;
+      }
+      if(!getResourceSizing().equals(that.getResourceSizing())) {
+        return false;
+      }
+      return getPlacementConstraint().equals(that.getPlacementConstraint());
     }
     return false;
   }
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 e651831266..d7fa600060 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
@@ -226,6 +226,22 @@ public void testSchedulingRequestUpdate() {
     schedulingRequest.getResourceSizing().setNumAllocations(10);
     allocator.updatePendingAsk(schedulerRequestKey, schedulingRequest, false);
 
+    // Update allocator with a newly constructed scheduling request different at
+    // #allocations, should succeeded.
+    SchedulingRequest newSchedulingRequest =
+        SchedulingRequest.newBuilder().executionType(
+            ExecutionTypeRequest.newInstance(ExecutionType.GUARANTEED))
+            .allocationRequestId(10L).priority(Priority.newInstance(1))
+            .placementConstraintExpression(PlacementConstraints
+                .targetNotIn(PlacementConstraints.NODE,
+                    PlacementConstraints.PlacementTargets.nodePartition(""),
+                    PlacementConstraints.PlacementTargets
+                        .allocationTag("mapper", "reducer"))
+                .build()).resourceSizing(
+            ResourceSizing.newInstance(11, Resource.newInstance(1024, 1)))
+            .build();
+    allocator.updatePendingAsk(schedulerRequestKey, newSchedulingRequest, false);
+
     // Update allocator with scheduling request different at resource,
     // should failed.
     schedulingRequest.getResourceSizing().setResources(