ZOOKEEPER-3682: Stop initializing new SSL connection if ZK server is …

…shutting down

Author: Jie Huang <jiehuang@fb.com>
Author: Fangmin Lyu <allenlyu@fb.com>

Reviewers: andor@apache.org, fangmin@apache.org

Closes #1210 from jhuan31/ZOOKEEPER-3682

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

@@ -555,13 +555,6 @@ public class NIOServerCnxn extends ServerCnxn {
         return true;
     }
 
-    /**
-     * @return true if the server is running, false otherwise.
-     */
-    boolean isZKServerRunning() {
-        return zkServer != null && zkServer.isRunning();
-    }
-
     /*
      * (non-Javadoc)
      *

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

@@ -224,6 +224,15 @@ public class NettyServerCnxnFactory extends ServerCnxnFactory {
             NettyServerCnxn cnxn = new NettyServerCnxn(channel, zkServer, NettyServerCnxnFactory.this);
             ctx.channel().attr(CONNECTION_ATTRIBUTE).set(cnxn);
 
+            // Check the zkServer assigned to the cnxn is still running,
+            // close it before starting the heavy TLS handshake
+            if (!cnxn.isZKServerRunning()) {
+                LOG.warn("Zookeeper server is not running, close the connection before starting the TLS handshake");
+                ServerMetrics.getMetrics().CNXN_CLOSED_WITHOUT_ZK_SERVER_RUNNING.add(1);
+                channel.close();
+                return;
+            }
+
             if (handshakeThrottlingEnabled) {
                 // Favor to check and throttling even in dual mode which
                 // accepts both secure and insecure connections, since

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

@@ -601,6 +601,13 @@ public abstract class ServerCnxn implements Stats, Watcher {
         }
     }
 
+    /**
+     * @return true if the server is running, false otherwise.
+     */
+    public boolean isZKServerRunning() {
+        return zkServer != null && zkServer.isRunning();
+    }
+
     /**
      * Returns the IP address or empty string.
      */

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

@@ -230,6 +230,8 @@ public final class ServerMetrics {
 
         DIGEST_MISMATCHES_COUNT = metricsContext.getCounter("digest_mismatches_count");
         TLS_HANDSHAKE_EXCEEDED = metricsContext.getCounter("tls_handshake_exceeded");
+
+        CNXN_CLOSED_WITHOUT_ZK_SERVER_RUNNING = metricsContext.getCounter("cnxn_closed_without_zk_server_running");
     }
 
     /**
@@ -444,6 +446,8 @@ public final class ServerMetrics {
 
     public final Counter TLS_HANDSHAKE_EXCEEDED;
 
+    public final Counter CNXN_CLOSED_WITHOUT_ZK_SERVER_RUNNING;
+
     private final MetricsProvider metricsProvider;
 
     public void resetAll() {