ZOOKEEPER-4367: Zookeeper#Login thread leak in case of Sasl AuthFailed.

https://issues.apache.org/jira/browse/ZOOKEEPER-4367

Things changed in this PR:
1. Moving `zooKeeperSaslClient` from `ClientCnxn` to `SendThread` since only SendThread is creating and accessing `zooKeeperSaslClient`
2. Moved closing of `zooKeeperSaslClient` from `ClientCnxn#disconnect` method to `SendThread#run` method. This will make sure that we will close zooKeeperSaslClient (and in turn close Login object) even if `ClientCnxn#disconnect` is not called which happens when we encounter AuthFailed Exceptions. Also it looks clean that whenever SendThread is terminating we clean up all the class variable.
3. Setting login to null when `ZooKeeperSaslClient#shutdown` is called. This helps me testing whether zooKeeperSaslClient object is shutdown and in turn Login object is shutdown.

Author: Rushabh Shah <rushabh.shah@rushabh-ltmflld.internal.salesforce.com>

Reviewers: Enrico Olivelli <eolivelli@apache.org>, Damien Diederen <ddiederen@apache.org>

Closes #1755 from shahrs87/ZOOKEEPER-4367

5 files changed, 27 insertions, 12 deletions

diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/ClientCnxn.java b/zookeeper-server/src/main/java/org/apache/zookeeper/ClientCnxn.java
index 076662447..347a68f80 100644
--- a/zookeeper-server/src/main/java/org/apache/zookeeper/ClientCnxn.java
+++ b/zookeeper-server/src/main/java/org/apache/zookeeper/ClientCnxn.java
@@ -214,8 +214,6 @@ public class ClientCnxn {
      */
     volatile boolean seenRwServerBefore = false;
 
-    public ZooKeeperSaslClient zooKeeperSaslClient;
-
     private final ZKClientConfig clientConfig;
     /**
      * If any request's response in not received in configured requestTimeout
@@ -870,6 +868,8 @@ public class ClientCnxn {
         private long lastPingSentNs;
         private final ClientCnxnSocket clientCnxnSocket;
         private boolean isFirstConnect = true;
+        private volatile ZooKeeperSaslClient zooKeeperSaslClient;
+
 
         void readResponse(ByteBuffer incomingBuffer) throws IOException {
             ByteBufferInputStream bbis = new ByteBufferInputStream(incomingBuffer);
@@ -1313,6 +1313,10 @@ public class ClientCnxn {
                 eventThread.queueEvent(new WatchedEvent(Event.EventType.None, Event.KeeperState.Disconnected, null));
             }
             eventThread.queueEvent(new WatchedEvent(Event.EventType.None, Event.KeeperState.Closed, null));
+
+            if (zooKeeperSaslClient != null) {
+                zooKeeperSaslClient.shutdown();
+            }
             ZooTrace.logTraceMessage(
                 LOG,
                 ZooTrace.getTextTraceLevel(),
@@ -1486,6 +1490,9 @@ public class ClientCnxn {
             clientCnxnSocket.sendPacket(p);
         }
 
+        public ZooKeeperSaslClient getZooKeeperSaslClient() {
+            return zooKeeperSaslClient;
+        }
     }
 
     /**
@@ -1504,9 +1511,6 @@ public class ClientCnxn {
             LOG.warn("Got interrupted while waiting for the sender thread to close", ex);
         }
         eventThread.queueEventOfDeath();
-        if (zooKeeperSaslClient != null) {
-            zooKeeperSaslClient.shutdown();
-        }
     }
 
     /**
@@ -1739,4 +1743,7 @@ public class ClientCnxn {
         }
     }
 
+    public ZooKeeperSaslClient getZooKeeperSaslClient() {
+        return sendThread.getZooKeeperSaslClient();
+    }
 }
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeper.java b/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeper.java
index a013dd3d5..127f5a2a3 100644
--- a/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeper.java
+++ b/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeper.java
@@ -226,7 +226,7 @@ public class ZooKeeper implements AutoCloseable {
     }
 
     public ZooKeeperSaslClient getSaslClient() {
-        return cnxn.zooKeeperSaslClient;
+        return cnxn.getZooKeeperSaslClient();
     }
 
     private final ZKClientConfig clientConfig;
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/client/ZooKeeperSaslClient.java b/zookeeper-server/src/main/java/org/apache/zookeeper/client/ZooKeeperSaslClient.java
index b0598bca8..940bcfe22 100644
--- a/zookeeper-server/src/main/java/org/apache/zookeeper/client/ZooKeeperSaslClient.java
+++ b/zookeeper-server/src/main/java/org/apache/zookeeper/client/ZooKeeperSaslClient.java
@@ -215,7 +215,7 @@ public class ZooKeeperSaslClient {
             // data[] contains the Zookeeper Server's SASL token.
             // ctx is the ZooKeeperSaslClient object. We use this object's respondToServer() method
             // to reply to the Zookeeper Server's SASL token
-            ZooKeeperSaslClient client = ((ClientCnxn) ctx).zooKeeperSaslClient;
+            ZooKeeperSaslClient client = ((ClientCnxn) ctx).getZooKeeperSaslClient();
             if (client == null) {
                 LOG.warn("sasl client was unexpectedly null: cannot respond to Zookeeper server.");
                 return;
@@ -458,6 +458,7 @@ public class ZooKeeperSaslClient {
     public void shutdown() {
         if (null != login) {
             login.shutdown();
+            login = null;
         }
     }
 
diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/ClientCnxnSocketFragilityTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/ClientCnxnSocketFragilityTest.java
index b5bb0cd94..953414468 100644
--- a/zookeeper-server/src/test/java/org/apache/zookeeper/ClientCnxnSocketFragilityTest.java
+++ b/zookeeper-server/src/test/java/org/apache/zookeeper/ClientCnxnSocketFragilityTest.java
@@ -338,9 +338,6 @@ public class ClientCnxnSocketFragilityTest extends QuorumPeerTestBase {
                 LOG.warn("Got interrupted while waiting for the sender thread to close", ex);
             }
             eventThread.queueEventOfDeath();
-            if (zooKeeperSaslClient != null) {
-                zooKeeperSaslClient.shutdown();
-            }
         }
     }
 
diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/SaslAuthTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/SaslAuthTest.java
index 1d0bffaaf..005cf4bac 100644
--- a/zookeeper-server/src/test/java/org/apache/zookeeper/SaslAuthTest.java
+++ b/zookeeper-server/src/test/java/org/apache/zookeeper/SaslAuthTest.java
@@ -18,6 +18,8 @@
 
 package org.apache.zookeeper;
 
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -33,6 +35,7 @@ import org.apache.zookeeper.ClientCnxn.EventThread;
 import org.apache.zookeeper.ClientCnxn.SendThread;
 import org.apache.zookeeper.Watcher.Event.KeeperState;
 import org.apache.zookeeper.ZooDefs.Ids;
+import org.apache.zookeeper.client.ZooKeeperSaslClient;
 import org.apache.zookeeper.data.ACL;
 import org.apache.zookeeper.data.Id;
 import org.apache.zookeeper.test.ClientBase;
@@ -232,8 +235,16 @@ public class SaslAuthTest extends ClientBase {
             Field eventThreadField = clientCnxn.getClass().getDeclaredField("eventThread");
             eventThreadField.setAccessible(true);
             EventThread eventThread = (EventThread) eventThreadField.get(clientCnxn);
+            ZooKeeperSaslClient zooKeeperSaslClient = clientCnxn.getZooKeeperSaslClient();
+            assertNotNull(zooKeeperSaslClient);
             sendThread.join(CONNECTION_TIMEOUT);
             eventThread.join(CONNECTION_TIMEOUT);
+            Field loginField = zooKeeperSaslClient.getClass().getDeclaredField("login");
+            loginField.setAccessible(true);
+            Login login = (Login) loginField.get(zooKeeperSaslClient);
+            // If login is null, this means ZooKeeperSaslClient#shutdown method has been called which in turns
+            // means that Login#shutdown has been called.
+            assertNull(login);
             assertFalse(sendThread.isAlive(), "SendThread did not shutdown after authFail");
             assertFalse(eventThread.isAlive(), "EventThread did not shutdown after authFail");
         } finally {
@@ -242,5 +253,4 @@ public class SaslAuthTest extends ClientBase {
             }
         }
     }
-
-}
+}
\ No newline at end of file