RATIS-2625. Streaming: support sending commands in Data Stream

ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamOutput.java

@@ -75,6 +75,17 @@ default CompletableFuture<DataStreamReply> writeAsync(File src, long position, l
    */
   CompletableFuture<DataStreamReply> writeAsync(FilePositionCount src, WriteOption... options);
 
+  /**
+   * Send a command asynchronously.
+   * Commands are ordered with data writes but do not advance the stream byte offset.
+   * The server state machine may handle the command via its data-stream command hook
+   * instead of writing bytes to the data channel.
+   *
+   * @param command the command payload
+   * @return a future of the reply
+   */
+  CompletableFuture<DataStreamReply> commandAsync(ByteBuffer command);
+
   /**
    * Return the future of the {@link RaftClientReply}
    * which will be received once this stream has been closed successfully.

ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java

@@ -173,6 +173,18 @@ private CompletableFuture<DataStreamReply> writeAsyncImpl(Object data, long leng
       return f;
     }
 
+    private CompletableFuture<DataStreamReply> commandAsyncImpl(Object data, long length) {
+      if (isClosed()) {
+        if (data instanceof ByteBuf) {
+          ((ByteBuf) data).release();
+        }
+        return JavaUtils.completeExceptionally(new AlreadyClosedException(
+            clientId + ": stream already closed, request=" + header));
+      }
+      return combineHeader(send(Type.STREAM_COMMAND, data, length,
+          Collections.singleton(StandardWriteOption.FLUSH)));
+    }
+
     public CompletableFuture<DataStreamReply> writeAsync(ByteBuf src, Iterable<WriteOption> options) {
       return writeAsyncImpl(src, src.readableBytes(), options);
     }
@@ -187,6 +199,15 @@ public CompletableFuture<DataStreamReply> writeAsync(FilePositionCount src, Writ
       return writeAsyncImpl(src, src.getCount(), Arrays.asList(options));
     }
 
+    @Override
+    public CompletableFuture<DataStreamReply> commandAsync(ByteBuffer src) {
+      return commandAsyncImpl(src, src.remaining());
+    }
+
+    public CompletableFuture<DataStreamReply> commandAsync(ByteBuf src) {
+      return commandAsyncImpl(src, src.readableBytes());
+    }
+
     boolean isClosed() {
       return closeFuture != null;
     }

ratis-netty/src/main/java/org/apache/ratis/netty/metrics/NettyServerStreamRpcMetrics.java

@@ -39,7 +39,8 @@ public class NettyServerStreamRpcMetrics extends RatisMetrics {
   private static final String METRICS_NUM_REQUESTS = "num_requests_%s";
 
   public enum RequestType {
-    CHANNEL_READ, HEADER, LOCAL_WRITE, REMOTE_WRITE, STATE_MACHINE_STREAM, START_TRANSACTION;
+    CHANNEL_READ, HEADER, LOCAL_WRITE, LOCAL_COMMAND, REMOTE_WRITE, REMOTE_COMMAND, STATE_MACHINE_STREAM,
+    START_TRANSACTION;
 
     private final String numRequestsString;
     private final String successCountString;

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

@@ -89,20 +89,30 @@ public class DataStreamManagement {
   static class LocalStream {
     private final CompletableFuture<DataStream> streamFuture;
     private final AtomicReference<CompletableFuture<Long>> writeFuture;
-    private final RequestMetrics metrics;
+    private final RequestMetrics writeMetrics;
+    private final RequestMetrics commandMetrics;
 
-    LocalStream(CompletableFuture<DataStream> streamFuture, RequestMetrics metrics) {
+    LocalStream(CompletableFuture<DataStream> streamFuture, RequestMetrics writeMetrics,
+        RequestMetrics commandMetrics) {
       this.streamFuture = streamFuture;
       this.writeFuture = new AtomicReference<>(streamFuture.thenApply(s -> 0L));
-      this.metrics = metrics;
+      this.writeMetrics = writeMetrics;
+      this.commandMetrics = commandMetrics;
     }
 
     CompletableFuture<Long> write(ByteBuf buf, Iterable<WriteOption> options,
                                   Executor executor) {
-      final Timekeeper.Context context = metrics.start();
+      final Timekeeper.Context context = writeMetrics.start();
       return composeAsync(writeFuture, executor,
           n -> streamFuture.thenCompose(stream -> writeToAsync(buf, options, stream, executor)
-              .whenComplete((l, e) -> metrics.stop(context, e == null))));
+              .whenComplete((l, e) -> writeMetrics.stop(context, e == null))));
+    }
+
+    CompletableFuture<Long> command(ByteBuf buf, long streamOffset, Executor executor) {
+      final Timekeeper.Context context = commandMetrics.start();
+      return composeAsync(writeFuture, executor,
+          n -> streamFuture.thenCompose(stream -> commandToAsync(buf, streamOffset, stream, executor)
+              .whenComplete((l, e) -> commandMetrics.stop(context, e == null))));
     }
 
     void cleanUp() {
@@ -114,10 +124,12 @@ static class RemoteStream {
     private final DataStreamOutputImpl out;
     private final AtomicReference<CompletableFuture<DataStreamReply>> sendFuture
         = new AtomicReference<>(CompletableFuture.completedFuture(null));
-    private final RequestMetrics metrics;
+    private final RequestMetrics writeMetrics;
+    private final RequestMetrics commandMetrics;
 
-    RemoteStream(DataStreamOutputImpl out, RequestMetrics metrics) {
-      this.metrics = metrics;
+    RemoteStream(DataStreamOutputImpl out, RequestMetrics writeMetrics, RequestMetrics commandMetrics) {
+      this.writeMetrics = writeMetrics;
+      this.commandMetrics = commandMetrics;
       this.out = out;
     }
 
@@ -130,10 +142,17 @@ static Iterable<WriteOption> addFlush(List<WriteOption> original) {
     }
 
     CompletableFuture<DataStreamReply> write(DataStreamRequestByteBuf request, Executor executor) {
-      final Timekeeper.Context context = metrics.start();
+      final Timekeeper.Context context = writeMetrics.start();
       return composeAsync(sendFuture, executor,
           n -> out.writeAsync(request.slice().retain(), addFlush(request.getWriteOptionList()))
-              .whenComplete((l, e) -> metrics.stop(context, e == null)));
+              .whenComplete((l, e) -> writeMetrics.stop(context, e == null)));
+    }
+
+    CompletableFuture<DataStreamReply> command(DataStreamRequestByteBuf request, Executor executor) {
+      final Timekeeper.Context context = commandMetrics.start();
+      return composeAsync(sendFuture, executor,
+          n -> out.commandAsync(request.slice().retain())
+              .whenComplete((l, e) -> commandMetrics.stop(context, e == null)));
     }
   }
 
@@ -152,12 +171,16 @@ static class StreamInfo {
         throws IOException {
       this.request = request;
       this.primary = primary;
-      this.local = new LocalStream(stream, metricsConstructor.apply(RequestType.LOCAL_WRITE));
+      this.local = new LocalStream(stream,
+          metricsConstructor.apply(RequestType.LOCAL_WRITE),
+          metricsConstructor.apply(RequestType.LOCAL_COMMAND));
       this.division = division;
       final Set<RaftPeer> successors = getSuccessors(division.getId());
       final Set<DataStreamOutputImpl> outs = getStreams.apply(request, successors);
       this.remotes = outs.stream()
-          .map(o -> new RemoteStream(o, metricsConstructor.apply(RequestType.REMOTE_WRITE)))
+          .map(o -> new RemoteStream(o,
+              metricsConstructor.apply(RequestType.REMOTE_WRITE),
+              metricsConstructor.apply(RequestType.REMOTE_COMMAND)))
           .collect(Collectors.toSet());
     }
 
@@ -303,6 +326,24 @@ static CompletableFuture<Long> writeToAsync(ByteBuf buf,
     return CompletableFuture.supplyAsync(() -> writeTo(buf, options, stream), e);
   }
 
+  static CompletableFuture<Long> commandToAsync(ByteBuf buf, long streamOffset, DataStream stream,
+      Executor defaultExecutor) {
+    final Executor e = Optional.ofNullable(stream.getExecutor()).orElse(defaultExecutor);
+    final ByteBuffer command = copyBuffer(buf);
+    return CompletableFuture.runAsync(() -> {}, e)
+        .thenCompose(v -> stream.onCommand(command, streamOffset))
+        .thenApply(v -> 0L);
+  }
+
+  static ByteBuffer copyBuffer(ByteBuf buf) {
+    final ByteBuffer copy = ByteBuffer.allocate(buf.readableBytes());
+    for (ByteBuffer buffer : buf.nioBuffers()) {
+      copy.put(buffer);
+    }
+    copy.flip();
+    return copy;
+  }
+
   static long writeTo(ByteBuf buf, Iterable<WriteOption> options,
                       DataStream stream) {
     final DataChannel channel = stream.getDataChannel();
@@ -477,6 +518,9 @@ private void readImpl(DataStreamRequestByteBuf request, ChannelHandlerContext ct
     } else if (request.getType() == Type.STREAM_DATA) {
       localWrite = info.getLocal().write(request.slice(), request.getWriteOptionList(), writeExecutor);
       remoteWrites = info.applyToRemotes(out -> out.write(request, requestExecutor));
+    } else if (request.getType() == Type.STREAM_COMMAND) {
+      localWrite = info.getLocal().command(request.slice(), request.getStreamOffset(), writeExecutor);
+      remoteWrites = info.applyToRemotes(out -> out.command(request, requestExecutor));
     } else {
       throw new IllegalStateException(this + ": Unexpected type " + request.getType() + ", request=" + request);
     }
@@ -485,6 +529,7 @@ private void readImpl(DataStreamRequestByteBuf request, ChannelHandlerContext ct
         .thenCombineAsync(localWrite, (v, bytesWritten) -> {
           if (request.getType() == Type.STREAM_HEADER
               || request.getType() == Type.STREAM_DATA
+              || request.getType() == Type.STREAM_COMMAND
               || close) {
             sendReply(remoteWrites, request, bytesWritten, info.getCommitInfos(), ctx);
           } else {

ratis-proto/src/main/proto/Raft.proto

@@ -343,6 +343,7 @@ message DataStreamPacketHeaderProto {
   enum Type {
     STREAM_HEADER = 0;
     STREAM_DATA = 1;
+    STREAM_COMMAND = 2;
   }
 
   enum Option {

ratis-server-api/src/main/java/org/apache/ratis/statemachine/StateMachine.java

@@ -385,6 +385,19 @@ interface DataStream {
     default Executor getExecutor() {
       return null;
     }
+
+    /**
+     * Handle a command received in the middle of a data stream.
+     * The {@code streamOffset} indicates the current byte offset in the stream
+     * (i.e. the total number of data bytes received so far).
+     *
+     * @param command the command payload
+     * @param streamOffset the current stream byte offset
+     * @return a future for the command task
+     */
+    default CompletableFuture<?> onCommand(ByteBuffer command, long streamOffset) {
+      return CompletableFuture.completedFuture(null);
+    }
   }
 
   /**

ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java

@@ -20,22 +20,31 @@
 import org.apache.ratis.client.DataStreamClient;
 import org.apache.ratis.client.DataStreamClientRpc;
 import org.apache.ratis.client.api.DataStreamInput;
+import org.apache.ratis.client.impl.DataStreamClientImpl.DataStreamOutputImpl;
 import org.apache.ratis.conf.RaftProperties;
 import org.apache.ratis.datastream.DataStreamObserver;
+import org.apache.ratis.datastream.impl.DataStreamPacketByteBuffer;
+import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer;
 import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer;
+import org.apache.ratis.io.StandardWriteOption;
+import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type;
 import org.apache.ratis.proto.RaftProtos.RaftClientRequestProto;
 import org.apache.ratis.protocol.ClientId;
 import org.apache.ratis.protocol.DataStreamReply;
 import org.apache.ratis.protocol.DataStreamRequest;
 import org.apache.ratis.protocol.RaftClientRequest;
 import org.apache.ratis.protocol.RaftGroupId;
 import org.apache.ratis.protocol.RaftPeer;
+import org.apache.ratis.protocol.exceptions.AlreadyClosedException;
 import org.apache.ratis.util.ReferenceCountedObject;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 import java.io.EOFException;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicReference;
@@ -88,6 +97,86 @@ private static DataStreamClient newDataStreamClient(
         ClientId.randomId(), RaftGroupId.randomId(), dataStreamServer, dataStreamClientRpc, properties);
   }
 
+  private static class AllRequestRecordingRpc implements DataStreamClientRpc {
+    private final List<DataStreamRequest> requests = Collections.synchronizedList(new ArrayList<>());
+
+    @Override
+    public CompletableFuture<DataStreamReply> streamAsync(DataStreamRequest request) {
+      requests.add(request);
+      final long bytesWritten = request.getType() == Type.STREAM_DATA ? request.getDataLength() : 0;
+      return CompletableFuture.completedFuture(DataStreamReplyByteBuffer.newBuilder()
+          .setDataStreamPacket(request)
+          .setBuffer(DataStreamPacketByteBuffer.EMPTY_BYTE_BUFFER)
+          .setSuccess(true)
+          .setBytesWritten(bytesWritten)
+          .build());
+    }
+
+    List<DataStreamRequest> getRequests() {
+      return requests;
+    }
+
+    @Override
+    public void close() {
+    }
+  }
+
+  private static DataStreamOutputImpl newDataStreamOutput(AllRequestRecordingRpc rpc) {
+    final RaftPeer server = newPeer("server");
+    final RaftProperties properties = new RaftProperties();
+    final DataStreamClientImpl client = new DataStreamClientImpl(
+        ClientId.randomId(), RaftGroupId.randomId(), server, rpc, properties);
+    return (DataStreamOutputImpl) client.stream((ByteBuffer) null);
+  }
+
+  @Test
+  public void testCommandAsyncSendsStreamCommandWithCurrentOffset() {
+    final AllRequestRecordingRpc rpc = new AllRequestRecordingRpc();
+    final DataStreamOutputImpl out = newDataStreamOutput(rpc);
+
+    out.getHeaderFuture().join();
+    out.writeAsync(ByteBuffer.allocate(5)).join();
+    out.commandAsync(ByteBuffer.wrap(new byte[] {'c', 't', 'r', 'l'})).join();
+
+    final List<DataStreamRequest> requests = rpc.getRequests();
+    Assertions.assertEquals(3, requests.size());
+    Assertions.assertEquals(Type.STREAM_HEADER, requests.get(0).getType());
+    Assertions.assertEquals(Type.STREAM_DATA, requests.get(1).getType());
+    Assertions.assertEquals(0, requests.get(1).getStreamOffset());
+    Assertions.assertEquals(Type.STREAM_COMMAND, requests.get(2).getType());
+    Assertions.assertEquals(5, requests.get(2).getStreamOffset());
+    Assertions.assertEquals(4, requests.get(2).getDataLength());
+  }
+
+  @Test
+  public void testCommandAsyncDoesNotAdvanceStreamOffset() {
+    final AllRequestRecordingRpc rpc = new AllRequestRecordingRpc();
+    final DataStreamOutputImpl out = newDataStreamOutput(rpc);
+
+    out.getHeaderFuture().join();
+    out.writeAsync(ByteBuffer.allocate(5)).join();
+    out.commandAsync(ByteBuffer.wrap(new byte[] {'c', 't', 'r', 'l'})).join();
+    out.writeAsync(ByteBuffer.allocate(3)).join();
+
+    final List<DataStreamRequest> requests = rpc.getRequests();
+    Assertions.assertEquals(4, requests.size());
+    Assertions.assertEquals(Type.STREAM_DATA, requests.get(3).getType());
+    Assertions.assertEquals(5, requests.get(3).getStreamOffset());
+  }
+
+  @Test
+  public void testCommandAsyncAfterCloseFails() {
+    final AllRequestRecordingRpc rpc = new AllRequestRecordingRpc();
+    final DataStreamOutputImpl out = newDataStreamOutput(rpc);
+
+    out.getHeaderFuture().join();
+    out.writeAsync(DataStreamPacketByteBuffer.EMPTY_BYTE_BUFFER, StandardWriteOption.CLOSE).join();
+
+    final ExecutionException exception = Assertions.assertThrows(ExecutionException.class,
+        () -> out.commandAsync(ByteBuffer.wrap(new byte[] {'c'})).get());
+    Assertions.assertInstanceOf(AlreadyClosedException.class, exception.getCause());
+  }
+
   @Test
   public void testReadOnlyInputCompletesPendingReadOnCompleted() throws Exception {
     final RaftPeer follower = newPeer("follower");

ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamAsyncClusterTests.java

@@ -19,6 +19,7 @@
 
 import org.apache.ratis.netty.client.NettyClientStreamRpc;
 import org.apache.ratis.protocol.RaftPeer;
+import org.apache.ratis.protocol.RoutingTable;
 import org.apache.ratis.server.impl.MiniRaftCluster;
 import org.apache.ratis.RaftTestUtil;
 import org.apache.ratis.client.RaftClient;
@@ -41,6 +42,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -50,9 +52,19 @@
 
 @Timeout(value = 300)
 public abstract class DataStreamAsyncClusterTests<CLUSTER extends MiniRaftCluster>
-    extends DataStreamClusterTests<CLUSTER> {
+    extends DataStreamClusterTests<CLUSTER> implements DataStreamCommandE2ETestCases {
   final Executor executor = Executors.newFixedThreadPool(16);
 
+  @Override
+  public MiniRaftCluster.Factory.Get<CLUSTER> getClusterFactory() {
+    return this;
+  }
+
+  @Override
+  public RoutingTable routingTable(Collection<RaftPeer> peers, RaftPeer primary) {
+    return getRoutingTable(peers, primary);
+  }
+
   @Test
   public void testSingleStreamsMultipleServers() throws Exception {
     Slf4jUtils.setLogLevel(NettyClientStreamRpc.LOG, Level.TRACE);