RATIS-2294. Add Timeout.

Author: slfan1989

ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientRpc.java

@@ -17,6 +17,7 @@
  */
 package org.apache.ratis.netty.client;
 
+import org.apache.ratis.client.RaftClientConfigKeys;
 import org.apache.ratis.client.impl.ClientProtoUtils;
 import org.apache.ratis.client.impl.RaftClientRpcWithProxy;
 import org.apache.ratis.conf.RaftProperties;
@@ -28,36 +29,97 @@
 import org.apache.ratis.proto.RaftProtos.GroupManagementRequestProto;
 import org.apache.ratis.proto.RaftProtos.SetConfigurationRequestProto;
 import org.apache.ratis.proto.netty.NettyProtos.RaftNettyServerRequestProto;
+import org.apache.ratis.protocol.exceptions.LeaderNotReadyException;
+import org.apache.ratis.protocol.exceptions.NotLeaderException;
+import org.apache.ratis.protocol.exceptions.TimeoutIOException;
 import org.apache.ratis.util.JavaUtils;
+import org.apache.ratis.util.TimeDuration;
+import org.apache.ratis.util.TimeoutExecutor;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
 
 public class NettyClientRpc extends RaftClientRpcWithProxy<NettyRpcProxy> {
+
+  public static final Logger LOG = LoggerFactory.getLogger(NettyClientRpc.class);
+
+  private Supplier<String> name;
+  private ClientId cId;
+  private final TimeDuration requestTimeoutDuration;
+  private final TimeoutExecutor scheduler = TimeoutExecutor.getInstance();
+
   public NettyClientRpc(ClientId clientId, RaftProperties properties) {
     super(new NettyRpcProxy.PeerMap(clientId.toString(), properties));
+    this.cId = clientId;
+    this.requestTimeoutDuration = RaftClientConfigKeys.Rpc.requestTimeout(properties);
   }
 
   @Override
   public CompletableFuture<RaftClientReply> sendRequestAsync(RaftClientRequest request) {
     final RaftPeerId serverId = request.getServerId();
+    long callId = request.getCallId();
     try {
+      name = JavaUtils.memoize(() -> cId + "->" + serverId);
       final NettyRpcProxy proxy = getProxies().getProxy(serverId);
       final RaftNettyServerRequestProto serverRequestProto = buildRequestProto(request);
-      return proxy.sendAsync(serverRequestProto).thenApply(replyProto -> {
+      final CompletableFuture<RaftClientReply> replyFuture = new CompletableFuture<>();
+
+      proxy.sendAsync(serverRequestProto).thenApply(replyProto -> {
         if (request instanceof GroupListRequest) {
           return ClientProtoUtils.toGroupListReply(replyProto.getGroupListReply());
         } else if (request instanceof GroupInfoRequest) {
           return ClientProtoUtils.toGroupInfoReply(replyProto.getGroupInfoReply());
         } else {
           return ClientProtoUtils.toRaftClientReply(replyProto.getRaftClientReply());
         }
+      }).thenCompose(raftReply -> {
+        final NotLeaderException nle = raftReply.getNotLeaderException();
+        if (nle != null) {
+          return failedFuture(nle);
+        }
+        final LeaderNotReadyException lnre = raftReply.getLeaderNotReadyException();
+        if (lnre != null) {
+          return failedFuture(lnre);
+        }
+        return CompletableFuture.completedFuture(raftReply);
+      }).whenComplete((raftReply, ex) -> {
+        if (ex != null) {
+          replyFuture.completeExceptionally(ex);
+        } else {
+          replyFuture.complete(raftReply);
+        }
       });
+
+      scheduler.onTimeout(requestTimeoutDuration,
+        () -> {
+          if (!replyFuture.isDone()) {
+            TimeoutIOException timeout = new TimeoutIOException(
+                getName()+ " request #" + callId + " timeout " +
+                requestTimeoutDuration.getDuration());
+            replyFuture.completeExceptionally(timeout);
+          }
+        }, LOG, () -> "Timeout check for client request #" + callId
+      );
+
+      return replyFuture;
     } catch (Throwable e) {
       return JavaUtils.completeExceptionally(e);
     }
   }
 
+  private String getName() {
+    return name.get();
+  }
+
+  private <T> CompletableFuture<T> failedFuture(Throwable ex) {
+    CompletableFuture<T> future = new CompletableFuture<>();
+    future.completeExceptionally(ex);
+    return future;
+  }
+
   @Override
   public RaftClientReply sendRequest(RaftClientRequest request) throws IOException {
     final RaftPeerId serverId = request.getServerId();

ratis-server/src/test/java/org/apache/ratis/RaftAsyncTests.java

@@ -52,6 +52,7 @@
 import org.apache.ratis.util.function.CheckedRunnable;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
 import org.slf4j.event.Level;
 
 import java.io.IOException;
@@ -69,6 +70,7 @@
 
 import static org.apache.ratis.RaftTestUtil.waitForLeader;
 
+@Timeout(100)
 public abstract class RaftAsyncTests<CLUSTER extends MiniRaftCluster> extends BaseTest
     implements MiniRaftCluster.Factory.Get<CLUSTER> {
   {
@@ -417,13 +419,13 @@ void runTestAppendEntriesTimeout(CLUSTER cluster) throws Exception {
     LOG.info("Running testAppendEntriesTimeout");
     final TimeDuration oldExpiryTime = RaftServerConfigKeys.RetryCache.expiryTime(getProperties());
     RaftServerConfigKeys.RetryCache.setExpiryTime(getProperties(), TimeDuration.valueOf(20, TimeUnit.SECONDS));
-    waitForLeader(cluster);
+    RaftPeerId id = waitForLeader(cluster).getId();
     long time = System.currentTimeMillis();
     long waitTime = 5000;
     try (final RaftClient client = cluster.createClient()) {
       // block append requests
       cluster.getServerAliveStream()
-          .filter(impl -> !impl.getInfo().isLeader())
+          .filter(impl -> !impl.getInfo().isLeader() && impl.getPeer().getId() != id)
           .map(SimpleStateMachine4Testing::get)
           .forEach(SimpleStateMachine4Testing::blockWriteStateMachineData);
 
@@ -433,7 +435,7 @@ void runTestAppendEntriesTimeout(CLUSTER cluster) throws Exception {
       Assertions.assertFalse(replyFuture.isDone());
       // unblock append request.
       cluster.getServerAliveStream()
-          .filter(impl -> !impl.getInfo().isLeader())
+          .filter(impl -> !impl.getInfo().isLeader() && impl.getPeer().getId() != id)
           .map(SimpleStateMachine4Testing::get)
           .forEach(SimpleStateMachine4Testing::unblockWriteStateMachineData);