From 5b557de721fa1587a7d0e197ea6b7120b53e8717 Mon Sep 17 00:00:00 2001
From: Rohith Sharma K S <rohithsharmaks@apache.org>
Date: Fri, 30 Aug 2019 10:35:09 +0530
Subject: [PATCH] YARN-9714. ZooKeeper connection in ZKRMStateStore leaks after
 RM transitioned to standby. Contributed by Tao Yang.

---
 .../yarn/server/resourcemanager/recovery/ZKRMStateStore.java    | 2 +-
 .../server/resourcemanager/recovery/TestZKRMStateStore.java     | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/ZKRMStateStore.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/ZKRMStateStore.java
index dcf891a48d..fd0a2141e6 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/ZKRMStateStore.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/ZKRMStateStore.java
@@ -453,7 +453,7 @@ protected synchronized void closeInternal() throws Exception {
       verifyActiveStatusThread.join(1000);
     }
 
-    if (!HAUtil.isHAEnabled(getConfig())) {
+    if (resourceManager.getZKManager() == null) {
       CuratorFramework curatorFramework = zkManager.getCurator();
       IOUtils.closeStream(curatorFramework);
     }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/TestZKRMStateStore.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/TestZKRMStateStore.java
index f21ff96209..6b6f0131cd 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/TestZKRMStateStore.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/TestZKRMStateStore.java
@@ -453,6 +453,7 @@ public void testZKRootPathAcls() throws Exception {
     rm = new MockRM(conf);
     rm.start();
     rm.getRMContext().getRMAdminService().transitionToActive(req);
+    stateStore = (ZKRMStateStore) rm.getRMContext().getStateStore();
     acls = stateStore.getACL(rootPath);
     assertEquals(acls.size(), 1);
     verifyZKACL("world", "anyone", Perms.ALL, acls);
@@ -463,6 +464,7 @@ public void testZKRootPathAcls() throws Exception {
     rm = new MockRM(conf);
     rm.start();
     rm.getRMContext().getRMAdminService().transitionToActive(req);
+    stateStore = (ZKRMStateStore) rm.getRMContext().getStateStore();
     acls = stateStore.getACL(rootPath);
     assertThat(acls).hasSize(2);
     verifyZKACL("digest", "localhost", Perms.CREATE | Perms.DELETE, acls);