RATIS-2400. Support timeout and interrupt handling in GrpcClientRpc.

Author: slfan1989

ratis-grpc/src/main/java/org/apache/ratis/grpc/client/GrpcClientRpc.java

@@ -17,13 +17,16 @@
  */
 package org.apache.ratis.grpc.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;
 import org.apache.ratis.grpc.GrpcConfigKeys;
 import org.apache.ratis.grpc.GrpcUtil;
 import org.apache.ratis.protocol.*;
 import org.apache.ratis.protocol.exceptions.AlreadyClosedException;
+import org.apache.ratis.protocol.exceptions.TimeoutIOException;
+import org.apache.ratis.thirdparty.io.grpc.Status;
 import org.apache.ratis.thirdparty.io.grpc.StatusRuntimeException;
 import org.apache.ratis.thirdparty.io.grpc.stub.StreamObserver;
 import org.apache.ratis.proto.RaftProtos.GroupInfoRequestProto;
@@ -39,26 +42,34 @@
 import org.apache.ratis.util.IOUtils;
 import org.apache.ratis.util.JavaUtils;
 import org.apache.ratis.util.PeerProxyMap;
+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.io.InterruptedIOException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 
 public class GrpcClientRpc extends RaftClientRpcWithProxy<GrpcClientProtocolClient> {
   public static final Logger LOG = LoggerFactory.getLogger(GrpcClientRpc.class);
 
   private final ClientId clientId;
   private final int maxMessageSize;
+  private final TimeDuration requestTimeoutDuration;
+  private final TimeDuration watchRequestTimeoutDuration;
+  private final TimeoutExecutor scheduler = TimeoutExecutor.getInstance();
 
   public GrpcClientRpc(ClientId clientId, RaftProperties properties,
       SslContext adminSslContext, SslContext clientSslContext) {
     super(new PeerProxyMap<>(clientId.toString(),
         p -> new GrpcClientProtocolClient(clientId, p, properties, adminSslContext, clientSslContext)));
     this.clientId = clientId;
     this.maxMessageSize = GrpcConfigKeys.messageSizeMax(properties, LOG::debug).getSizeInt();
+    this.requestTimeoutDuration = RaftClientConfigKeys.Rpc.requestTimeout(properties);
+    this.watchRequestTimeoutDuration = RaftClientConfigKeys.Rpc.watchRequestTimeout(properties);
   }
 
   @Override
@@ -122,11 +133,11 @@ public RaftClientReply sendRequest(RaftClientRequest request)
       return ClientProtoUtils.toRaftClientReply(proxy.leaderElectionManagement(proto));
     } else {
       final CompletableFuture<RaftClientReply> f = sendRequest(request, proxy);
-      // TODO: timeout support
       try {
         return f.get();
       } catch (InterruptedException e) {
         Thread.currentThread().interrupt();
+        f.cancel(true);
         throw new InterruptedIOException(
             "Interrupted while waiting for response of request " + request);
       } catch (ExecutionException e) {
@@ -143,6 +154,7 @@ private CompletableFuture<RaftClientReply> sendRequest(
     final RaftClientRequestProto requestProto =
         toRaftClientRequestProto(request);
     final CompletableFuture<RaftClientReplyProto> replyFuture = new CompletableFuture<>();
+    final TimeDuration timeout = getTimeoutDuration(request);
     // create a new grpc stream for each non-async call.
     final StreamObserver<RaftClientRequestProto> requestObserver =
         proxy.unorderedWithTimeout(new StreamObserver<RaftClientReplyProto>() {
@@ -164,12 +176,35 @@ public void onCompleted() {
             }
           }
         });
+    replyFuture.whenComplete((reply, exception) -> {
+      if (replyFuture.isCancelled()) {
+        requestObserver.onError(Status.CANCELLED
+            .withDescription(clientId + ": request #" + request.getCallId() + " cancelled")
+            .asRuntimeException());
+      }
+    });
+    scheduler.onTimeout(timeout, () -> {
+          if (replyFuture.completeExceptionally(new TimeoutIOException(
+              clientId + ": request #" + request.getCallId() + " timeout " + timeout))) {
+            requestObserver.onError(Status.DEADLINE_EXCEEDED
+                .withDescription(clientId + ": request #" + request.getCallId() + " timeout " + timeout)
+                .asRuntimeException());
+          }
+        }, LOG, () -> "Timeout check failed for client request #" + request.getCallId());
     requestObserver.onNext(requestProto);
     requestObserver.onCompleted();
 
     return replyFuture.thenApply(ClientProtoUtils::toRaftClientReply);
   }
 
+  private TimeDuration getTimeoutDuration(RaftClientRequest request) {
+    final long timeoutMs = request.getTimeoutMs();
+    if (timeoutMs > 0) {
+      return TimeDuration.valueOf(timeoutMs, TimeUnit.MILLISECONDS);
+    }
+    return request.is(RaftClientRequestProto.TypeCase.WATCH) ? watchRequestTimeoutDuration : requestTimeoutDuration;
+  }
+
   private RaftClientRequestProto toRaftClientRequestProto(RaftClientRequest request) throws IOException {
     final RaftClientRequestProto proto = ClientProtoUtils.toRaftClientRequestProto(request);
     if (proto.getSerializedSize() > maxMessageSize) {

ratis-test/src/test/java/org/apache/ratis/grpc/TestRaftServerWithGrpc.java

@@ -69,6 +69,7 @@
 import javax.net.ssl.KeyManager;
 import javax.net.ssl.TrustManager;
 import java.io.IOException;
+import java.io.InterruptedIOException;
 import java.nio.channels.OverlappingFileLockException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -81,6 +82,7 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class TestRaftServerWithGrpc extends BaseTest implements MiniRaftClusterWithGrpc.FactoryGet {
   {
@@ -238,6 +240,58 @@ public void testRaftClientMetrics(Boolean separateHeartbeat) throws Exception {
     runWithNewCluster(3, this::testRaftClientRequestMetrics);
   }
 
+  @ParameterizedTest
+  @MethodSource("data")
+  public void testGrpcClientRpcSyncTimeout(Boolean separateHeartbeat) throws Exception {
+    GrpcConfigKeys.Server.setHeartbeatChannel(getProperties(), separateHeartbeat);
+    runWithNewCluster(3, cluster -> {
+      final RaftPeerId leaderId = RaftTestUtil.waitForLeader(cluster).getId();
+      try (RaftClient client = cluster.createClient(leaderId, RetryPolicies.noRetry())) {
+        final SimpleStateMachine4Testing stateMachine = SimpleStateMachine4Testing.get(cluster.getLeader());
+        stateMachine.blockStartTransaction();
+        try {
+          Assertions.assertThrows(TimeoutIOException.class,
+              () -> client.io().send(new SimpleMessage("sync-timeout")));
+        } finally {
+          stateMachine.unblockStartTransaction();
+        }
+      }
+    });
+  }
+
+  @ParameterizedTest
+  @MethodSource("data")
+  public void testGrpcClientRpcSyncCancelOnInterrupt(Boolean separateHeartbeat) throws Exception {
+    RaftClientConfigKeys.Rpc.setRequestTimeout(getProperties(), TimeDuration.valueOf(10, TimeUnit.SECONDS));
+    GrpcConfigKeys.Server.setHeartbeatChannel(getProperties(), separateHeartbeat);
+    runWithNewCluster(3, cluster -> {
+      final RaftPeerId leaderId = RaftTestUtil.waitForLeader(cluster).getId();
+      try (RaftClient client = cluster.createClient(leaderId, RetryPolicies.noRetry())) {
+        final SimpleStateMachine4Testing stateMachine = SimpleStateMachine4Testing.get(cluster.getLeader());
+        stateMachine.blockStartTransaction();
+        try {
+          final AtomicReference<Throwable> error = new AtomicReference<>();
+          final Thread t = new Thread(() -> {
+            try {
+              client.io().send(new SimpleMessage("sync-cancel"));
+            } catch (Throwable e) {
+              error.set(e);
+            }
+          });
+          t.start();
+          Thread.sleep(200);
+          t.interrupt();
+          t.join(5000);
+          Assertions.assertFalse(t.isAlive(), "request thread should exit after interrupt");
+          Assertions.assertTrue(error.get() instanceof InterruptedIOException,
+              "expected InterruptedIOException but got " + error.get());
+        } finally {
+          stateMachine.unblockStartTransaction();
+        }
+      }
+    });
+  }
+
   @ParameterizedTest
   @MethodSource("data")
   public void testRaftServerMetrics(Boolean separateHeartbeat) throws Exception {