RATIS-2795: fix: gate stale read when snapshots installing

ratis-server/src/main/java/org/apache/ratis/server/impl/ReadRequestHandler.java

@@ -81,6 +81,10 @@ RaftGroupMemberId getMemberId() {
   }
 
   CompletableFuture<RaftClientReply> staleReadAsync(RaftClientRequest request) {
+    final ReadException ex = getReadExceptionIfInstallingSnapshot("stale read");
+    if (ex != null) {
+      return CompletableFuture.completedFuture(server.newExceptionReply(request, ex));
+    }
     final long minIndex = request.getType().getStaleRead().getMinIndex();
     final long commitIndex = state.getLog().getLastCommittedIndex();
     LOG.debug("{}: minIndex={}, commitIndex={} from {}", getMemberId(), minIndex, commitIndex, request.getClientId());
@@ -94,19 +98,20 @@ CompletableFuture<RaftClientReply> staleReadAsync(RaftClientRequest request) {
   }
 
   ReadException getReadException(String op, long installSnapshot, boolean started) {
-    return installSnapshot == RaftLog.INVALID_LOG_INDEX ? null : new ReadException(getMemberId() + ": Failed to " + op
-        + " readIndex as snapshot (" + installSnapshot + ") installation is " + (started ? "started" : "in progress"));
+    return installSnapshot == RaftLog.INVALID_LOG_INDEX ? null : new ReadException(getMemberId() + ": Failed " + op
+        + " as snapshot (" + installSnapshot + ") installation is " + (started ? "started" : "in progress"));
   }
 
-  private ReadException getReadExceptionIfInstallingSnapshot() {
-    return getReadException("add", server.getSnapshotInstallationHandler().getInProgressInstallSnapshotIndex(), false);
+  private ReadException getReadExceptionIfInstallingSnapshot(String op) {
+    final long installSnapshotIndex = server.getSnapshotInstallationHandler().getInProgressInstallSnapshotIndex();
+    return getReadException(op, installSnapshotIndex, false);
   }
 
   private CompletableFuture<ReadIndexReplyProto> sendReadIndexAsync(
       ClientId clientId, ReadRequestTypeProto readRequestType) {
-    final long installSnapshot = server.getSnapshotInstallationHandler().getInProgressInstallSnapshotIndex();
-    if (installSnapshot != RaftLog.INVALID_LOG_INDEX) {
-      return JavaUtils.completeExceptionally(getReadException("get", installSnapshot, false));
+    final ReadException readException = getReadExceptionIfInstallingSnapshot("to send readIndex request");
+    if (readException != null) {
+      return JavaUtils.completeExceptionally(readException);
     }
     final RaftPeerId leaderId = server.getInfo().getLeaderId();
     if (leaderId == null) {
@@ -154,7 +159,7 @@ CompletableFuture<Message> readOnlyAsync(ClientId clientId, ReadRequestTypeProto
       final LeaderStateImpl leader = server.getRole().getLeaderState().orElse(null);
       return getReadIndex(clientId, readRequestType, leader)
           .thenCompose(readIndex -> state.getReadRequests().waitToAdvance(readIndex,
-              this::getReadExceptionIfInstallingSnapshot))
+              () -> getReadExceptionIfInstallingSnapshot("local read")))
           .thenCompose(readIndex -> stateMachine.query(queryMessage));
     } else {
       throw new IllegalStateException("Unexpected read option: " + readOption);
@@ -177,7 +182,7 @@ CompletableFuture<RaftClientReply> readAsync(RaftClientRequest request) {
 
       return replyFuture
           .thenCompose(readIndex -> state.getReadRequests().waitToAdvance(readIndex,
-              this::getReadExceptionIfInstallingSnapshot))
+              () -> getReadExceptionIfInstallingSnapshot("read")))
           .thenCompose(readIndex -> isDummyRead(request)
               ? CompletableFuture.completedFuture(server.newSuccessReply(request)) : queryStateMachine(request))
           .exceptionally(e -> readException2Reply(request, e));

ratis-server/src/main/java/org/apache/ratis/server/impl/SnapshotInstallationHandler.java

@@ -43,6 +43,7 @@
 import org.apache.ratis.util.CodeInjectionForTesting;
 import org.apache.ratis.util.LifeCycle;
 import org.apache.ratis.util.Preconditions;
+import org.apache.ratis.thirdparty.com.google.common.annotations.VisibleForTesting;
 import org.apache.ratis.util.ProtoUtils;
 import org.apache.ratis.util.Timestamp;
 import org.slf4j.Logger;
@@ -106,6 +107,11 @@ long getInProgressInstallSnapshotIndex() {
     return inProgressInstallSnapshotIndex.get();
   }
 
+  @VisibleForTesting
+  void setInProgressInstallSnapshotIndex(long index) {
+    inProgressInstallSnapshotIndex.set(index);
+  }
+
   InstallSnapshotReplyProto installSnapshot(InstallSnapshotRequestProto request) throws IOException {
     BatchLogger.print(BatchLogKey.INSTALL_SNAPSHOT_REQUEST, getMemberId(),
         suffix -> LOG.info("{}: receive installSnapshot: {} {}",
@@ -234,10 +240,19 @@ private CompletableFuture<InstallSnapshotReplyProto> checkAndInstallSnapshot(Ins
         // update the committed index
         // re-load the state machine if this is the last chunk
         if (snapshotChunkRequest.getDone()) {
-          stateMachine.pause(); // pause the SM right before publishing the snapshot atomically
-          snapshotManager.finalizeSnapshot(request);
-          state.reloadStateMachine(lastIncluded);
-          chunk0CallId.set(-1);
+          final long i = inProgressInstallSnapshotIndex.getAndUpdate(
+              previous -> previous == INVALID_LOG_INDEX ? lastIncludedIndex : previous);
+          if (i != INVALID_LOG_INDEX) {
+            throw new IOException("Snapshot (index " + i + ") installation is in progress");
+          }
+          try {
+            stateMachine.pause(); // pause the SM right before publishing the snapshot atomically
+            snapshotManager.finalizeSnapshot(request);
+            state.reloadStateMachine(lastIncluded);
+            chunk0CallId.set(-1);
+          } finally {
+            inProgressInstallSnapshotIndex.set(INVALID_LOG_INDEX);
+          }
         }
       } finally {
         server.updateLastRpcTime(UpdateType.INSTALL_SNAPSHOT_COMPLETE);
@@ -289,7 +304,7 @@ private CompletableFuture<InstallSnapshotReplyProto> notifyStateMachineToInstall
           return future.thenApply(dummy -> reply);
         }
         server.getState().getReadRequests().fail(
-            server.getReadRequestHandler().getReadException("wait for", firstAvailableLogIndex, true));
+            server.getReadRequestHandler().getReadException("to wait for readIndex", firstAvailableLogIndex, true));
 
         final RaftPeerProto leaderProto;
         if (!request.hasLastRaftConfigurationLogEntryProto()) {

ratis-server/src/test/java/org/apache/ratis/server/impl/TestStaleReadDuringSnapshotInstallation.java

@@ -0,0 +1,118 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.ratis.server.impl;
+
+import org.apache.ratis.RaftTestUtil;
+import org.apache.ratis.client.RaftClient;
+import org.apache.ratis.conf.RaftProperties;
+import org.apache.ratis.protocol.Message;
+import org.apache.ratis.protocol.RaftClientReply;
+import org.apache.ratis.protocol.exceptions.ReadException;
+import org.apache.ratis.retry.RetryPolicies;
+import org.apache.ratis.server.RaftServer;
+import org.apache.ratis.server.raftlog.RaftLog;
+import org.apache.ratis.server.simulation.MiniRaftClusterWithSimulatedRpc;
+import org.apache.ratis.statemachine.StateMachine;
+import org.apache.ratis.statemachine.impl.SimpleStateMachine4Testing;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.io.IOException;
+
+/**
+ * Test stale reads being rejected while a snapshot installation is in progress: the follower
+ * serving the stale read is exactly the node whose state machine is paused/reinitialized by the
+ * installation, so the read must be rejected with a retriable error instead of querying a paused
+ * state machine (see {@link ReadRequestHandler#staleReadAsync}).
+ */
+public class TestStaleReadDuringSnapshotInstallation {
+  private static final int NUM_SERVERS = 3;
+
+  private final MiniRaftClusterWithSimulatedRpc cluster = MiniRaftClusterWithSimulatedRpc
+      .FACTORY.newCluster(NUM_SERVERS, getProperties());
+
+  private static RaftProperties getProperties() {
+    final RaftProperties properties = new RaftProperties();
+    properties.setClass(MiniRaftCluster.STATEMACHINE_CLASS_KEY,
+        SimpleStateMachine4Testing.class, StateMachine.class);
+    return properties;
+  }
+
+  @BeforeEach
+  public void setup() throws IOException {
+    cluster.start();
+  }
+
+  @AfterEach
+  public void tearDown() {
+    cluster.shutdown();
+  }
+
+  @Test
+  @Timeout(120)
+  public void testStaleReadRejectedWhileSnapshotInstallationInProgress() throws Exception {
+    RaftTestUtil.waitForLeader(cluster);
+
+    try (final RaftClient client = cluster.createClient()) {
+      for (int i = 0; i < 5; i++) {
+        Assertions.assertTrue(client.io().send(Message.valueOf("m" + i)).isSuccess());
+      }
+
+      final RaftServer.Division follower = cluster.getFollowers().get(0);
+      final SnapshotInstallationHandler handler =
+          ((RaftServerImpl) follower).getSnapshotInstallationHandler();
+      final Message message = Message.valueOf("m4");
+
+      // send() only guarantees that the leader has committed; the follower applies asynchronously.
+      // Wait until the follower has applied m4, otherwise a stale read may hit the state machine
+      // before the queried entry is applied (see SimpleStateMachine4Testing#query).
+      final long m4Index = cluster.getLeader().getRaftLog().getLastCommittedIndex();
+      for (int i = 0; follower.getInfo().getLastAppliedIndex() < m4Index; i++) {
+        Assertions.assertTrue(i < 100, () -> follower.getId() + ": fails to catch up, lastAppliedIndex="
+            + follower.getInfo().getLastAppliedIndex() + " < " + m4Index);
+        Thread.sleep(100);
+      }
+
+      // Regression guard: a stale read on a healthy follower still succeeds.
+      final RaftClientReply before = client.io().sendStaleRead(message, m4Index, follower.getId());
+      Assertions.assertTrue(before.isSuccess());
+
+      // Simulate a snapshot installation in progress on the follower: the gate must reject
+      // the stale read with a retriable ReadException instead of querying the paused state
+      // machine. Use a no-retry client so the exception surfaces immediately (the default
+      // test policy would keep retrying it -- see BlockingImpl#sendRequestWithRetry).
+      handler.setInProgressInstallSnapshotIndex(1);
+      try (final RaftClient noRetryClient = cluster.createClient(follower.getId(), RetryPolicies.noRetry())) {
+        final IOException e = Assertions.assertThrows(IOException.class,
+            () -> noRetryClient.io().sendStaleRead(message, m4Index, follower.getId()));
+        Assertions.assertTrue(e instanceof ReadException, () -> "Unexpected exception: " + e);
+        Assertions.assertTrue(e.getMessage().contains("stale read"), () -> e.getMessage());
+        Assertions.assertTrue(e.getMessage().contains("in progress"), () -> e.getMessage());
+      } finally {
+        handler.setInProgressInstallSnapshotIndex(RaftLog.INVALID_LOG_INDEX);
+      }
+
+      // After the (simulated) installation completed: stale reads work again.
+      final RaftClientReply after = client.io().sendStaleRead(message, m4Index, follower.getId());
+      Assertions.assertTrue(after.isSuccess());
+    }
+  }
+}