From 6926fd0ec634df2576bbc9f45e9636b99260db72 Mon Sep 17 00:00:00 2001
From: Wangda Tan <wangda@apache.org>
Date: Thu, 4 Oct 2018 10:31:33 -0700
Subject: [PATCH] YARN-8758. Support getting PreemptionMessage when using
 AMRMClientAsyn. (Zian Chen via wangda)

Change-Id: Ibf5d165f49957b582eeadeb41dc285c84d2f05e7
---
 .../hadoop/yarn/client/api/async/AMRMClientAsync.java | 11 +++++++++++
 .../client/api/async/impl/AMRMClientAsyncImpl.java    |  9 +++++++++
 2 files changed, 20 insertions(+)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/async/AMRMClientAsync.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/async/AMRMClientAsync.java
index 3dd2f718ba..c9f4e5f79a 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/async/AMRMClientAsync.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/async/AMRMClientAsync.java
@@ -39,6 +39,7 @@ import org.apache.hadoop.yarn.api.records.ContainerUpdateType;
 import org.apache.hadoop.yarn.api.records.ExecutionType;
 import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
 import org.apache.hadoop.yarn.api.records.NodeReport;
+import org.apache.hadoop.yarn.api.records.PreemptionMessage;
 import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.RejectedSchedulingRequest;
 import org.apache.hadoop.yarn.api.records.Resource;
@@ -553,6 +554,16 @@ extends AbstractService {
     public void onRequestsRejected(
         List<RejectedSchedulingRequest> rejectedSchedulingRequests) {
     }
+
+    /**
+     * Called when the RM responds to a heartbeat with preemption message
+     * @param preemptionMessage
+     */
+    @Public
+    @Unstable
+    public void onPreemptionMessageReceived(
+        PreemptionMessage preemptionMessage) {
+    }
   }
 
   /**
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/async/impl/AMRMClientAsyncImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/async/impl/AMRMClientAsyncImpl.java
index 3cf2c3496e..922b185a1a 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/async/impl/AMRMClientAsyncImpl.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/async/impl/AMRMClientAsyncImpl.java
@@ -37,6 +37,7 @@ import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.hadoop.yarn.api.records.ContainerStatus;
 import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
 import org.apache.hadoop.yarn.api.records.NodeReport;
+import org.apache.hadoop.yarn.api.records.PreemptionMessage;
 import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.RejectedSchedulingRequest;
 import org.apache.hadoop.yarn.api.records.Resource;
@@ -401,6 +402,14 @@ extends AMRMClientAsync<T> {
             handler.onContainersAllocated(allocated);
           }
 
+          PreemptionMessage preemptionMessage = response.getPreemptionMessage();
+          if (preemptionMessage != null) {
+            if (handler instanceof AMRMClientAsync.AbstractCallbackHandler) {
+              ((AMRMClientAsync.AbstractCallbackHandler) handler)
+                  .onPreemptionMessageReceived(preemptionMessage);
+            }
+          }
+
           if (!response.getContainersFromPreviousAttempts().isEmpty()) {
             if (handler instanceof AMRMClientAsync.AbstractCallbackHandler) {
               ((AMRMClientAsync.AbstractCallbackHandler) handler)