RATIS-2603. Add DataStreamApi.Resolver for group-independent read-only requests

ratis-netty/src/main/java/org/apache/ratis/netty/server/NettyServerStreamRpc.java

@@ -36,6 +36,7 @@
 import org.apache.ratis.protocol.RaftClientRequest;
 import org.apache.ratis.protocol.RaftPeer;
 import org.apache.ratis.security.TlsConf;
+import org.apache.ratis.server.DataStreamReadResolver;
 import org.apache.ratis.server.DataStreamServerRpc;
 import org.apache.ratis.server.RaftServer;
 import org.apache.ratis.server.RaftServerConfigKeys;
@@ -165,7 +166,7 @@ public NettyServerStreamRpc(RaftServer server, Parameters parameters) {
     this.name = server.getId() + "-" + JavaUtils.getClassSimpleName(getClass());
     this.metrics = new NettyServerStreamRpcMetrics(this.name);
     this.requests = new DataStreamManagement(server, metrics);
-    this.reads = new ReadStreamManagement(server);
+    this.reads = new ReadStreamManagement(server, DataStreamReadResolver.get(parameters));
 
     final RaftProperties properties = server.getProperties();
 

ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java

@@ -29,8 +29,11 @@
 import org.apache.ratis.protocol.RaftClientReply;
 import org.apache.ratis.protocol.RaftClientRequest;
 import org.apache.ratis.protocol.exceptions.AlreadyClosedException;
+import org.apache.ratis.protocol.exceptions.DataStreamException;
+import org.apache.ratis.server.DataStreamReadResolver;
 import org.apache.ratis.server.RaftServer;
 import org.apache.ratis.server.RaftServerConfigKeys;
+import org.apache.ratis.statemachine.StateMachine;
 import org.apache.ratis.thirdparty.com.google.protobuf.InvalidProtocolBufferException;
 import org.apache.ratis.thirdparty.io.netty.channel.ChannelFuture;
 import org.apache.ratis.thirdparty.io.netty.channel.ChannelHandlerContext;
@@ -91,8 +94,16 @@ public boolean isOpen() {
 
     @Override
     public synchronized void close() throws IOException {
+      close(terminalReply);
+    }
+
+    private synchronized void fail(RaftClientReply reply) throws IOException {
+      close(newReadStreamTerminalReply(clientId, streamId, reply));
+    }
+
+    private void close(DataStreamReplyByteBuffer reply) throws IOException {
       if (closed.complete(null)) {
-        writeAndFlush(terminalReply, 0);
+        writeAndFlush(reply, 0);
       }
     }
 
@@ -137,11 +148,17 @@ private synchronized DataStreamReplyByteBuffer newReply(ByteBuffer buffer) {
   }
 
   private final RaftServer server;
+  private final DataStreamReadResolver readResolver;
   private final String name;
   private final ExecutorService requestExecutor;
 
   ReadStreamManagement(RaftServer server) {
+    this(server, null);
+  }
+
+  ReadStreamManagement(RaftServer server, DataStreamReadResolver readResolver) {
     this.server = server;
+    this.readResolver = readResolver;
     this.name = server.getId() + "-" + JavaUtils.getClassSimpleName(getClass());
 
     final RaftProperties properties = server.getProperties();
@@ -182,6 +199,22 @@ private boolean processImpl(DataStreamRequestByteBuf requestBuf, ChannelHandlerC
       return false;
     }
 
+    if (readResolver != null) {
+      final StateMachine.DataApi dataApi;
+      try {
+        dataApi = readResolver.resolve(request);
+      } catch (Throwable t) {
+        replyDataStreamException(server, t, request, requestBuf, ctx);
+        return true;
+      }
+      if (dataApi != null) {
+        final RaftClientReply reply = RaftClientReply.newBuilder().setRequest(request).setSuccess().build();
+        final long streamId = requestBuf.getStreamId();
+        requestExecutor.execute(() -> queryResolved(dataApi, request, streamId, ctx, reply));
+        return true;
+      }
+    }
+
     final RaftServer.Division division;
     try {
       division = server.getDivision(request.getRaftGroupId());
@@ -209,16 +242,42 @@ private boolean processImpl(DataStreamRequestByteBuf requestBuf, ChannelHandlerC
         return;
       }
 
-      final ReadStream stream = new ReadStream(request, requestBuf.getStreamId(), ctx, readCheckReply);
-      try {
-        division.getStateMachine().data().query(request.getMessage(), stream);
-      } catch (Throwable t) {
-        LOG.error("{}: Failed read-only data stream query for {}", this, request, t);
-      }
+      query(division.getStateMachine().data(), request, requestBuf.getStreamId(), ctx, readCheckReply);
     }, requestExecutor);
     return true;
   }
 
+  private void queryResolved(StateMachine.DataApi dataApi, RaftClientRequest request, long streamId,
+      ChannelHandlerContext ctx, RaftClientReply terminalReply) {
+    final ReadStream stream = new ReadStream(request, streamId, ctx, terminalReply);
+    try {
+      dataApi.query(request.getMessage(), stream);
+    } catch (Throwable t) {
+      LOG.error("{}: Failed resolved read-only data stream query for {}", this, request, t);
+      if (stream.isOpen()) {
+        final RaftClientReply failure = RaftClientReply.newBuilder()
+            .setRequest(request)
+            .setException(new DataStreamException(server.getId(), t))
+            .build();
+        try {
+          stream.fail(failure);
+        } catch (IOException e) {
+          LOG.warn("{}: Failed to send resolved read-only data stream error for {}", this, request, e);
+        }
+      }
+    }
+  }
+
+  private void query(StateMachine.DataApi dataApi, RaftClientRequest request, long streamId,
+      ChannelHandlerContext ctx, RaftClientReply terminalReply) {
+    final ReadStream stream = new ReadStream(request, streamId, ctx, terminalReply);
+    try {
+      dataApi.query(request.getMessage(), stream);
+    } catch (Throwable t) {
+      LOG.error("{}: Failed read-only data stream query for {}", this, request, t);
+    }
+  }
+
   private static RaftClientRequest newDummyReadRequest(RaftClientRequest request) {
     final RaftProtos.ReadRequestTypeProto original = request.getType().getRead();
     return RaftClientRequest.newBuilder()

ratis-server-api/src/main/java/org/apache/ratis/server/DataStreamReadResolver.java

@@ -0,0 +1,47 @@
+/*
+ * 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;
+
+import org.apache.ratis.conf.Parameters;
+import org.apache.ratis.protocol.RaftClientRequest;
+import org.apache.ratis.statemachine.StateMachine;
+
+import java.io.IOException;
+
+/** Resolves read-only data stream requests independently of a Raft division. */
+@FunctionalInterface
+public interface DataStreamReadResolver {
+  String PARAMETER_KEY = DataStreamReadResolver.class.getName();
+
+  /** Set the resolver in the given server parameters. */
+  static void set(Parameters parameters, DataStreamReadResolver resolver) {
+    parameters.put(PARAMETER_KEY, resolver, DataStreamReadResolver.class);
+  }
+
+  /** Get the resolver from the given server parameters, if configured. */
+  static DataStreamReadResolver get(Parameters parameters) {
+    return parameters != null
+        ? parameters.get(PARAMETER_KEY, DataStreamReadResolver.class)
+        : null;
+  }
+
+  /**
+   * @return the data API handling this request, or null to use the Raft division associated with the request.
+   */
+  StateMachine.DataApi resolve(RaftClientRequest request) throws IOException;
+}

ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java

@@ -127,6 +127,175 @@ public void query(Message request, WritableByteChannel stream) {
     }
   }
 
+  @Test
+  void readResolverHandlesRequestWithoutDivision() throws Exception {
+    final RaftPeerId serverId = RaftPeerId.valueOf("s1");
+    final ClientId clientId = ClientId.randomId();
+    final RaftGroupId groupId = RaftGroupId.randomId();
+    final ByteString query = ByteString.copyFromUtf8("query");
+    final ByteString response = ByteString.copyFromUtf8("response");
+    final AtomicBoolean readCheckSubmitted = new AtomicBoolean();
+    final AtomicReference<RaftClientRequest> resolvedRequest = new AtomicReference<>();
+    final AtomicReference<Message> messageRef = new AtomicReference<>();
+    final AtomicReference<WritableByteChannel> streamRef = new AtomicReference<>();
+
+    final DataApi dataApi = new DataApi() {
+      @Override
+      public void query(Message request, WritableByteChannel stream) {
+        messageRef.set(request);
+        streamRef.set(stream);
+      }
+    };
+    final RaftServer server = newRaftServer(serverId, new RaftProperties(), null, null, request -> {
+      readCheckSubmitted.set(true);
+      return successReply(request);
+    });
+    final ReadStreamManagement management = new ReadStreamManagement(server, request -> {
+      resolvedRequest.set(request);
+      return dataApi;
+    });
+    final EmbeddedChannel embeddedChannel = new EmbeddedChannel(new ChannelInboundHandlerAdapter());
+    final ReadOnlyRequest readOnlyRequest = newReadOnlyRequest(clientId, serverId, groupId, 1L, query);
+
+    try {
+      assertTrue(management.process(readOnlyRequest.request, embeddedChannel.pipeline().firstContext()));
+      assertEquals(0, readOnlyRequest.headerBuf.refCnt());
+
+      JavaUtils.attempt(() -> assertNotNull(streamRef.get()), 10,
+          TimeDuration.valueOf(100, TimeUnit.MILLISECONDS), "resolved read-only stream", null);
+      streamRef.get().write(response.asReadOnlyByteBuffer());
+      streamRef.get().close();
+
+      final List<DataStreamReply> replies = new ArrayList<>();
+      JavaUtils.attempt(() -> {
+        for (Object outbound; (outbound = embeddedChannel.readOutbound()) != null;) {
+          replies.add((DataStreamReply) outbound);
+        }
+        assertEquals(2, replies.size());
+      }, 10, TimeDuration.valueOf(100, TimeUnit.MILLISECONDS), "resolved read-only replies", null);
+
+      assertEquals(query, resolvedRequest.get().getMessage().getContent());
+      assertEquals(query, messageRef.get().getContent());
+      assertFalse(readCheckSubmitted.get(), "a resolved request should bypass the Raft read check");
+      assertSuccessReply(Type.STREAM_DATA, response.size(), replies.get(0));
+      assertSuccessReply(Type.STREAM_HEADER, 0, replies.get(1));
+      assertTrue(ClientProtoUtils.getRaftClientReply(replies.get(1)).isSuccess());
+    } finally {
+      embeddedChannel.finishAndReleaseAll();
+      management.shutdown();
+    }
+  }
+
+  @Test
+  void declinedReadResolverUsesDivisionAndReadCheck() throws Exception {
+    final RaftPeerId serverId = RaftPeerId.valueOf("s1");
+    final ClientId clientId = ClientId.randomId();
+    final RaftGroupId groupId = RaftGroupId.randomId();
+    final AtomicBoolean resolverCalled = new AtomicBoolean();
+    final AtomicReference<RaftClientRequest> submittedReadCheck = new AtomicReference<>();
+    final AtomicReference<WritableByteChannel> streamRef = new AtomicReference<>();
+
+    final DataApi dataApi = new DataApi() {
+      @Override
+      public void query(Message request, WritableByteChannel stream) {
+        streamRef.set(stream);
+      }
+    };
+    final StateMachine stateMachine = new BaseStateMachine() {
+      @Override
+      public DataApi data() {
+        return dataApi;
+      }
+    };
+    final RaftServer server = newRaftServer(serverId, new RaftProperties(), groupId, newDivision(stateMachine),
+        request -> {
+          submittedReadCheck.set(request);
+          return successReply(request);
+        });
+    final ReadStreamManagement management = new ReadStreamManagement(server, request -> {
+      resolverCalled.set(true);
+      return null;
+    });
+    final EmbeddedChannel embeddedChannel = new EmbeddedChannel(new ChannelInboundHandlerAdapter());
+    final ReadOnlyRequest readOnlyRequest = newReadOnlyRequest(
+        clientId, serverId, groupId, 1L, ByteString.copyFromUtf8("query"));
+
+    try {
+      assertTrue(management.process(readOnlyRequest.request, embeddedChannel.pipeline().firstContext()));
+      JavaUtils.attempt(() -> assertNotNull(streamRef.get()), 10,
+          TimeDuration.valueOf(100, TimeUnit.MILLISECONDS), "declined read-only stream", null);
+
+      assertTrue(resolverCalled.get());
+      assertEquals(OrderedAsync.DUMMY.getContent(), submittedReadCheck.get().getMessage().getContent());
+      streamRef.get().close();
+    } finally {
+      embeddedChannel.finishAndReleaseAll();
+      management.shutdown();
+    }
+  }
+
+  @Test
+  void readResolverExceptionReturnsFailure() throws Exception {
+    final RaftPeerId serverId = RaftPeerId.valueOf("s1");
+    final ClientId clientId = ClientId.randomId();
+    final RaftGroupId groupId = RaftGroupId.randomId();
+    final RaftServer server = newRaftServer(serverId, new RaftProperties());
+    final ReadStreamManagement management = new ReadStreamManagement(server, request -> {
+      throw new IOException("resolve failed");
+    });
+    final EmbeddedChannel embeddedChannel = new EmbeddedChannel(new ChannelInboundHandlerAdapter());
+    final ReadOnlyRequest readOnlyRequest = newReadOnlyRequest(
+        clientId, serverId, groupId, 1L, ByteString.copyFromUtf8("query"));
+
+    try {
+      assertTrue(management.process(readOnlyRequest.request, embeddedChannel.pipeline().firstContext()));
+      final DataStreamReply reply = embeddedChannel.readOutbound();
+      assertNotNull(reply);
+      assertEquals(Type.STREAM_HEADER, reply.getType());
+      assertFalse(reply.isSuccess());
+      assertNotNull(ClientProtoUtils.getRaftClientReply(reply).getDataStreamException());
+    } finally {
+      embeddedChannel.finishAndReleaseAll();
+      management.shutdown();
+    }
+  }
+
+  @Test
+  void resolvedQueryExceptionReturnsFailure() throws Exception {
+    final RaftPeerId serverId = RaftPeerId.valueOf("s1");
+    final ClientId clientId = ClientId.randomId();
+    final RaftGroupId groupId = RaftGroupId.randomId();
+    final DataApi dataApi = new DataApi() {
+      @Override
+      public void query(Message request, WritableByteChannel stream) {
+        throw new IllegalStateException("query failed");
+      }
+    };
+    final RaftServer server = newRaftServer(serverId, new RaftProperties());
+    final ReadStreamManagement management = new ReadStreamManagement(server, request -> dataApi);
+    final EmbeddedChannel embeddedChannel = new EmbeddedChannel(new ChannelInboundHandlerAdapter());
+    final ReadOnlyRequest readOnlyRequest = newReadOnlyRequest(
+        clientId, serverId, groupId, 1L, ByteString.copyFromUtf8("query"));
+
+    try {
+      assertTrue(management.process(readOnlyRequest.request, embeddedChannel.pipeline().firstContext()));
+      assertEquals(0, readOnlyRequest.headerBuf.refCnt());
+      final AtomicReference<DataStreamReply> replyRef = new AtomicReference<>();
+      JavaUtils.attempt(() -> {
+        replyRef.set(embeddedChannel.readOutbound());
+        assertNotNull(replyRef.get());
+      }, 10, TimeDuration.valueOf(100, TimeUnit.MILLISECONDS), "resolved query failure", null);
+
+      final DataStreamReply reply = replyRef.get();
+      assertEquals(Type.STREAM_HEADER, reply.getType());
+      assertFalse(reply.isSuccess());
+      assertNotNull(ClientProtoUtils.getRaftClientReply(reply).getDataStreamException());
+    } finally {
+      embeddedChannel.finishAndReleaseAll();
+      management.shutdown();
+    }
+  }
+
   @Test
   void readOnlyRequestWaitsForLinearizableCheck() throws Exception {
     final RaftPeerId serverId = RaftPeerId.valueOf("s1");