Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
szetszwo merged 14 commits into from
Aug 4, 2026
Merged

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

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
ab19857
update
amaliujia Jul 29, 2026
2cac57a
fix
amaliujia Jul 29, 2026
14fc6c3
update
amaliujia Jul 29, 2026
eea3b1b
update
amaliujia Jul 29, 2026
00fd13b
update
amaliujia Jul 29, 2026
c2a7221
update
amaliujia Jul 30, 2026
ffad563
remove filestore change
amaliujia Jul 30, 2026
29d0df7
update
amaliujia Jul 30, 2026
51d0d14
address review comments
amaliujia Jul 31, 2026
3fcff8d
update
amaliujia Jul 31, 2026
2b80b08
address comments
Aug 3, 2026
1b69a7a
fix findsBug
Aug 3, 2026
d417115
address comments
amaliujia Aug 3, 2026
02ee091
Make one call to the map
amaliujia Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamOutput.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
17 changes: 17 additions & 0 deletions ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -187,6 +199,11 @@ public CompletableFuture<DataStreamReply> writeAsync(FilePositionCount src, Writ
return writeAsyncImpl(src, src.getCount(), Arrays.asList(options));
}

@Override
public CompletableFuture<DataStreamReply> commandAsync(ByteBuffer src) {

@ss77892 ss77892 Aug 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What would happen on the client side if it issues two commands without data in between? As far as I can see, NettyClientReply, when it maps requests, considers only the stream offset and type, so the second command with the same offset wouldn't have a ReplyEntry.

@amaliujia amaliujia Aug 3, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very good call. It seems to me that the first command's ReplyEntry will be used for the second command and the second command's RequestEntry won't be in the map. Anyway it looks like a mess without a handling.

So I instead fail the commands if there is already one at the same stream offset. After all, we do not expect the caller issue multiple commands at the same stream offset for now.

return commandAsyncImpl(src, src.remaining());

@szetszwo szetszwo Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

commandAsyncImpl is used just once. Let's inline the code.

    public CompletableFuture<DataStreamReply> commandAsync(ByteBuffer src) {
      if (isClosed()) {
        return JavaUtils.completeExceptionally(new AlreadyClosedException(
            clientId + ": stream already closed, request=" + header));
      }
      return combineHeader(send(Type.STREAM_COMMAND, src, src.remaining(),
          Collections.singleton(StandardWriteOption.FLUSH)));
    }

}

boolean isClosed() {
return closeFuture != null;
}
Expand Down
3 changes: 2 additions & 1 deletion ratis-netty/src/main/java/org/apache/ratis/netty/metrics/NettyServerStreamRpcMetrics.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
160 changes: 130 additions & 30 deletions ratis-netty/src/main/java/org/apache/ratis/netty/server/DataStreamManagement.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@
import org.apache.ratis.util.MemoizedSupplier;
import org.apache.ratis.util.Preconditions;
import org.apache.ratis.util.ReferenceCountedObject;
import org.apache.ratis.util.StringUtils;
import org.apache.ratis.util.TimeDuration;
import org.apache.ratis.util.TimeoutExecutor;
import org.apache.ratis.util.UncheckedAutoCloseable;
Expand All @@ -86,23 +87,65 @@
public class DataStreamManagement {
public static final Logger LOG = LoggerFactory.getLogger(DataStreamManagement.class);

static final class LocalResult {
static final LocalResult ZERO = of(0);

static LocalResult of(long byteWritten) {
return new LocalResult(byteWritten, null);
}

static LocalResult of(ByteBuffer reply) {
return new LocalResult(0, reply);
}

/** For {@link Type#STREAM_DATA}. */
private final long byteWritten;
/** For {@link Type#STREAM_COMMAND}. */
private final ByteBuffer commandReply;

private LocalResult(long byteWritten, ByteBuffer commandReply) {
this.byteWritten = byteWritten;
this.commandReply = commandReply;
}

long getByteWritten() {
return byteWritten;
}

ByteBuffer getCommandReply() {
return commandReply;
}
}

static class LocalStream {
private final CompletableFuture<DataStream> streamFuture;
private final AtomicReference<CompletableFuture<Long>> writeFuture;
private final RequestMetrics metrics;
private final AtomicReference<CompletableFuture<LocalResult>> writeFuture;

@szetszwo szetszwo Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's rename it to resultFuture.

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.writeFuture = new AtomicReference<>(streamFuture.thenApply(s -> LocalResult.ZERO));
this.writeMetrics = writeMetrics;
this.commandMetrics = commandMetrics;
}

CompletableFuture<Long> write(ByteBuf buf, Iterable<WriteOption> options,
CompletableFuture<LocalResult> write(ByteBuf buf, Iterable<WriteOption> options,
Executor executor) {

@szetszwo szetszwo Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make it a single line.

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)))
.thenApply(LocalResult::of));
}

CompletableFuture<LocalResult> command(ByteBuffer command, long streamOffset, Executor executor) {
final Timekeeper.Context context = commandMetrics.start();
return composeAsync(writeFuture, executor,
n -> streamFuture.thenCompose(stream -> commandToAsync(command, streamOffset, stream, executor)
.whenComplete((l, e) -> commandMetrics.stop(context, e == null)))
.thenApply(LocalResult::of));
}

void cleanUp() {
Expand All @@ -114,10 +157,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;
}

Expand All @@ -130,10 +175,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(ByteBuffer command, Executor executor) {
final Timekeeper.Context context = commandMetrics.start();
return composeAsync(sendFuture, executor,
n -> out.commandAsync(command)
.whenComplete((l, e) -> commandMetrics.stop(context, e == null)));
}
}

Expand All @@ -152,12 +204,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());
}

Expand Down Expand Up @@ -303,6 +359,22 @@ static CompletableFuture<Long> writeToAsync(ByteBuf buf,
return CompletableFuture.supplyAsync(() -> writeTo(buf, options, stream), e);
}

static CompletableFuture<ByteBuffer> commandToAsync(ByteBuffer command, long streamOffset,
DataStream stream, Executor defaultExecutor) {
final Executor e = Optional.ofNullable(stream.getExecutor()).orElse(defaultExecutor);
return CompletableFuture.runAsync(() -> {}, e)
.thenCompose(ignored -> stream.onCommand(command, streamOffset));

@szetszwo szetszwo Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using completedFuture(null) is more efficient since it doesn't have to run an empty task.

    return CompletableFuture.completedFuture(null)
        .thenComposeAsync(ignored -> stream.onCommand(command, streamOffset), e);

@amaliujia amaliujia Aug 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. I really learned something from this comment.

}

static ByteBuffer copyBuffer(ByteBuf buf) {
final ByteBuffer copy = ByteBuffer.allocate(buf.readableBytes());
for (ByteBuffer buffer : buf.nioBuffers()) {
copy.put(buffer);
}
copy.flip();
return copy;

@szetszwo szetszwo Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make it readonly

    return copy.asReadOnlyBuffer();

}

static long writeTo(ByteBuf buf, Iterable<WriteOption> options,
DataStream stream) {
final DataChannel channel = stream.getDataChannel();
Expand Down Expand Up @@ -357,15 +429,18 @@ static DataStreamReplyByteBuffer newDataStreamReplyByteBuffer(DataStreamRequestB
}

private void sendReply(List<CompletableFuture<DataStreamReply>> remoteWrites,
DataStreamRequestByteBuf request, long bytesWritten, Collection<CommitInfoProto> commitInfos,
DataStreamRequestByteBuf request, LocalResult localResult, Collection<CommitInfoProto> commitInfos,
ChannelHandlerContext ctx) {
final boolean success = checkSuccessRemoteWrite(remoteWrites, bytesWritten, request);
final boolean success = checkSuccessRemoteWrite(remoteWrites, localResult, request);
final DataStreamReplyByteBuffer.Builder builder = DataStreamReplyByteBuffer.newBuilder()
.setDataStreamPacket(request)
.setSuccess(success)
.setCommitInfos(commitInfos);
if (success) {
builder.setBytesWritten(bytesWritten);
builder.setBytesWritten(localResult.getByteWritten());
if (localResult.getCommandReply() != null) {
builder.setBuffer(localResult.getCommandReply());
}
}
ctx.writeAndFlush(builder.build());
}
Expand Down Expand Up @@ -469,24 +544,31 @@ private void readImpl(DataStreamRequestByteBuf request, ChannelHandlerContext ct
() -> new IllegalStateException("Failed to get StreamInfo for " + request));
}

final CompletableFuture<Long> localWrite;
final CompletableFuture<LocalResult> localResult;
final List<CompletableFuture<DataStreamReply>> remoteWrites;
if (request.getType() == Type.STREAM_HEADER) {
localWrite = CompletableFuture.completedFuture(0L);
localResult = CompletableFuture.completedFuture(LocalResult.ZERO);
remoteWrites = Collections.emptyList();
} else if (request.getType() == Type.STREAM_DATA) {
localWrite = info.getLocal().write(request.slice(), request.getWriteOptionList(), writeExecutor);
localResult = info.getLocal().write(request.slice(), request.getWriteOptionList(), writeExecutor);
remoteWrites = info.applyToRemotes(out -> out.write(request, requestExecutor));
} else if (request.getType() == Type.STREAM_COMMAND) {
// command is supposed to have small data size, just copy it
final ByteBuffer command = copyBuffer(request.slice());
localResult = info.getLocal().command(command, request.getStreamOffset(), writeExecutor);
remoteWrites = info.applyToRemotes(out -> out.command(
copyBuffer(request.slice()), requestExecutor));

@szetszwo szetszwo Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

duplicate() and reuse command, i.e.

      final ByteBuffer command = copyBuffer(request.slice());
      localResult = info.getLocal().command(command.duplicate(), request.getStreamOffset(), writeExecutor);
      remoteWrites = info.applyToRemotes(out -> out.command(command, requestExecutor));

} else {
throw new IllegalStateException(this + ": Unexpected type " + request.getType() + ", request=" + request);
}

composeAsync(info.getPrevious(), requestExecutor, n -> JavaUtils.allOf(remoteWrites)
.thenCombineAsync(localWrite, (v, bytesWritten) -> {
.thenCombineAsync(localResult, (v, local) -> {
if (request.getType() == Type.STREAM_HEADER
|| request.getType() == Type.STREAM_DATA
|| request.getType() == Type.STREAM_COMMAND
|| close) {
sendReply(remoteWrites, request, bytesWritten, info.getCommitInfos(), ctx);
sendReply(remoteWrites, request, local, info.getCommitInfos(), ctx);
} else {
throw new IllegalStateException(this + ": Unexpected type " + request.getType() + ", request=" + request);
}
Expand Down Expand Up @@ -519,30 +601,48 @@ static void assertReplyCorrespondingToRequest(
Preconditions.assertTrue(request.getStreamOffset() == reply.getStreamOffset());
}

private boolean checkSuccessRemoteWrite(List<CompletableFuture<DataStreamReply>> replyFutures, long bytesWritten,
final DataStreamRequestByteBuf request) {
private boolean checkSuccessRemoteWrite(List<CompletableFuture<DataStreamReply>> replyFutures,
LocalResult localResult, final DataStreamRequestByteBuf request) {
for (CompletableFuture<DataStreamReply> replyFuture : replyFutures) {
final DataStreamReply reply;
try {
reply = replyFuture.get(requestTimeout.getDuration(), requestTimeout.getUnit());
} catch (Exception e) {
throw new CompletionException("Failed to get reply for bytesWritten=" + bytesWritten + ", " + request, e);
throw new CompletionException("Failed to get reply for " + localResult + ", " + request, e);

@szetszwo szetszwo Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add toString()

    @Override
    public String toString() {
      return commandReply != null ? "commandReply:" + StringUtils.bytes2HexString(commandReply)
          : "byteWritten:" + byteWritten;
    }

}
assertReplyCorrespondingToRequest(request, reply);
if (!reply.isSuccess()) {
LOG.warn("reply is not success, request: {}", request);
return false;
}
if (reply.getBytesWritten() != bytesWritten) {
LOG.warn(
"reply written bytes not match, local size: {} remote size: {} request: {}",
bytesWritten, reply.getBytesWritten(), request);
if (reply.getBytesWritten() != localResult.getByteWritten()) {
LOG.warn("reply written bytes not match, local size: {} remote size: {} request: {}",
localResult.getByteWritten(), reply.getBytesWritten(), request);
return false;
}
if (!commandReplyEquals(reply.nioBuffer(), localResult.getCommandReply())) {
LOG.warn("reply buffer not match, local: {} remote: {} request: {}",
commandReplyString(localResult.getCommandReply()),
commandReplyString(reply.nioBuffer()),
request);
return false;
}
}
return true;
}

private static boolean commandReplyEquals(ByteBuffer a, ByteBuffer b) {
if (a == null || a.remaining() == 0) {
return b == null || b.remaining() == 0;
}
return a.equals(b);
}

private static String commandReplyString(ByteBuffer buffer) {
return buffer == null || buffer.remaining() == 0
? "null" : StringUtils.bytes2HexString(buffer);
}

NettyServerStreamRpcMetrics getMetrics() {
return nettyServerStreamRpcMetrics;
}
Expand Down
1 change: 1 addition & 0 deletions ratis-proto/src/main/proto/Raft.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -343,6 +343,7 @@ message DataStreamPacketHeaderProto {
enum Type {
STREAM_HEADER = 0;
STREAM_DATA = 1;
STREAM_COMMAND = 2;
}

enum Option {
Expand Down
13 changes: 13 additions & 0 deletions ratis-server-api/src/main/java/org/apache/ratis/statemachine/StateMachine.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<ByteBuffer> onCommand(ByteBuffer command, long streamOffset) {
return CompletableFuture.completedFuture(null);
}
}

/**
Expand Down
Loading
Loading