ZOOKEEPER-4712: Fix partially shutdown of ZooKeeperServer and its processors

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

@@ -294,7 +294,7 @@ public long loadDataBase() throws IOException {
     }
 
     /**
-     * Fast forward the database adding transactions from the committed log into memory.
+     * Fast-forward the database adding transactions from the committed log into memory.
      * @return the last valid zxid.
      * @throws IOException
      */

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

@@ -913,7 +913,7 @@ protected void setState(State state) {
      * @return true if the server is running or server hits an error, false
      *         otherwise.
      */
-    protected boolean canShutdown() {
+    private boolean canShutdown() {
         return state == State.RUNNING || state == State.ERROR;
     }
 
@@ -930,21 +930,43 @@ public final void shutdown() {
 
     /**
      * Shut down the server instance
-     * @param fullyShutDown true if another server using the same database will not replace this one in the same process
+     * @param fullyShutDown true when no other server will use the same database to replace this one
      */
-    public synchronized void shutdown(boolean fullyShutDown) {
-        if (!canShutdown()) {
-            if (fullyShutDown && zkDb != null) {
-                zkDb.clear();
+    public final synchronized void shutdown(boolean fullyShutDown) {
+        if (canShutdown()) {
+            LOG.info("Shutting down");
+
+            shutdownComponents();
+
+            if (zkDb != null && !fullyShutDown) {
+                // There is no need to clear the database if we are going to reuse it:
+                //  * When a new quorum is established we can still apply the diff
+                //    on top of the same zkDb data
+                //  * If we fetch a new snapshot from leader, the zkDb will be
+                //    cleared anyway before loading the snapshot
+                try {
+                    // This will fast-forward the database to the last recorded transaction
+                    zkDb.fastForwardDataBase();
+                } catch (IOException e) {
+                    LOG.error("Error updating DB", e);
+                    fullyShutDown = true;
+                }
             }
+            setState(State.SHUTDOWN);
+        } else {
             LOG.debug("ZooKeeper server is not running, so not proceeding to shutdown!");
-            return;
         }
-        LOG.info("shutting down");
-
-        // new RuntimeException("Calling shutdown").printStackTrace();
-        setState(State.SHUTDOWN);
+        if (zkDb != null && fullyShutDown) {
+            zkDb.clear();
+        }
+    }
 
+    /**
+     * @implNote
+     * Shuts down components owned by this class;
+     * remember to call super.shutdownComponents() when overriding!
+     */
+    protected void shutdownComponents() {
         // unregister all metrics that are keeping a strong reference to this object
         // subclasses will do their specific clean up
         unregisterMetrics();
@@ -953,9 +975,8 @@ public synchronized void shutdown(boolean fullyShutDown) {
             requestThrottler.shutdown();
         }
 
-        // Since sessionTracker and syncThreads poll we just have to
-        // set running to false and they will detect it during the poll
-        // interval.
+        // Since sessionTracker and syncThreads poll we just have to set running to false,
+        // and they will detect it during the poll interval.
         if (sessionTracker != null) {
             sessionTracker.shutdown();
         }
@@ -966,25 +987,6 @@ public synchronized void shutdown(boolean fullyShutDown) {
             jvmPauseMonitor.serviceStop();
         }
 
-        if (zkDb != null) {
-            if (fullyShutDown) {
-                zkDb.clear();
-            } else {
-                // else there is no need to clear the database
-                //  * When a new quorum is established we can still apply the diff
-                //    on top of the same zkDb data
-                //  * If we fetch a new snapshot from leader, the zkDb will be
-                //    cleared anyway before loading the snapshot
-                try {
-                    //This will fast forward the database to the latest recorded transactions
-                    zkDb.fastForwardDataBase();
-                } catch (IOException e) {
-                    LOG.error("Error updating DB", e);
-                    zkDb.clear();
-                }
-            }
-        }
-
         requestPathMetricsCollector.shutdown();
         unregisterJMX();
     }

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

@@ -192,9 +192,7 @@ public void runFromConfig(ServerConfig config) throws IOException, AdminServerEx
             if (secureCnxnFactory != null) {
                 secureCnxnFactory.join();
             }
-            if (zkServer.canShutdown()) {
-                zkServer.shutdown(true);
-            }
+            zkServer.shutdown(true);
         } catch (InterruptedException e) {
             // warn, but generally this is ok
             LOG.warn("Server interrupted", e);

zookeeper-server/src/main/java/org/apache/zookeeper/server/persistence/FileTxnSnapLog.java

@@ -313,7 +313,7 @@ public long restore(DataTree dt, Map<Long, Integer> sessions, PlayBackListener l
     }
 
     /**
-     * This function will fast forward the server database to have the latest
+     * This function will fast-forward the server database to have the latest
      * transactions in it.  This is the same as restore, but only reads from
      * the transaction logs and not restores from a snapshot.
      * @param dt the datatree to write transactions to.

zookeeper-server/src/main/java/org/apache/zookeeper/server/quorum/LeaderZooKeeperServer.java

@@ -155,11 +155,11 @@ protected void unregisterMetrics() {
     }
 
     @Override
-    public synchronized void shutdown(boolean fullyShutDown) {
+    protected void shutdownComponents() {
         if (containerManager != null) {
             containerManager.stop();
         }
-        super.shutdown(fullyShutDown);
+        super.shutdownComponents();
     }
 
     @Override

zookeeper-server/src/main/java/org/apache/zookeeper/server/quorum/LearnerZooKeeperServer.java

@@ -152,13 +152,7 @@ protected void unregisterJMX(Learner peer) {
     }
 
     @Override
-    public synchronized void shutdown(boolean fullyShutDown) {
-        if (!canShutdown()) {
-            super.shutdown(fullyShutDown);
-            LOG.debug("ZooKeeper server is not running, so not proceeding to shutdown!");
-            return;
-        }
-        LOG.info("Shutting down");
+    protected void shutdownComponents() {
         try {
             if (syncProcessor != null) {
                 syncProcessor.shutdown();
@@ -167,7 +161,7 @@ public synchronized void shutdown(boolean fullyShutDown) {
             LOG.warn("Ignoring unexpected exception in syncprocessor shutdown", e);
         }
         try {
-            super.shutdown(fullyShutDown);
+            super.shutdownComponents();
         } catch (Exception e) {
             LOG.warn("Ignoring unexpected exception during shutdown", e);
         }

zookeeper-server/src/main/java/org/apache/zookeeper/server/quorum/ReadOnlyZooKeeperServer.java

@@ -190,12 +190,7 @@ public long getServerId() {
     }
 
     @Override
-    public synchronized void shutdown(boolean fullyShutDown) {
-        if (!canShutdown()) {
-            super.shutdown(fullyShutDown);
-            LOG.debug("ZooKeeper server is not running, so not proceeding to shutdown!");
-            return;
-        }
+    protected void shutdownComponents() {
         shutdown = true;
         unregisterJMX(this);
 
@@ -207,7 +202,7 @@ public synchronized void shutdown(boolean fullyShutDown) {
         self.adminServer.setZooKeeperServer(null);
 
         // shutdown the server itself
-        super.shutdown(fullyShutDown);
+        super.shutdownComponents();
     }
 
     @Override

zookeeper-server/src/test/java/org/apache/zookeeper/server/quorum/WatchLeakTest.java

@@ -97,7 +97,7 @@ public Boolean answer(InvocationOnMock invocation) throws Throwable {
             }
         });
 
-        ZKDatabase database = new ZKDatabase(null);
+        ZKDatabase database = new ZKDatabase(mock(FileTxnSnapLog.class));
         database.setlastProcessedZxid(2L);
         QuorumPeer quorumPeer = mock(QuorumPeer.class);
         FileTxnSnapLog logfactory = mock(FileTxnSnapLog.class);