ZOOKEEPER-4236 Java Client SendThread create many unnecessary Login objects (#2128)

zookeeper-server/src/main/java/org/apache/zookeeper/ClientCnxn.java

@@ -42,6 +42,7 @@ import java.util.concurrent.CopyOnWriteArraySet;
 import java.util.concurrent.LinkedBlockingDeque;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.security.auth.login.LoginException;
 import javax.security.sasl.SaslException;
 import org.apache.jute.BinaryInputArchive;
@@ -842,6 +843,7 @@ public class ClientCnxn {
         private final ClientCnxnSocket clientCnxnSocket;
         private boolean isFirstConnect = true;
         private volatile ZooKeeperSaslClient zooKeeperSaslClient;
+        private final AtomicReference<Login> loginRef = new AtomicReference<>();
 
         void readResponse(ByteBuffer incomingBuffer) throws IOException {
             ByteBufferInputStream bbis = new ByteBufferInputStream(incomingBuffer);
@@ -1087,10 +1089,8 @@ public class ClientCnxn {
             setName(getName().replaceAll("\\(.*\\)", "(" + hostPort + ")"));
             if (clientConfig.isSaslClientEnabled()) {
                 try {
-                    if (zooKeeperSaslClient != null) {
-                        zooKeeperSaslClient.shutdown();
-                    }
-                    zooKeeperSaslClient = new ZooKeeperSaslClient(SaslServerPrincipal.getServerPrincipal(addr, clientConfig), clientConfig);
+                    zooKeeperSaslClient = new ZooKeeperSaslClient(
+                        SaslServerPrincipal.getServerPrincipal(addr, clientConfig), clientConfig, loginRef);
                 } catch (LoginException e) {
                     // An authentication error occurred when the SASL client tried to initialize:
                     // for Kerberos this means that the client failed to authenticate with the KDC.
@@ -1258,8 +1258,9 @@ public class ClientCnxn {
             }
             eventThread.queueEvent(new WatchedEvent(Event.EventType.None, Event.KeeperState.Closed, null));
 
-            if (zooKeeperSaslClient != null) {
-                zooKeeperSaslClient.shutdown();
+            Login l = loginRef.getAndSet(null);
+            if (l != null) {
+                l.shutdown();
             }
             ZooTrace.logTraceMessage(
                 LOG,
@@ -1431,6 +1432,11 @@ public class ClientCnxn {
         public ZooKeeperSaslClient getZooKeeperSaslClient() {
             return zooKeeperSaslClient;
         }
+
+        // VisibleForTesting
+        Login getLogin() {
+            return loginRef.get();
+        }
     }
 
     /**

zookeeper-server/src/main/java/org/apache/zookeeper/client/ZooKeeperSaslClient.java

@@ -21,6 +21,7 @@ package org.apache.zookeeper.client;
 import java.io.IOException;
 import java.security.PrivilegedActionException;
 import java.security.PrivilegedExceptionAction;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.security.auth.Subject;
 import javax.security.auth.login.AppConfigurationEntry;
 import javax.security.auth.login.Configuration;
@@ -64,7 +65,6 @@ public class ZooKeeperSaslClient {
      */
     @Deprecated
     public static final String ENABLE_CLIENT_SASL_DEFAULT = "true";
-    private volatile boolean initializedLogin = false;
 
     /**
      * Returns true if the SASL client is enabled. By default, the client
@@ -112,7 +112,7 @@ public class ZooKeeperSaslClient {
         return null;
     }
 
-    public ZooKeeperSaslClient(final String serverPrincipal, ZKClientConfig clientConfig) throws LoginException {
+    public ZooKeeperSaslClient(final String serverPrincipal, ZKClientConfig clientConfig, AtomicReference<Login> loginRef) throws LoginException {
         /**
          * ZOOKEEPER-1373: allow system property to specify the JAAS
          * configuration section that the zookeeper client should use.
@@ -136,7 +136,8 @@ public class ZooKeeperSaslClient {
         }
         if (entries != null) {
             this.configStatus = "Will attempt to SASL-authenticate using Login Context section '" + clientSection + "'";
-            this.saslClient = createSaslClient(serverPrincipal, clientSection);
+            this.saslClient = createSaslClient(serverPrincipal, clientSection, loginRef);
+            this.login = loginRef.get();
         } else {
             // Handle situation of clientSection's being null: it might simply because the client does not intend to
             // use SASL, so not necessarily an error.
@@ -234,26 +235,25 @@ public class ZooKeeperSaslClient {
 
     private SaslClient createSaslClient(
         final String servicePrincipal,
-        final String loginContext) throws LoginException {
+        final String loginContext,
+        final AtomicReference<Login> loginRef) throws LoginException {
         try {
-            if (!initializedLogin) {
-                synchronized (this) {
-                    if (login == null) {
-                        LOG.debug("JAAS loginContext is: {}", loginContext);
-                        // note that the login object is static: it's shared amongst all zookeeper-related connections.
-                        // in order to ensure the login is initialized only once, it must be synchronized the code snippet.
-                        login = new Login(loginContext, new SaslClientCallbackHandler(null, "Client"), clientConfig);
-                        login.startThreadIfNeeded();
-                        initializedLogin = true;
-                    }
+            if (loginRef.get() == null) {
+                LOG.debug("JAAS loginContext is: {}", loginContext);
+                // note that the login object is static: it's shared amongst all zookeeper-related connections.
+                // in order to ensure the login is initialized only once, it must be synchronized the code snippet.
+                Login l = new Login(loginContext, new SaslClientCallbackHandler(null, "Client"), clientConfig);
+                if (loginRef.compareAndSet(null, l)) {
+                    l.startThreadIfNeeded();
                 }
             }
-            return SecurityUtils.createSaslClient(login.getSubject(), servicePrincipal, "zookeeper", "zk-sasl-md5", LOG, "Client");
+            return SecurityUtils.createSaslClient(loginRef.get().getSubject(),
+                servicePrincipal, "zookeeper", "zk-sasl-md5", LOG, "Client");
         } catch (LoginException e) {
             // We throw LoginExceptions...
             throw e;
         } catch (Exception e) {
-            // ..but consume (with a log message) all other types of exceptions.
+            // ...but consume (with a log message) all other types of exceptions.
             LOG.error("Exception while trying to create SASL client.", e);
             return null;
         }
@@ -451,15 +451,4 @@ public class ZooKeeperSaslClient {
             return false;
         }
     }
-
-    /**
-     * close login thread if running
-     */
-    public void shutdown() {
-        if (null != login) {
-            login.shutdown();
-            login = null;
-        }
-    }
-
 }

zookeeper-server/src/test/java/org/apache/zookeeper/SaslAuthTest.java

@@ -20,6 +20,7 @@ package org.apache.zookeeper;
 
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -239,12 +240,9 @@ public class SaslAuthTest extends ClientBase {
             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);
+            assertNull(sendThread.getLogin());
             assertFalse(sendThread.isAlive(), "SendThread did not shutdown after authFail");
             assertFalse(eventThread.isAlive(), "EventThread did not shutdown after authFail");
         } finally {
@@ -253,4 +251,37 @@ public class SaslAuthTest extends ClientBase {
             }
         }
     }
+
+    @Test
+    public void testDisconnectNotCreatingLoginThread() throws Exception {
+        MyWatcher watcher = new MyWatcher();
+        ZooKeeper zk = null;
+        try {
+            zk = new ZooKeeper(hostPort, CONNECTION_TIMEOUT, watcher);
+            watcher.waitForConnected(CONNECTION_TIMEOUT);
+            zk.getData("/", false, null);
+
+            Field cnxnField = zk.getClass().getDeclaredField("cnxn");
+            cnxnField.setAccessible(true);
+            ClientCnxn clientCnxn = (ClientCnxn) cnxnField.get(zk);
+            Field sendThreadField = clientCnxn.getClass().getDeclaredField("sendThread");
+            sendThreadField.setAccessible(true);
+            SendThread sendThread = (SendThread) sendThreadField.get(clientCnxn);
+
+            Login l1 = sendThread.getLogin();
+            assertNotNull(l1);
+
+            stopServer();
+            watcher.waitForDisconnected(CONNECTION_TIMEOUT);
+            startServer();
+            watcher.waitForConnected(CONNECTION_TIMEOUT);
+            zk.getData("/", false, null);
+
+            assertSame("Login thread should not been recreated on disconnect", l1, sendThread.getLogin());
+        } finally {
+            if (zk != null) {
+                zk.close();
+            }
+        }
+    }
 }