[To dev/1.3] [Test] Fix Ratis mini-cluster shutdown race (#18879)

iotdb-core/consensus/src/test/java/org/apache/iotdb/consensus/ratis/RecoverReadTest.java

@@ -24,7 +24,9 @@
 import org.apache.iotdb.consensus.common.request.IConsensusRequest;
 import org.apache.iotdb.consensus.config.RatisConfig;
 import org.apache.iotdb.consensus.exception.RatisReadUnavailableException;
+import org.apache.iotdb.consensus.ratis.utils.Utils;
 
+import org.apache.ratis.server.RaftServer;
 import org.apache.ratis.util.TimeDuration;
 import org.junit.After;
 import org.junit.Assert;
@@ -35,6 +37,10 @@
 
 import java.io.File;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -46,6 +52,7 @@ public class RecoverReadTest {
   private static class SlowRecoverStateMachine extends TestUtils.IntegerCounter {
     private final TimeDuration stallDuration;
     private final AtomicBoolean stallApply = new AtomicBoolean(false);
+    private final CountDownLatch stopped = new CountDownLatch(1);
 
     private SlowRecoverStateMachine(TimeDuration stallDuration) {
       this.stallDuration = stallDuration;
@@ -70,6 +77,11 @@ public TSStatus write(IConsensusRequest request) {
       return super.write(request);
     }
 
+    @Override
+    public void stop() {
+      stopped.countDown();
+    }
+
     @Override
     public boolean takeSnapshot(File snapshotDir) {
       // allow no snapshot to ensure the raft log be replayed
@@ -126,6 +138,64 @@ public void tearUp() throws Exception {
     }
   }
 
+  @Test
+  public void stopWaitsForInFlightLeaderCheck() throws Exception {
+    final ConsensusGroupId gid = miniCluster.getGid();
+    for (RatisConsensus server : miniCluster.getServers()) {
+      server.createLocalPeer(gid, miniCluster.getPeers());
+    }
+    miniCluster.waitUntilActiveLeaderElectedAndReady();
+
+    int leaderIndex = 0;
+    while (!miniCluster.getServer(leaderIndex).isLeader(gid)) {
+      leaderIndex++;
+    }
+    final RaftServer.Division leader =
+        miniCluster
+            .getServer(leaderIndex)
+            .getServer()
+            .getDivision(Utils.fromConsensusGroupIdToRaftGroupId(gid));
+    final SlowRecoverStateMachine stateMachine =
+        (SlowRecoverStateMachine) miniCluster.getStateMachines().get(leaderIndex);
+    final CountDownLatch checkStarted = new CountDownLatch(1);
+    final CountDownLatch finishCheck = new CountDownLatch(1);
+    final CountDownLatch stopStarted = new CountDownLatch(1);
+    final ExecutorService executor = Executors.newFixedThreadPool(2);
+    try {
+      final Future<?> check =
+          executor.submit(
+              () -> {
+                // LeaderStateImpl performs its leadership checks under this monitor.
+                synchronized (leader) {
+                  checkStarted.countDown();
+                  Assert.assertTrue(finishCheck.await(30, TimeUnit.SECONDS));
+                  leader.getRaftLog().getLastEntryTermIndex();
+                }
+                return null;
+              });
+      Assert.assertTrue(checkStarted.await(30, TimeUnit.SECONDS));
+      final Future<?> stop =
+          executor.submit(
+              () -> {
+                stopStarted.countDown();
+                miniCluster.stop();
+                return null;
+              });
+      Assert.assertTrue(stopStarted.await(30, TimeUnit.SECONDS));
+      Assert.assertFalse(
+          "The state machine must remain open until the in-flight leader check finishes",
+          stateMachine.stopped.await(5, TimeUnit.SECONDS));
+      finishCheck.countDown();
+      check.get(30, TimeUnit.SECONDS);
+      stop.get(30, TimeUnit.SECONDS);
+      Assert.assertEquals(0, stateMachine.stopped.getCount());
+    } finally {
+      finishCheck.countDown();
+      executor.shutdown();
+      Assert.assertTrue(executor.awaitTermination(30, TimeUnit.SECONDS));
+    }
+  }
+
   /* mimics the situation before this patch */
   @Test
   public void inconsistentReadAfterRestart() throws Exception {

iotdb-core/consensus/src/test/java/org/apache/iotdb/consensus/ratis/TestUtils.java

@@ -39,9 +39,16 @@
 import org.apache.iotdb.consensus.ratis.utils.Retriable;
 import org.apache.iotdb.consensus.ratis.utils.Utils;
 
+import org.apache.ratis.protocol.ClientId;
+import org.apache.ratis.protocol.LeaderElectionManagementRequest;
+import org.apache.ratis.protocol.RaftClientReply;
+import org.apache.ratis.protocol.RaftGroupId;
+import org.apache.ratis.protocol.TransferLeadershipRequest;
+import org.apache.ratis.server.RaftServer;
 import org.apache.ratis.thirdparty.com.google.common.base.Preconditions;
 import org.apache.ratis.util.ExitUtils;
 import org.apache.ratis.util.FileUtils;
+import org.apache.ratis.util.LifeCycle;
 import org.apache.ratis.util.TimeDuration;
 import org.apache.ratis.util.Timestamp;
 import org.junit.Assert;
@@ -305,13 +312,53 @@ void start() throws IOException {
     }
 
     void stop() throws IOException {
-      for (RatisConsensus server : servers) {
-        server.stop();
+      if (isStopped.get()) {
+        return;
+      }
+      try {
+        stepDownBeforeStopping();
+      } finally {
+        for (RatisConsensus server : servers) {
+          server.stop();
+        }
       }
       waitUntilPortsReleased(peerPorts);
       isStopped.set(true);
     }
 
+    private void stepDownBeforeStopping() throws IOException {
+      final ClientId clientId = ClientId.randomId();
+      final List<RaftServer> runningServers = new ArrayList<>();
+      // Ratis 3.2.1 closes the log without joining the leader event processor. Step down through
+      // that processor first, so an in-flight leadership check cannot access an already closed log.
+      // Pause elections on every peer before stepping down to prevent a replacement leader.
+      for (RatisConsensus server : servers) {
+        final RaftServer raftServer = server.getServer();
+        if (raftServer.getLifeCycleState() != LifeCycle.State.RUNNING) {
+          continue;
+        }
+        runningServers.add(raftServer);
+        for (RaftGroupId groupId : raftServer.getGroupIds()) {
+          final RaftClientReply reply =
+              raftServer.leaderElectionManagement(
+                  LeaderElectionManagementRequest.newPause(
+                      clientId, raftServer.getId(), groupId, 0));
+          Assert.assertTrue(reply.toString(), reply.isSuccess());
+        }
+      }
+      for (RaftServer raftServer : runningServers) {
+        for (RaftGroupId groupId : raftServer.getGroupIds()) {
+          if (raftServer.getDivision(groupId).getInfo().isLeader()) {
+            final RaftClientReply reply =
+                raftServer.transferLeadership(
+                    new TransferLeadershipRequest(
+                        clientId, raftServer.getId(), groupId, 1, null, 10_000));
+            Assert.assertTrue(reply.toString(), reply.isSuccess());
+          }
+        }
+      }
+    }
+
     void cleanUp() throws IOException {
       stop();
       for (File storage : peerStorage) {