From 5c4280fe5e87fadeaa7a3beb29c3d9633aab0548 Mon Sep 17 00:00:00 2001
From: zhengchenyu001 <zhengchenyu001@ke.com>
Date: Wed, 3 Aug 2022 17:46:38 +0800
Subject: [PATCH] TEZ-4440. When tez app run in yarn fed cluster, may throw NPE

---
 .../org/apache/tez/dag/app/rm/DagAwareYarnTaskScheduler.java | 3 ++-
 .../org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java  | 5 +++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/DagAwareYarnTaskScheduler.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/DagAwareYarnTaskScheduler.java
index 29f350f9d0..1b8e19176d 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/DagAwareYarnTaskScheduler.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/DagAwareYarnTaskScheduler.java
@@ -780,7 +780,8 @@ public void onError(Throwable e) {
 
   @Override
   public Resource getAvailableResources() {
-    return client.getAvailableResources();
+    Resource resource = client.getAvailableResources();
+    return resource == null ? Resource.newInstance(0, 0) : resource;
   }
 
   @Override
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java
index 6f51566f75..ea0bdb92e8 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java
@@ -254,7 +254,8 @@ public YarnTaskSchedulerService(TaskSchedulerContext taskSchedulerContext) {
 
   @Override
   public Resource getAvailableResources() {
-    return amRmClient.getAvailableResources();
+    Resource resource = amRmClient.getAvailableResources();
+    return resource == null ? Resource.newInstance(0, 0) : resource;
   }
 
   @Override
@@ -1166,7 +1167,7 @@ boolean preemptIfNeeded() {
     ContainerId[] preemptedContainers = null;
     int numPendingRequestsToService = 0;
     synchronized (this) {
-      Resource freeResources = amRmClient.getAvailableResources();
+      Resource freeResources = this.getAvailableResources();
       if (LOG.isDebugEnabled()) {
         LOG.debug(constructPreemptionPeriodicLog(freeResources));
       } else {