From 92abd994504b3756ee83b2e968ce8d382a6250cb Mon Sep 17 00:00:00 2001
From: Ashutosh Gupta <ashutosh.gupta@st.niituniversity.in>
Date: Mon, 8 Aug 2022 21:29:27 +0100
Subject: [PATCH] YARN-11237. Fix Bug while disabling proxy failover with
 Federation (#4658)

---
 .../TestFederationRMFailoverProxyProvider.java |  3 ++-
 .../org/apache/hadoop/yarn/client/RMProxy.java | 18 +++++++++++++++---
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/TestFederationRMFailoverProxyProvider.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/TestFederationRMFailoverProxyProvider.java
index 0a7ee3ff734..74cf5152d5a 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/TestFederationRMFailoverProxyProvider.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/TestFederationRMFailoverProxyProvider.java
@@ -107,7 +107,8 @@ public class TestFederationRMFailoverProxyProvider {
     conf.setBoolean(YarnConfiguration.FEDERATION_FLUSH_CACHE_FOR_RM_ADDR,
         facadeFlushCache);
 
-    conf.setBoolean(YarnConfiguration.RM_HA_ENABLED, true);
+    conf.setBoolean(YarnConfiguration.FEDERATION_ENABLED, true);
+    conf.setBoolean(YarnConfiguration.FEDERATION_FAILOVER_ENABLED, true);
     conf.setBoolean(YarnConfiguration.AUTO_FAILOVER_ENABLED, false);
     conf.set(YarnConfiguration.RM_CLUSTER_ID, "cluster1");
     conf.set(YarnConfiguration.RM_HA_IDS, "rm1,rm2,rm3");
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/client/RMProxy.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/client/RMProxy.java
index eb39cc13445..95c8f2a01ce 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/client/RMProxy.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/client/RMProxy.java
@@ -98,8 +98,7 @@ public class RMProxy<T> {
     YarnConfiguration conf = (configuration instanceof YarnConfiguration)
         ? (YarnConfiguration) configuration
         : new YarnConfiguration(configuration);
-    RetryPolicy retryPolicy = createRetryPolicy(conf,
-        (HAUtil.isHAEnabled(conf) || HAUtil.isFederationFailoverEnabled(conf)));
+    RetryPolicy retryPolicy = createRetryPolicy(conf, isFailoverEnabled(conf));
     return newProxyInstance(conf, protocol, instance, retryPolicy);
   }
 
@@ -126,7 +125,7 @@ public class RMProxy<T> {
       final Class<T> protocol, RMProxy<T> instance, RetryPolicy retryPolicy)
           throws IOException{
     RMFailoverProxyProvider<T> provider;
-    if (HAUtil.isHAEnabled(conf) || HAUtil.isFederationEnabled(conf)) {
+    if (isFailoverEnabled(conf)) {
       provider = instance.createRMFailoverProxyProvider(conf, protocol);
     } else {
       provider = instance.createNonHaRMFailoverProxyProvider(conf, protocol);
@@ -303,4 +302,17 @@ public class RMProxy<T> {
     return RetryPolicies.retryOtherThanRemoteAndSaslException(
         RetryPolicies.TRY_ONCE_THEN_FAIL, exceptionToPolicyMap);
   }
+
+  private static boolean isFailoverEnabled(YarnConfiguration conf) {
+    if (HAUtil.isHAEnabled(conf)) {
+      // Considering Resource Manager HA is enabled.
+      return true;
+    }
+    if (HAUtil.isFederationEnabled(conf) && HAUtil.isFederationFailoverEnabled(conf)) {
+      // Considering both federation and federation failover are enabled.
+      return true;
+    }
+    return false;
+  }
+
 }