ZOOKEEPER-2623: [ADDENDUM] Fix flaky CheckTest::testCluster (#2067)

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

@@ -19,13 +19,17 @@
 package org.apache.zookeeper.test;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import java.io.File;
 import org.apache.zookeeper.KeeperException;
 import org.apache.zookeeper.TestableZooKeeper;
 import org.apache.zookeeper.ZooDefs;
+import org.apache.zookeeper.client.ZKClientConfig;
 import org.apache.zookeeper.data.Stat;
 import org.apache.zookeeper.proto.CheckVersionRequest;
 import org.apache.zookeeper.proto.ReplyHeader;
 import org.apache.zookeeper.proto.RequestHeader;
+import org.apache.zookeeper.server.ZKDatabase;
+import org.apache.zookeeper.server.persistence.FileTxnSnapLog;
 import org.apache.zookeeper.server.quorum.QuorumPeer;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -36,6 +40,7 @@ public class CheckTest extends ClientBase {
 
     @BeforeEach
     public void setUp(TestInfo testInfo) throws Exception {
+        System.setProperty(ZKClientConfig.ZOOKEEPER_REQUEST_TIMEOUT, "2000");
         if (testInfo.getDisplayName().contains("Cluster")) {
             return;
         }
@@ -44,6 +49,7 @@ public class CheckTest extends ClientBase {
 
     @AfterEach
     public void tearDown(TestInfo testInfo) throws Exception {
+        System.clearProperty(ZKClientConfig.ZOOKEEPER_REQUEST_TIMEOUT);
         if (testInfo.getDisplayName().contains("Cluster")) {
             return;
         }
@@ -83,11 +89,18 @@ public class CheckTest extends ClientBase {
 
     @Test
     public void testStandalone() throws Exception {
-        TestableZooKeeper zk = createClient();
-        testOperations(zk);
+        testOperations(createClient());
+    }
+
+    @Test
+    public void testStandaloneDatabaseReloadAfterCheck() throws Exception {
+        try {
+            testOperations(createClient());
+        } catch (Exception ignored) {
+            // Ignore to test database reload after check
+        }
         stopServer();
         startServer();
-        createClient();
     }
 
     @Test
@@ -98,13 +111,34 @@ public class CheckTest extends ClientBase {
             testOperations(qb.createClient(new CountdownWatcher(), QuorumPeer.ServerState.OBSERVING));
             testOperations(qb.createClient(new CountdownWatcher(), QuorumPeer.ServerState.FOLLOWING));
             testOperations(qb.createClient(new CountdownWatcher(), QuorumPeer.ServerState.LEADING));
-            int leaderIndex = qb.getLeaderIndex();
-            int leaderPort = qb.getLeaderClientPort();
-            qb.shutdown(qb.getLeaderQuorumPeer());
-            qb.setupServer(leaderIndex + 1);
-            QuorumPeer quorumPeer = qb.getPeerList().get(leaderIndex);
-            quorumPeer.start();
-            qb.createClient("localhost:" + leaderPort, 2 * CONNECTION_TIMEOUT);
+        } finally {
+            try {
+                qb.tearDown();
+            } catch (Exception ignored) {}
+        }
+    }
+
+    @Test
+    public void testClusterDatabaseReloadAfterCheck() throws Exception {
+        QuorumBase qb = new QuorumBase();
+        try {
+            qb.setUp(true, true);
+
+            // Get leader before possible damaging operations to
+            // reduce chance of leader migration and log truncation.
+            File dataDir = qb.getLeaderDataDir();
+            QuorumPeer leader = qb.getLeaderQuorumPeer();
+
+            try {
+                testOperations(qb.createClient(new CountdownWatcher(), QuorumPeer.ServerState.LEADING));
+            } catch (Exception ignored) {
+                // Ignore to test database reload after check
+            }
+            qb.shutdown(leader);
+
+            FileTxnSnapLog txnSnapLog = new FileTxnSnapLog(dataDir, dataDir);
+            ZKDatabase database = new ZKDatabase(txnSnapLog);
+            database.loadDataBase();
         } finally {
             try {
                 qb.tearDown();

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

@@ -342,6 +342,21 @@ public class QuorumBase extends ClientBase {
       return null;
     }
 
+    public File getLeaderDataDir() {
+        if (s1.getPeerState() == ServerState.LEADING) {
+            return s1dir;
+        } else if (s2.getPeerState() == ServerState.LEADING) {
+            return s2dir;
+        } else if (s3.getPeerState() == ServerState.LEADING) {
+            return s3dir;
+        } else if (s4.getPeerState() == ServerState.LEADING) {
+            return s4dir;
+        } else if (s5.getPeerState() == ServerState.LEADING) {
+            return s5dir;
+        }
+        return null;
+    }
+
     public QuorumPeer getFirstObserver() {
       if (s1.getLearnerType() == LearnerType.OBSERVER) {
         return s1;