From bd30f5171f4108a35cbae13128c181fe6aa44d27 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 13 Jun 2026 21:24:20 +0800 Subject: [PATCH 01/13] RATIS-2560 Signed-off-by: peterxcli --- dev-support/checkstyle.xml | 4 + .../dev-support/findbugsExcludeFile.xml | 6 +- .../ratis/client/DataStreamClientRpc.java | 8 + .../ratis/client/api/DataStreamApi.java | 10 + .../ratis/client/api/DataStreamInput.java | 37 ++++ .../ratis/client/impl/ClientProtoUtils.java | 11 +- .../dev-support/findbugsExcludeFile.xml | 10 +- .../impl/DataStreamReplyByteBuf.java | 174 ++++++++++++++++++ .../ratis/protocol/DataStreamReply.java | 15 +- .../ratis/netty/NettyDataStreamUtils.java | 44 ++++- .../netty/server/NettyServerStreamRpc.java | 19 +- .../netty/server/ReadStreamManagement.java | 49 +++-- .../client/impl/TestDataStreamClientImpl.java | 118 ++++++++++++ .../datastream/DataStreamClusterTests.java | 83 +++++++++ .../server/TestDataStreamManagement.java | 4 +- 15 files changed, 567 insertions(+), 25 deletions(-) create mode 100644 ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java create mode 100644 ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamReplyByteBuf.java create mode 100644 ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java diff --git a/dev-support/checkstyle.xml b/dev-support/checkstyle.xml index db4954fb49..f7e168b029 100644 --- a/dev-support/checkstyle.xml +++ b/dev-support/checkstyle.xml @@ -60,6 +60,10 @@ + + + + diff --git a/ratis-client/dev-support/findbugsExcludeFile.xml b/ratis-client/dev-support/findbugsExcludeFile.xml index 3a808c4486..d7e97b5893 100644 --- a/ratis-client/dev-support/findbugsExcludeFile.xml +++ b/ratis-client/dev-support/findbugsExcludeFile.xml @@ -27,6 +27,10 @@ + + + + @@ -39,4 +43,4 @@ - \ No newline at end of file + diff --git a/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java b/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java index a9bcd9d58a..fd5bd9538e 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java @@ -24,6 +24,7 @@ import java.io.Closeable; import java.util.concurrent.CompletableFuture; +import java.util.function.Consumer; /** * A client interface for sending stream requests. @@ -36,4 +37,11 @@ default CompletableFuture streamAsync(DataStreamRequest request throw new UnsupportedOperationException(getClass() + " does not support " + JavaUtils.getCurrentStackTraceElement().getMethodName()); } + + /** Async call to send a request and receive multiple replies for the request. */ + default CompletableFuture streamAsync( + DataStreamRequest request, Consumer replyConsumer) { + throw new UnsupportedOperationException(getClass() + " does not support " + + JavaUtils.getCurrentStackTraceElement().getMethodName()); + } } diff --git a/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamApi.java b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamApi.java index 9e5e2438cb..800337d8f5 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamApi.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamApi.java @@ -50,4 +50,14 @@ default DataStreamOutput stream() { /** Create a stream by providing a customized header message and route table. */ DataStreamOutput stream(ByteBuffer headerMessage, RoutingTable routingTable); + + /** + * Create a stream to read data for readonly requests. + */ + default DataStreamInput streamReadOnly() { + return streamReadOnly(null); + } + + /** Create a stream by providing a customized header message for readonly requests. */ + DataStreamInput streamReadOnly(ByteBuffer message); } diff --git a/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java new file mode 100644 index 0000000000..0356a6a97f --- /dev/null +++ b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java @@ -0,0 +1,37 @@ +/* + * 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.client.api; + +import org.apache.ratis.protocol.DataStreamReply; + +import java.io.Closeable; +import java.util.concurrent.CompletableFuture; + +/** + * An asynchronous input stream supporting zero buffer copying. + */ +public interface DataStreamInput extends Closeable { + /** + * Read the next chunk in the stream asynchronously. + * The caller owns the returned {@link DataStreamReply} and should call + * {@link DataStreamReply#release()} after consuming it. + * + * @return a future of the reply. + */ + CompletableFuture readAsync(); +} diff --git a/ratis-client/src/main/java/org/apache/ratis/client/impl/ClientProtoUtils.java b/ratis-client/src/main/java/org/apache/ratis/client/impl/ClientProtoUtils.java index d2146a521f..1126453c27 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/impl/ClientProtoUtils.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/impl/ClientProtoUtils.java @@ -18,6 +18,7 @@ package org.apache.ratis.client.impl; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.proto.RaftProtos.AlreadyClosedExceptionProto; import org.apache.ratis.proto.RaftProtos.ClientMessageEntryProto; import org.apache.ratis.proto.RaftProtos.GroupAddRequestProto; @@ -379,11 +380,13 @@ static GroupInfoReplyProto toGroupInfoReplyProto(GroupInfoReply reply) { } static RaftClientReply getRaftClientReply(DataStreamReply reply) { - if (!(reply instanceof DataStreamReplyByteBuffer)) { - throw new IllegalStateException("Unexpected " + reply.getClass() + ": reply is " + reply); - } try { - return toRaftClientReply(((DataStreamReplyByteBuffer) reply).slice()); + if (reply instanceof DataStreamReplyByteBuffer) { + return toRaftClientReply(((DataStreamReplyByteBuffer) reply).slice()); + } else if (reply instanceof DataStreamReplyByteBuf) { + return toRaftClientReply(((DataStreamReplyByteBuf) reply).slice().nioBuffer()); + } + throw new IllegalStateException("Unexpected " + reply.getClass() + ": reply is " + reply); } catch (InvalidProtocolBufferException e) { throw new IllegalStateException("Failed to getRaftClientReply from " + reply, e); } diff --git a/ratis-common/dev-support/findbugsExcludeFile.xml b/ratis-common/dev-support/findbugsExcludeFile.xml index 882f08b7fa..cec157b527 100644 --- a/ratis-common/dev-support/findbugsExcludeFile.xml +++ b/ratis-common/dev-support/findbugsExcludeFile.xml @@ -19,6 +19,10 @@ + + + + @@ -43,6 +47,10 @@ + + + + @@ -111,4 +119,4 @@ - \ No newline at end of file + diff --git a/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamReplyByteBuf.java b/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamReplyByteBuf.java new file mode 100644 index 0000000000..d94d3f0c33 --- /dev/null +++ b/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamReplyByteBuf.java @@ -0,0 +1,174 @@ +/* + * 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.datastream.impl; + +import org.apache.ratis.proto.RaftProtos.CommitInfoProto; +import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; +import org.apache.ratis.protocol.ClientId; +import org.apache.ratis.protocol.DataStreamPacket; +import org.apache.ratis.protocol.DataStreamReply; +import org.apache.ratis.protocol.DataStreamReplyHeader; +import org.apache.ratis.thirdparty.io.netty.buffer.ByteBuf; +import org.apache.ratis.thirdparty.io.netty.buffer.Unpooled; + +import java.util.Collection; +import java.util.Collections; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Implements {@link DataStreamReply} with {@link ByteBuf}. + */ +public final class DataStreamReplyByteBuf extends DataStreamPacketImpl implements DataStreamReply { + public static final class Builder { + private ClientId clientId; + private Type type; + private long streamId; + private long streamOffset; + private ByteBuf buffer; + + private boolean success; + private long bytesWritten; + private Collection commitInfos; + + private Builder() { + } + + public Builder setClientId(ClientId clientId) { + this.clientId = clientId; + return this; + } + + public Builder setType(Type type) { + this.type = type; + return this; + } + + public Builder setStreamId(long streamId) { + this.streamId = streamId; + return this; + } + + public Builder setStreamOffset(long streamOffset) { + this.streamOffset = streamOffset; + return this; + } + + public Builder setBuffer(ByteBuf buffer) { + this.buffer = buffer; + return this; + } + + public Builder setSuccess(boolean success) { + this.success = success; + return this; + } + + public Builder setBytesWritten(long bytesWritten) { + this.bytesWritten = bytesWritten; + return this; + } + + public Builder setCommitInfos(Collection commitInfos) { + this.commitInfos = commitInfos; + return this; + } + + public Builder setDataStreamReplyHeader(DataStreamReplyHeader header) { + return setDataStreamPacket(header) + .setSuccess(header.isSuccess()) + .setBytesWritten(header.getBytesWritten()) + .setCommitInfos(header.getCommitInfos()); + } + + public Builder setDataStreamPacket(DataStreamPacket packet) { + return setClientId(packet.getClientId()) + .setType(packet.getType()) + .setStreamId(packet.getStreamId()) + .setStreamOffset(packet.getStreamOffset()); + } + + public DataStreamReplyByteBuf build() { + return new DataStreamReplyByteBuf( + clientId, type, streamId, streamOffset, buffer, success, bytesWritten, commitInfos); + } + } + + public static Builder newBuilder() { + return new Builder(); + } + + private final AtomicReference buffer; + private final boolean success; + private final long bytesWritten; + private final Collection commitInfos; + + @SuppressWarnings("parameternumber") + private DataStreamReplyByteBuf(ClientId clientId, Type type, long streamId, long streamOffset, ByteBuf buffer, + boolean success, long bytesWritten, Collection commitInfos) { + super(clientId, type, streamId, streamOffset); + this.buffer = new AtomicReference<>(buffer != null ? buffer.asReadOnly() : Unpooled.EMPTY_BUFFER); + this.success = success; + this.bytesWritten = bytesWritten; + this.commitInfos = commitInfos != null ? commitInfos : Collections.emptyList(); + } + + private ByteBuf getBuffer() { + return Optional.ofNullable(buffer.get()).orElseThrow( + () -> new IllegalStateException("buffer is already released in " + this)); + } + + @Override + public long getDataLength() { + return getBuffer().readableBytes(); + } + + public ByteBuf slice() { + return getBuffer().slice(); + } + + @Override + public boolean isSuccess() { + return success; + } + + @Override + public long getBytesWritten() { + return bytesWritten; + } + + @Override + public Collection getCommitInfos() { + return commitInfos; + } + + @Override + public void release() { + final ByteBuf got = buffer.getAndSet(null); + if (got != null && got != Unpooled.EMPTY_BUFFER) { + got.release(); + } + } + + @Override + public String toString() { + return super.toString() + + "," + (success ? "SUCCESS" : "FAILED") + + ",bytesWritten=" + bytesWritten; + } +} diff --git a/ratis-common/src/main/java/org/apache/ratis/protocol/DataStreamReply.java b/ratis-common/src/main/java/org/apache/ratis/protocol/DataStreamReply.java index 459aee363c..cf1f826867 100644 --- a/ratis-common/src/main/java/org/apache/ratis/protocol/DataStreamReply.java +++ b/ratis-common/src/main/java/org/apache/ratis/protocol/DataStreamReply.java @@ -22,7 +22,7 @@ import java.util.Collection; -public interface DataStreamReply extends DataStreamPacket { +public interface DataStreamReply extends DataStreamPacket, AutoCloseable { boolean isSuccess(); @@ -30,4 +30,15 @@ public interface DataStreamReply extends DataStreamPacket { /** @return the commit information when the reply is created. */ Collection getCommitInfos(); -} \ No newline at end of file + + /** + * Release resources owned by this reply. + */ + default void release() { + } + + @Override + default void close() { + release(); + } +} diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java b/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java index 583d6e3e94..bd30339b33 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java @@ -18,6 +18,7 @@ package org.apache.ratis.netty; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; import org.apache.ratis.datastream.impl.DataStreamRequestFilePositionCount; import org.apache.ratis.io.FilePositionCount; @@ -30,6 +31,7 @@ import org.apache.ratis.protocol.ClientId; import org.apache.ratis.protocol.DataStreamPacketHeader; import org.apache.ratis.protocol.DataStreamReplyHeader; +import org.apache.ratis.protocol.DataStreamReply; import org.apache.ratis.protocol.DataStreamRequest; import org.apache.ratis.protocol.DataStreamRequestHeader; import org.apache.ratis.thirdparty.com.google.protobuf.InvalidProtocolBufferException; @@ -80,7 +82,7 @@ static ByteBuffer getDataStreamRequestHeaderProtoByteBuffer(DataStreamRequest re .asReadOnlyByteBuffer(); } - static ByteBuffer getDataStreamReplyHeaderProtoByteBuf(DataStreamReplyByteBuffer reply) { + static ByteBuffer getDataStreamReplyHeaderProtoByteBuf(DataStreamReply reply) { DataStreamPacketHeaderProto.Builder b = DataStreamPacketHeaderProto .newBuilder() .setClientId(reply.getClientId().toByteString()) @@ -159,6 +161,22 @@ static void encodeDataStreamReplyByteBuffer(DataStreamReplyByteBuffer reply, Con out.accept(Unpooled.wrappedBuffer(reply.slice())); } + static void encodeDataStreamReplyByteBuf(DataStreamReplyByteBuf reply, Consumer out, + ByteBufAllocator allocator) { + ByteBuffer headerBuf = getDataStreamReplyHeaderProtoByteBuf(reply); + final ByteBuf headerLenBuf = allocator.ioBuffer(DataStreamPacketHeader.getSizeOfHeaderLen()); + headerLenBuf.writeInt(headerBuf.remaining()); + out.accept(headerLenBuf); + out.accept(Unpooled.wrappedBuffer(headerBuf)); + + final ByteBuf data = reply.slice(); + if (data.readableBytes() == 0) { + out.accept(Unpooled.EMPTY_BUFFER); + } else { + out.accept(data.retain()); + } + } + static DataStreamRequestByteBuf decodeDataStreamRequestByteBuf(ByteBuf buf) { return Optional.ofNullable(decodeDataStreamRequestHeader(buf)) .map(header -> checkHeader(header, buf)) @@ -224,6 +242,30 @@ static DataStreamReplyByteBuffer decodeDataStreamReplyByteBuffer(ByteBuf buf) { .orElse(null); } + static DataStreamReplyByteBuffer toDataStreamReplyByteBuffer(DataStreamReplyByteBuf reply) { + try { + return DataStreamReplyByteBuffer.newBuilder() + .setDataStreamPacket(reply) + .setBuffer(copy(reply.slice())) + .setSuccess(reply.isSuccess()) + .setBytesWritten(reply.getBytesWritten()) + .setCommitInfos(reply.getCommitInfos()) + .build(); + } finally { + reply.release(); + } + } + + static DataStreamReplyByteBuf decodeDataStreamReplyByteBuf(ByteBuf buf) { + return Optional.ofNullable(decodeDataStreamReplyHeader(buf)) + .map(header -> checkHeader(header, buf)) + .map(header -> DataStreamReplyByteBuf.newBuilder() + .setDataStreamReplyHeader(header) + .setBuffer(decodeData(buf, header, ByteBuf::retainedSlice)) + .build()) + .orElse(null); + } + static DataStreamReplyHeader decodeDataStreamReplyHeader(ByteBuf buf) { if (DataStreamPacketHeader.getSizeOfHeaderLen() > buf.readableBytes()) { return null; diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/server/NettyServerStreamRpc.java b/ratis-netty/src/main/java/org/apache/ratis/netty/server/NettyServerStreamRpc.java index 643ed15a0f..3aff714da2 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/server/NettyServerStreamRpc.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/server/NettyServerStreamRpc.java @@ -23,6 +23,7 @@ import org.apache.ratis.client.impl.DataStreamClientImpl.DataStreamOutputImpl; import org.apache.ratis.conf.Parameters; import org.apache.ratis.conf.RaftProperties; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuf; import org.apache.ratis.netty.NettyConfigKeys; @@ -31,6 +32,7 @@ import org.apache.ratis.netty.metrics.NettyServerStreamRpcMetrics; import org.apache.ratis.protocol.ClientId; import org.apache.ratis.protocol.DataStreamPacket; +import org.apache.ratis.protocol.DataStreamReply; import org.apache.ratis.protocol.RaftClientRequest; import org.apache.ratis.protocol.RaftPeer; import org.apache.ratis.security.TlsConf; @@ -286,13 +288,22 @@ protected void decode(ChannelHandlerContext context, ByteBuf buf, List o }; } - static final MessageToMessageEncoder ENCODER = new Encoder(); + static final MessageToMessageEncoder ENCODER = new Encoder(); @ChannelHandler.Sharable - static class Encoder extends MessageToMessageEncoder { + static class Encoder extends MessageToMessageEncoder { @Override - protected void encode(ChannelHandlerContext context, DataStreamReplyByteBuffer reply, List out) { - NettyDataStreamUtils.encodeDataStreamReplyByteBuffer(reply, out::add, context.alloc()); + protected void encode(ChannelHandlerContext context, DataStreamReply reply, List out) { + if (reply instanceof DataStreamReplyByteBuffer) { + NettyDataStreamUtils.encodeDataStreamReplyByteBuffer((DataStreamReplyByteBuffer) reply, out::add, + context.alloc()); + } else if (reply instanceof DataStreamReplyByteBuf) { + try (DataStreamReplyByteBuf byteBufReply = (DataStreamReplyByteBuf) reply) { + NettyDataStreamUtils.encodeDataStreamReplyByteBuf(byteBufReply, out::add, context.alloc()); + } + } else { + throw new IllegalArgumentException("Unexpected DataStreamReply class " + reply.getClass()); + } } } diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java b/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java index bcdced1eac..4bdb807ff6 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java @@ -23,6 +23,7 @@ import org.apache.ratis.proto.RaftProtos.RaftClientRequestProto; import org.apache.ratis.proto.RaftProtos.RaftClientRequestProto.TypeCase; import org.apache.ratis.protocol.ClientId; +import org.apache.ratis.protocol.RaftClientReply; import org.apache.ratis.protocol.RaftClientRequest; import org.apache.ratis.protocol.exceptions.AlreadyClosedException; import org.apache.ratis.server.RaftServer; @@ -40,6 +41,7 @@ import java.util.concurrent.CompletableFuture; import static org.apache.ratis.client.impl.ClientProtoUtils.toRaftClientRequest; +import static org.apache.ratis.client.impl.ClientProtoUtils.toRaftClientReplyProto; import static org.apache.ratis.netty.server.DataStreamManagement.replyDataStreamException; public class ReadStreamManagement { @@ -48,13 +50,15 @@ public class ReadStreamManagement { static class ReadStream implements WritableByteChannel { private final ClientId clientId; private final long streamId; + private final RaftClientRequest request; private final ChannelHandlerContext ctx; private final CompletableFuture closed = new CompletableFuture<>(); private long streamOffset; - ReadStream(DataStreamRequestByteBuf request, ChannelHandlerContext ctx) { - this.clientId = request.getClientId(); - this.streamId = request.getStreamId(); + ReadStream(DataStreamRequestByteBuf requestBuf, RaftClientRequest request, ChannelHandlerContext ctx) { + this.clientId = requestBuf.getClientId(); + this.streamId = requestBuf.getStreamId(); + this.request = request; this.ctx = ctx; } @@ -64,8 +68,10 @@ public boolean isOpen() { } @Override - public void close() { - closed.complete(null); + public synchronized void close() throws IOException { + if (closed.complete(null)) { + writeAndFlush(newTerminalReply(), 0); + } } @Override @@ -76,19 +82,23 @@ public synchronized int write(ByteBuffer buffer) throws IOException { buffer = buffer.asReadOnlyBuffer(); final int length = buffer.remaining(); final DataStreamReplyByteBuffer reply = newReply(buffer); + writeAndFlush(reply, length); + streamOffset += length; + return length; + } + + private synchronized void writeAndFlush(DataStreamReplyByteBuffer reply, int length) throws IOException { + final long offset = reply.getStreamOffset(); final ChannelFuture future = ctx.writeAndFlush(reply); try { future.await(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); - throw new InterruptedIOException( - "Interrupted while writing " + length + " bytes at offset " + streamOffset); + throw new InterruptedIOException("Interrupted while writing " + length + " bytes at offset " + offset); } if (!future.isSuccess()) { - throw new IOException("Failed to write " + length + " bytes at offset " + streamOffset, future.cause()); + throw new IOException("Failed to write " + length + " bytes at offset " + offset, future.cause()); } - streamOffset += length; - return length; } private synchronized DataStreamReplyByteBuffer newReply(ByteBuffer buffer) { @@ -102,6 +112,23 @@ private synchronized DataStreamReplyByteBuffer newReply(ByteBuffer buffer) { .setBytesWritten(buffer.remaining()) .build(); } + + private synchronized DataStreamReplyByteBuffer newTerminalReply() { + final RaftClientReply reply = RaftClientReply.newBuilder() + .setRequest(request) + .setSuccess() + .build(); + final ByteBuffer buffer = toRaftClientReplyProto(reply).toByteString().asReadOnlyByteBuffer(); + return DataStreamReplyByteBuffer.newBuilder() + .setClientId(clientId) + .setType(Type.STREAM_HEADER) + .setStreamId(streamId) + .setStreamOffset(0) + .setBuffer(buffer) + .setSuccess(true) + .setBytesWritten(0) + .build(); + } } private final RaftServer server; @@ -146,7 +173,7 @@ private boolean processImpl(DataStreamRequestByteBuf requestBuf, ChannelHandlerC return true; } - final ReadStream stream = new ReadStream(requestBuf, ctx); + final ReadStream stream = new ReadStream(requestBuf, request, ctx); division.getStateMachine().data().query(request.getMessage(), stream); return true; } diff --git a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java new file mode 100644 index 0000000000..d06264184e --- /dev/null +++ b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java @@ -0,0 +1,118 @@ +/* + * 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.client.impl; + +import org.apache.ratis.client.DataStreamClient; +import org.apache.ratis.client.DataStreamClientRpc; +import org.apache.ratis.client.RaftClient; +import org.apache.ratis.client.RaftClientRpc; +import org.apache.ratis.client.api.DataStreamInput; +import org.apache.ratis.conf.RaftProperties; +import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; +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.RaftClientReply; +import org.apache.ratis.protocol.RaftClientRequest; +import org.apache.ratis.protocol.RaftGroup; +import org.apache.ratis.protocol.RaftGroupId; +import org.apache.ratis.protocol.RaftPeer; +import org.apache.ratis.retry.RetryPolicies; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Collection; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; + +public class TestDataStreamClientImpl { + private static RaftPeer newPeer(String id) { + return RaftPeer.newBuilder().setId(id).build(); + } + + private static class NoOpRaftClientRpc implements RaftClientRpc { + @Override + public void addRaftPeers(Collection peers) { + } + + @Override + public RaftClientReply sendRequest(RaftClientRequest request) throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public void close() { + } + } + + private static class RecordingDataStreamClientRpc implements DataStreamClientRpc { + private final AtomicReference request = new AtomicReference<>(); + + @Override + public CompletableFuture streamAsync( + DataStreamRequest dataStreamRequest, Consumer replyConsumer) { + try { + final ByteBuffer buffer = ((DataStreamRequestByteBuffer) dataStreamRequest).slice(); + request.set(ClientProtoUtils.toRaftClientRequest(RaftClientRequestProto.parseFrom(buffer))); + } catch (Exception e) { + throw new IllegalStateException(e); + } + return new CompletableFuture<>(); + } + + RaftClientRequest getRequest() { + return request.get(); + } + + @Override + public void close() { + } + } + + @Test + public void testStreamReadOnlyUsesPrimaryDataStreamServer() throws Exception { + final RaftProperties properties = new RaftProperties(); + final RaftPeer leader = newPeer("leader"); + final RaftPeer follower = newPeer("follower"); + final RaftGroup group = RaftGroup.valueOf(RaftGroupId.randomId(), leader, follower); + final RecordingDataStreamClientRpc dataStreamClientRpc = new RecordingDataStreamClientRpc(); + + try (RaftClient client = RaftClient.newBuilder() + .setClientId(ClientId.randomId()) + .setRaftGroup(group) + .setLeaderId(leader.getId()) + .setPrimaryDataStreamServer(follower) + .setClientRpc(new NoOpRaftClientRpc()) + .setProperties(properties) + .setRetryPolicy(RetryPolicies.noRetry()) + .build(); + DataStreamClient dataStreamClient = DataStreamClient.newBuilder(client) + .setDataStreamServer(follower) + .setDataStreamClientRpc(dataStreamClientRpc) + .setProperties(properties) + .build(); + DataStreamInput input = dataStreamClient.streamReadOnly(ByteBuffer.wrap(new byte[] {1}))) { + Assertions.assertNotNull(input); + Assertions.assertEquals(follower.getId(), dataStreamClientRpc.getRequest().getServerId()); + } + } +} diff --git a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java index dabc93dda2..01972e7d67 100644 --- a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java +++ b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java @@ -18,20 +18,27 @@ package org.apache.ratis.datastream; import org.apache.ratis.BaseTest; +import org.apache.ratis.client.api.DataStreamInput; import org.apache.ratis.io.StandardWriteOption; import org.apache.ratis.protocol.RaftPeer; import org.apache.ratis.protocol.RoutingTable; import org.apache.ratis.server.impl.MiniRaftCluster; import org.apache.ratis.client.RaftClient; +import org.apache.ratis.client.impl.ClientProtoUtils; import org.apache.ratis.client.impl.DataStreamClientImpl.DataStreamOutputImpl; import org.apache.ratis.datastream.DataStreamTestUtils.MultiDataStreamStateMachine; import org.apache.ratis.datastream.DataStreamTestUtils.SingleDataStream; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; import org.apache.ratis.proto.RaftProtos.ReplicationLevel; import org.apache.ratis.protocol.DataStreamReply; import org.apache.ratis.protocol.RaftClientReply; import org.apache.ratis.protocol.RaftClientRequest; +import org.apache.ratis.retry.RetryPolicies; import org.apache.ratis.server.RaftServer; +import org.apache.ratis.server.RaftServerConfigKeys; +import org.apache.ratis.thirdparty.com.google.protobuf.ByteString; import org.apache.ratis.util.CollectionUtils; import org.apache.ratis.util.FileUtils; import org.apache.ratis.util.Timestamp; @@ -39,12 +46,16 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.io.EOFException; import java.io.File; +import java.nio.ByteBuffer; import java.nio.channels.FileChannel; import java.nio.file.StandardOpenOption; import java.util.Collection; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; import static org.apache.ratis.RaftTestUtil.waitForLeader; @@ -71,6 +82,17 @@ public void testStreamWithInvalidRoutingTable() throws Exception { runWithNewCluster(NUM_SERVERS, this::runTestInvalidPrimaryInRoutingTable); } + @Test + public void testStreamReadOnly() throws Exception { + runWithNewCluster(NUM_SERVERS, this::runTestStreamReadOnly); + } + + @Test + public void testStreamReadOnlyWithNonLeaderPrimary() throws Exception { + RaftServerConfigKeys.Read.setOption(getProperties(), RaftServerConfigKeys.Read.Option.LINEARIZABLE); + runWithNewCluster(NUM_SERVERS, this::runTestStreamReadOnlyWithNonLeaderPrimary); + } + void testStreamWrites(CLUSTER cluster) throws Exception { waitForLeader(cluster); runTestDataStreamOutput(cluster); @@ -105,6 +127,67 @@ void runTestDataStreamOutput(CLUSTER cluster) throws Exception { assertLogEntry(cluster, request); } + void runTestStreamReadOnly(CLUSTER cluster) throws Exception { + final RaftPeer leader = waitForLeader(cluster).getPeer(); + runTestStreamReadOnly(cluster, leader, leader); + } + + void runTestStreamReadOnlyWithNonLeaderPrimary(CLUSTER cluster) throws Exception { + final RaftPeer leader = waitForLeader(cluster).getPeer(); + RaftPeer nonLeaderPrimary = null; + for (RaftPeer peer : cluster.getGroup().getPeers()) { + if (!peer.getId().equals(leader.getId())) { + nonLeaderPrimary = peer; + break; + } + } + Assertions.assertNotNull(nonLeaderPrimary, "Cannot find non-leader peer"); + runTestStreamReadOnly(cluster, leader, nonLeaderPrimary); + } + + void runTestStreamReadOnly(CLUSTER cluster, RaftPeer leader, RaftPeer primaryServer) throws Exception { + final ByteString query = ByteString.copyFromUtf8("stream-read-only"); + try (RaftClient client = cluster.createClient(leader.getId(), cluster.getGroup(), RetryPolicies.noRetry(), + primaryServer); + DataStreamInput in = client.getDataStreamApi().streamReadOnly(query.asReadOnlyByteBuffer())) { + for (int i = 0; i < MultiDataStreamStateMachine.READ_ONLY_STREAM_CHUNKS; i++) { + final ByteString chunk = MultiDataStreamStateMachine.getReadOnlyStreamChunk(query, i); + final DataStreamReply data = in.readAsync().join(); + DataStreamTestUtils.assertSuccessReply(Type.STREAM_DATA, chunk.size(), data); + Assertions.assertEquals(chunk, toByteString(data)); + } + + final DataStreamReply reply = in.readAsync().join(); + DataStreamTestUtils.assertSuccessReply(Type.STREAM_HEADER, 0, reply); + + try { + final RaftClientReply clientReply = ClientProtoUtils.getRaftClientReply(reply); + Assertions.assertTrue(clientReply.isSuccess()); + Assertions.assertEquals(primaryServer.getId(), clientReply.getServerId()); + } finally { + reply.release(); + } + + final ExecutionException eof = Assertions.assertThrows(ExecutionException.class, + () -> in.readAsync().get(1, TimeUnit.SECONDS)); + Assertions.assertInstanceOf(EOFException.class, eof.getCause()); + } + } + + private static ByteString toByteString(DataStreamReply reply) { + try { + if (reply instanceof DataStreamReplyByteBuffer) { + final ByteBuffer buffer = ((DataStreamReplyByteBuffer) reply).slice(); + return ByteString.copyFrom(buffer); + } else if (reply instanceof DataStreamReplyByteBuf) { + return ByteString.copyFrom(((DataStreamReplyByteBuf) reply).slice().nioBuffer()); + } + throw new AssertionError("Unexpected reply " + reply); + } finally { + reply.release(); + } + } + void runTestInvalidPrimaryInRoutingTable(CLUSTER cluster) throws Exception { final RaftPeer primaryServer = CollectionUtils.random(cluster.getGroup().getPeers()); diff --git a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java index 1573a2a283..516f8c1714 100644 --- a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java +++ b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java @@ -124,12 +124,14 @@ public DataApi data() { for (Object outbound; (outbound = embeddedChannel.readOutbound()) != null;) { replies.add((DataStreamReply) outbound); } - assertEquals(1, replies.size()); + assertEquals(2, replies.size()); }, 10, TimeDuration.valueOf(100, TimeUnit.MILLISECONDS), "read-only replies", null); assertEquals(query, messageRef.get().getContent()); assertFalse(streamRef.get().isOpen(), "state machine should close the streaming query channel"); 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(); } From 357cb71e7a4c379d78ecfae194612c6087a05160 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 13 Jun 2026 21:56:01 +0800 Subject: [PATCH 02/13] stream client handling Signed-off-by: peterxcli --- .../client/impl/DataStreamClientImpl.java | 132 +++++++++++++++- .../netty/client/NettyClientStreamRpc.java | 147 +++++++++++++++++- .../client/impl/TestDataStreamClientImpl.java | 82 +++++----- 3 files changed, 316 insertions(+), 45 deletions(-) diff --git a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java index 313131cbda..03fb253163 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java @@ -23,9 +23,11 @@ import org.apache.ratis.client.DataStreamClientRpc; import org.apache.ratis.client.DataStreamOutputRpc; import org.apache.ratis.client.RaftClient; +import org.apache.ratis.client.api.DataStreamInput; import org.apache.ratis.conf.RaftProperties; 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.FilePositionCount; import org.apache.ratis.io.StandardWriteOption; import org.apache.ratis.io.WriteOption; @@ -34,16 +36,17 @@ import org.apache.ratis.protocol.ClientInvocationId; import org.apache.ratis.protocol.DataStreamReply; import org.apache.ratis.protocol.DataStreamRequestHeader; +import org.apache.ratis.protocol.Message; import org.apache.ratis.protocol.RaftClientReply; import org.apache.ratis.protocol.RaftClientRequest; import org.apache.ratis.protocol.RaftGroupId; import org.apache.ratis.protocol.RaftPeer; +import org.apache.ratis.protocol.RoutingTable; import org.apache.ratis.protocol.exceptions.AlreadyClosedException; import org.apache.ratis.rpc.CallId; +import org.apache.ratis.thirdparty.com.google.protobuf.ByteString; import org.apache.ratis.thirdparty.io.netty.buffer.ByteBuf; import org.apache.ratis.util.IOUtils; -import org.apache.ratis.protocol.*; -import org.apache.ratis.thirdparty.com.google.protobuf.ByteString; import org.apache.ratis.util.JavaUtils; import org.apache.ratis.util.MemoizedSupplier; import org.apache.ratis.util.Preconditions; @@ -51,13 +54,16 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.io.EOFException; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.WritableByteChannel; +import java.util.ArrayDeque; import java.util.Arrays; import java.util.Collections; import java.util.Objects; import java.util.Optional; +import java.util.Queue; import java.util.concurrent.CompletableFuture; /** @@ -237,6 +243,113 @@ private CompletableFuture sendForward(DataStreamReply writeRepl } } + public final class DataStreamInputImpl implements DataStreamInput { + private final RaftClientRequest header; + private final CompletableFuture replyFuture; + private final Queue replies = new ArrayDeque<>(); + private final Queue> pendingReads = new ArrayDeque<>(); + private Throwable readException; + private boolean endOfStream; + private boolean closed; + + private DataStreamInputImpl(RaftClientRequest request) { + this.header = request; + final ByteBuffer buffer = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(header); + final DataStreamRequestHeader h = new DataStreamRequestHeader(header.getClientId(), Type.STREAM_HEADER, + header.getCallId(), 0, buffer.remaining(), StandardWriteOption.FLUSH, StandardWriteOption.CLOSE); + this.replyFuture = dataStreamClientRpc.streamAsync(new DataStreamRequestByteBuffer(h, buffer), this::receive); + replyFuture.whenComplete((reply, exception) -> { + if (exception != null) { + failReads(exception); + } else { + markEndOfStream(); + } + }); + } + + private void receive(DataStreamReply reply) { + for (;;) { + final CompletableFuture pending; + synchronized (this) { + if (closed) { + reply.release(); + return; + } + pending = pendingReads.poll(); + if (pending == null) { + replies.add(reply); + return; + } + } + if (pending.complete(reply)) { + return; + } + } + } + + private void failReads(Throwable t) { + for (;;) { + final CompletableFuture pending; + synchronized (this) { + readException = t; + pending = pendingReads.poll(); + if (pending == null) { + return; + } + } + pending.completeExceptionally(t); + } + } + + private EOFException newEndOfStreamException() { + return new EOFException(clientId + ": end of stream, request=" + header); + } + + private void markEndOfStream() { + for (;;) { + final CompletableFuture pending; + synchronized (this) { + endOfStream = true; + pending = pendingReads.poll(); + if (pending == null) { + return; + } + } + pending.completeExceptionally(newEndOfStreamException()); + } + } + + @Override + public synchronized CompletableFuture readAsync() { + if (closed) { + return JavaUtils.completeExceptionally(new AlreadyClosedException( + clientId + ": stream already closed, request=" + header)); + } + final DataStreamReply reply = replies.poll(); + if (reply != null) { + return CompletableFuture.completedFuture(reply); + } + if (readException != null) { + return JavaUtils.completeExceptionally(readException); + } + if (endOfStream) { + return JavaUtils.completeExceptionally(newEndOfStreamException()); + } + final CompletableFuture f = new CompletableFuture<>(); + pendingReads.add(f); + return f; + } + + @Override + public synchronized void close() { + closed = true; + for (DataStreamReply reply; (reply = replies.poll()) != null;) { + reply.release(); + } + failReads(new AlreadyClosedException(clientId + ": stream already closed, request=" + header)); + } + } + @Override public DataStreamClientRpc getClientRpc() { return dataStreamClientRpc; @@ -274,6 +387,21 @@ public DataStreamOutputRpc stream(ByteBuffer headerMessage, RoutingTable routing return new DataStreamOutputImpl(request); } + @Override + public DataStreamInput streamReadOnly(ByteBuffer headerMessage) { + final Message message = + Optional.ofNullable(headerMessage).map(ByteString::copyFrom).map(Message::valueOf).orElse(null); + final RaftClientRequest request = RaftClientRequest.newBuilder() + .setClientId(clientId) + .setServerId(dataStreamServer.getId()) + .setGroupId(groupId) + .setCallId(CallId.getAndIncrement()) + .setMessage(message) + .setType(RaftClientRequest.readRequestType()) + .build(); + return new DataStreamInputImpl(request); + } + @Override public void close() throws IOException { dataStreamClientRpc.close(); diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java index 5b673d5189..59b62eff82 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java @@ -23,6 +23,7 @@ import org.apache.ratis.conf.RaftProperties; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuf; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.datastream.impl.DataStreamRequestFilePositionCount; import org.apache.ratis.io.StandardWriteOption; import org.apache.ratis.io.WriteOption; @@ -55,6 +56,7 @@ import org.apache.ratis.thirdparty.io.netty.handler.codec.ByteToMessageDecoder; import org.apache.ratis.thirdparty.io.netty.handler.codec.MessageToMessageEncoder; import org.apache.ratis.thirdparty.io.netty.handler.ssl.SslContext; +import org.apache.ratis.thirdparty.io.netty.util.concurrent.ScheduledFuture; import org.apache.ratis.util.JavaUtils; import org.apache.ratis.util.MemoizedFunction; import org.apache.ratis.util.MemoizedSupplier; @@ -73,10 +75,13 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; @@ -320,10 +325,66 @@ synchronized boolean shouldFlush(int countMin, SizeInBytes bytesMin, DataStreamR } } + static class ReadOnlyStreamingReply { + private final NettyClientReplies.RequestEntry terminalEntry; + private final CompletableFuture replyFuture; + private final Consumer replyConsumer; + private Supplier> timeoutScheduler; + private ScheduledFuture timeoutFuture; + + ReadOnlyStreamingReply(DataStreamRequest request, CompletableFuture replyFuture, + Consumer replyConsumer) { + this.terminalEntry = new NettyClientReplies.RequestEntry(request); + this.replyFuture = replyFuture; + this.replyConsumer = replyConsumer; + } + + synchronized boolean receiveReply(DataStreamReply reply) { + NettyClientReplies.ReplyEntry.cancel(timeoutFuture); + final boolean terminal = !reply.isSuccess() || terminalEntry.equals(new NettyClientReplies.RequestEntry(reply)); + final DataStreamReply replyToComplete = terminal && reply instanceof DataStreamReplyByteBuf ? + NettyDataStreamUtils.toDataStreamReplyByteBuffer((DataStreamReplyByteBuf) reply) : reply; + try { + replyConsumer.accept(replyToComplete); + } catch (Throwable t) { + if (replyToComplete == reply) { + reply.release(); + } + completeExceptionally(t); + return true; + } + + if (terminal) { + replyFuture.complete(replyToComplete); + return true; + } + scheduleTimeout(); + return false; + } + + synchronized void completeExceptionally(Throwable t) { + NettyClientReplies.ReplyEntry.cancel(timeoutFuture); + replyFuture.completeExceptionally(t); + } + + synchronized void scheduleTimeout(Supplier> scheduleMethod) { + timeoutScheduler = scheduleMethod; + scheduleTimeout(); + } + + private void scheduleTimeout() { + if (!replyFuture.isDone() && timeoutScheduler != null) { + timeoutFuture = timeoutScheduler.get(); + } + } + } + private final String name; private final Connection connection; private final NettyClientReplies replies = new NettyClientReplies(); + private final ConcurrentMap readOnlyStreamingReplies + = new ConcurrentHashMap<>(); private final TimeDuration requestTimeout; private final TimeDuration closeTimeout; @@ -361,17 +422,46 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { LOG.debug("{}: read {}", name, reply); final ClientInvocationId clientInvocationId = ClientInvocationId.valueOf( reply.getClientId(), reply.getStreamId()); + final ReadOnlyStreamingReply readOnlyStreamingReply = readOnlyStreamingReplies.get(clientInvocationId); + if (readOnlyStreamingReply != null) { + try { + if (readOnlyStreamingReply.receiveReply(reply)) { + readOnlyStreamingReplies.remove(clientInvocationId, readOnlyStreamingReply); + } + } catch (Throwable cause) { + LOG.warn("{} : channelRead error:", name, cause); + readOnlyStreamingReplies.remove(clientInvocationId, readOnlyStreamingReply); + readOnlyStreamingReply.completeExceptionally(cause); + } + return; + } + final NettyClientReplies.ReplyMap replyMap = replies.getReplyMap(clientInvocationId); if (replyMap == null) { LOG.error("{}: {} replyMap not found for reply: {}", name, clientInvocationId, reply); + reply.release(); return; } + final DataStreamReply replyToReceive; try { - replyMap.receiveReply(reply); + replyToReceive = reply instanceof DataStreamReplyByteBuf ? + NettyDataStreamUtils.toDataStreamReplyByteBuffer((DataStreamReplyByteBuf) reply) : reply; } catch (Throwable cause) { - LOG.warn("{} : channelRead error:", name, cause); - replyMap.completeExceptionally(cause); + try (DataStreamReply replyToClose = reply) { + LOG.warn("{} : channelRead error for {}:", name, replyToClose.getClass().getSimpleName(), cause); + replyMap.completeExceptionally(cause); + } + return; + } + + try { + replyMap.receiveReply(replyToReceive); + } catch (Throwable cause) { + try (DataStreamReply replyToClose = replyToReceive) { + LOG.warn("{} : channelRead error for {}:", name, replyToClose.getClass().getSimpleName(), cause); + replyMap.completeExceptionally(cause); + } } } @@ -456,7 +546,7 @@ static ByteToMessageDecoder newDecoder() { @Override protected void decode(ChannelHandlerContext context, ByteBuf buf, List out) { - Optional.ofNullable(NettyDataStreamUtils.decodeDataStreamReplyByteBuffer(buf)).ifPresent(out::add); + Optional.ofNullable(NettyDataStreamUtils.decodeDataStreamReplyByteBuf(buf)).ifPresent(out::add); } }; } @@ -507,6 +597,55 @@ public CompletableFuture streamAsync(DataStreamRequest request) return f; } + @Override + public CompletableFuture streamAsync( + DataStreamRequest request, Consumer replyConsumer) { + final CompletableFuture f = new CompletableFuture<>(); + final ClientInvocationId clientInvocationId = ClientInvocationId.valueOf(request.getClientId(), + request.getStreamId()); + final ReadOnlyStreamingReply replyEntry = new ReadOnlyStreamingReply(request, f, replyConsumer); + if (readOnlyStreamingReplies.putIfAbsent(clientInvocationId, replyEntry) != null) { + f.completeExceptionally(new AlreadyClosedException(this + ": A read-only stream already exists for " + + clientInvocationId)); + return f; + } + + final ChannelFuture channelFuture; + final Channel channel; + LOG.debug("{}: write read-only stream begin {}", this, request); + synchronized (replyEntry) { + channel = connection.getChannelUninterruptibly(); + if (channel == null) { + readOnlyStreamingReplies.remove(clientInvocationId, replyEntry); + f.completeExceptionally(new AlreadyClosedException(this + ": Failed to send " + request)); + return f; + } + final Function writeMethod = outstandingRequests.shouldFlush( + flushRequestCountMin, flushRequestBytesMin, request)? channel::writeAndFlush: channel::write; + channelFuture = writeMethod.apply(request); + } + channelFuture.addListener(future -> { + if (!future.isSuccess()) { + readOnlyStreamingReplies.remove(clientInvocationId, replyEntry); + final IOException e = new IOException(this + ": Failed to send " + request + " to " + channel.remoteAddress(), + future.cause()); + replyEntry.completeExceptionally(e); + LOG.error("Channel write failed", e); + } else { + LOG.debug("{}: write read-only stream after {}", this, request); + + replyEntry.scheduleTimeout(() -> channel.eventLoop().schedule(() -> { + if (!f.isDone()) { + readOnlyStreamingReplies.remove(clientInvocationId, replyEntry); + replyEntry.completeExceptionally(new TimeoutIOException( + "Timeout " + requestTimeout + ": Failed to send " + request + " via channel " + channel)); + } + }, requestTimeout.getDuration(), requestTimeout.getUnit())); + } + }); + return f; + } + @Override public void close() { final boolean flush = outstandingRequests.shouldFlush(0, SizeInBytes.ZERO, null); diff --git a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java index d06264184e..498572c0df 100644 --- a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java +++ b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java @@ -19,27 +19,22 @@ import org.apache.ratis.client.DataStreamClient; import org.apache.ratis.client.DataStreamClientRpc; -import org.apache.ratis.client.RaftClient; -import org.apache.ratis.client.RaftClientRpc; import org.apache.ratis.client.api.DataStreamInput; import org.apache.ratis.conf.RaftProperties; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; +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.RaftClientReply; import org.apache.ratis.protocol.RaftClientRequest; -import org.apache.ratis.protocol.RaftGroup; import org.apache.ratis.protocol.RaftGroupId; import org.apache.ratis.protocol.RaftPeer; -import org.apache.ratis.retry.RetryPolicies; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; -import java.io.IOException; import java.nio.ByteBuffer; -import java.util.Collection; import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; @@ -49,23 +44,9 @@ private static RaftPeer newPeer(String id) { return RaftPeer.newBuilder().setId(id).build(); } - private static class NoOpRaftClientRpc implements RaftClientRpc { - @Override - public void addRaftPeers(Collection peers) { - } - - @Override - public RaftClientReply sendRequest(RaftClientRequest request) throws IOException { - throw new UnsupportedOperationException(); - } - - @Override - public void close() { - } - } - private static class RecordingDataStreamClientRpc implements DataStreamClientRpc { private final AtomicReference request = new AtomicReference<>(); + private final AtomicReference> replyConsumer = new AtomicReference<>(); @Override public CompletableFuture streamAsync( @@ -76,6 +57,7 @@ public CompletableFuture streamAsync( } catch (Exception e) { throw new IllegalStateException(e); } + this.replyConsumer.set(replyConsumer); return new CompletableFuture<>(); } @@ -83,36 +65,58 @@ RaftClientRequest getRequest() { return request.get(); } + void receive(DataStreamReply reply) { + replyConsumer.get().accept(reply); + } + @Override public void close() { } } + private static DataStreamClient newDataStreamClient( + RaftPeer dataStreamServer, RecordingDataStreamClientRpc dataStreamClientRpc) { + final RaftProperties properties = new RaftProperties(); + return new DataStreamClientImpl( + ClientId.randomId(), RaftGroupId.randomId(), dataStreamServer, dataStreamClientRpc, properties); + } + @Test public void testStreamReadOnlyUsesPrimaryDataStreamServer() throws Exception { - final RaftProperties properties = new RaftProperties(); - final RaftPeer leader = newPeer("leader"); final RaftPeer follower = newPeer("follower"); - final RaftGroup group = RaftGroup.valueOf(RaftGroupId.randomId(), leader, follower); final RecordingDataStreamClientRpc dataStreamClientRpc = new RecordingDataStreamClientRpc(); - try (RaftClient client = RaftClient.newBuilder() - .setClientId(ClientId.randomId()) - .setRaftGroup(group) - .setLeaderId(leader.getId()) - .setPrimaryDataStreamServer(follower) - .setClientRpc(new NoOpRaftClientRpc()) - .setProperties(properties) - .setRetryPolicy(RetryPolicies.noRetry()) - .build(); - DataStreamClient dataStreamClient = DataStreamClient.newBuilder(client) - .setDataStreamServer(follower) - .setDataStreamClientRpc(dataStreamClientRpc) - .setProperties(properties) - .build(); + try (DataStreamClient dataStreamClient = newDataStreamClient(follower, dataStreamClientRpc); DataStreamInput input = dataStreamClient.streamReadOnly(ByteBuffer.wrap(new byte[] {1}))) { Assertions.assertNotNull(input); Assertions.assertEquals(follower.getId(), dataStreamClientRpc.getRequest().getServerId()); } } + + @Test + public void testReceiveSkipsCancelledPendingRead() throws Exception { + final RaftPeer follower = newPeer("follower"); + final RecordingDataStreamClientRpc dataStreamClientRpc = new RecordingDataStreamClientRpc(); + + try (DataStreamClient dataStreamClient = newDataStreamClient(follower, dataStreamClientRpc); + DataStreamInput input = dataStreamClient.streamReadOnly(ByteBuffer.wrap(new byte[] {1}))) { + final CompletableFuture cancelled = input.readAsync(); + final CompletableFuture active = input.readAsync(); + cancelled.cancel(false); + + final DataStreamReply reply = DataStreamReplyByteBuffer.newBuilder() + .setClientId(ClientId.randomId()) + .setType(Type.STREAM_DATA) + .setStreamId(1) + .setStreamOffset(0) + .setBuffer(ByteBuffer.allocate(0)) + .setSuccess(true) + .build(); + dataStreamClientRpc.receive(reply); + + Assertions.assertTrue(active.isDone()); + Assertions.assertSame(reply, active.getNow(null)); + reply.release(); + } + } } From 65e6b83568a902f059799349ae396259a2534ff4 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 13 Jun 2026 22:25:51 +0800 Subject: [PATCH 03/13] open another thread for blocking sm write Signed-off-by: peterxcli --- .../netty/server/NettyServerStreamRpc.java | 6 ++ .../netty/server/ReadStreamManagement.java | 28 +++++++- .../server/TestDataStreamManagement.java | 65 +++++++++++++++++++ 3 files changed, 98 insertions(+), 1 deletion(-) diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/server/NettyServerStreamRpc.java b/ratis-netty/src/main/java/org/apache/ratis/netty/server/NettyServerStreamRpc.java index 3aff714da2..95cbd9aafb 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/server/NettyServerStreamRpc.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/server/NettyServerStreamRpc.java @@ -332,6 +332,12 @@ public void close() { LOG.error(this + ": Failed to shutdown request service.", e); } + try { + reads.shutdown(); + } catch (Exception e) { + LOG.error(this + ": Failed to shutdown read service.", e); + } + try { channelFuture.channel().close().sync(); bossGroup.shutdownGracefully(0, 100, TimeUnit.MILLISECONDS); diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java b/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java index 4bdb807ff6..4c0c81c901 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java @@ -17,6 +17,7 @@ */ package org.apache.ratis.netty.server; +import org.apache.ratis.conf.RaftProperties; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuf; import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; @@ -27,10 +28,13 @@ import org.apache.ratis.protocol.RaftClientRequest; import org.apache.ratis.protocol.exceptions.AlreadyClosedException; import org.apache.ratis.server.RaftServer; +import org.apache.ratis.server.RaftServerConfigKeys; 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; +import org.apache.ratis.util.ConcurrentUtils; import org.apache.ratis.util.JavaUtils; +import org.apache.ratis.util.TimeDuration; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -39,6 +43,7 @@ import java.nio.ByteBuffer; import java.nio.channels.WritableByteChannel; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; import static org.apache.ratis.client.impl.ClientProtoUtils.toRaftClientRequest; import static org.apache.ratis.client.impl.ClientProtoUtils.toRaftClientReplyProto; @@ -133,10 +138,25 @@ private synchronized DataStreamReplyByteBuffer newTerminalReply() { private final RaftServer server; private final String name; + private final ExecutorService requestExecutor; ReadStreamManagement(RaftServer server) { this.server = server; this.name = server.getId() + "-" + JavaUtils.getClassSimpleName(getClass()); + + final RaftProperties properties = server.getProperties(); + final boolean useCachedThreadPool = + RaftServerConfigKeys.DataStream.asyncRequestThreadPoolCached(properties); + this.requestExecutor = ConcurrentUtils.newThreadPoolWithMax( + useCachedThreadPool, + RaftServerConfigKeys.DataStream.asyncRequestThreadPoolSize(properties), + name + "-request-"); + } + + void shutdown() { + ConcurrentUtils.shutdownAndWait(TimeDuration.ONE_SECOND, requestExecutor, + timeout -> LOG.warn("{}: requestExecutor shutdown timeout in {}", + this, timeout)); } boolean process(DataStreamRequestByteBuf requestBuf, ChannelHandlerContext ctx) { @@ -174,7 +194,13 @@ private boolean processImpl(DataStreamRequestByteBuf requestBuf, ChannelHandlerC } final ReadStream stream = new ReadStream(requestBuf, request, ctx); - division.getStateMachine().data().query(request.getMessage(), stream); + requestExecutor.execute(() -> { + try { + division.getStateMachine().data().query(request.getMessage(), stream); + } catch (Throwable t) { + LOG.error("{}: Failed read-only data stream query for {}", this, request, t); + } + }); return true; } diff --git a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java index 516f8c1714..1605d18fcf 100644 --- a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java +++ b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java @@ -42,7 +42,10 @@ import org.apache.ratis.thirdparty.io.netty.channel.ChannelHandlerContext; import org.apache.ratis.thirdparty.io.netty.channel.ChannelId; import org.apache.ratis.thirdparty.io.netty.channel.ChannelInboundHandlerAdapter; +import org.apache.ratis.thirdparty.io.netty.channel.EventLoop; +import org.apache.ratis.thirdparty.io.netty.channel.EventLoopGroup; import org.apache.ratis.thirdparty.io.netty.channel.embedded.EmbeddedChannel; +import org.apache.ratis.thirdparty.io.netty.channel.nio.NioEventLoopGroup; import org.apache.ratis.util.JavaUtils; import org.apache.ratis.util.TimeDuration; import org.apache.ratis.util.function.CheckedBiFunction; @@ -56,7 +59,9 @@ import java.util.Collections; import java.util.List; import java.util.Set; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -137,6 +142,66 @@ public DataApi data() { } } + @Test + void readOnlyQueryDoesNotRunOnNettyEventLoop() throws Exception { + final RaftPeerId serverId = RaftPeerId.valueOf("s1"); + final ClientId clientId = ClientId.randomId(); + final RaftGroupId groupId = RaftGroupId.randomId(); + final CountDownLatch queryDone = new CountDownLatch(1); + final AtomicBoolean queryInEventLoop = new AtomicBoolean(); + final EventLoopGroup eventLoopGroup = new NioEventLoopGroup(1); + final EventLoop eventLoop = eventLoopGroup.next(); + final EmbeddedChannel embeddedChannel = new EmbeddedChannel(new ChannelInboundHandlerAdapter()); + final ChannelHandlerContext ctx = embeddedChannel.pipeline().firstContext(); + assertNotNull(ctx, "ChannelHandlerContext should be initialized"); + + final DataApi dataApi = new DataApi() { + @Override + public void query(Message request, WritableByteChannel stream) { + queryInEventLoop.set(eventLoop.inEventLoop()); + queryDone.countDown(); + } + }; + final StateMachine stateMachine = new BaseStateMachine() { + @Override + public DataApi data() { + return dataApi; + } + }; + final RaftServer server = newRaftServer(serverId, new RaftProperties(), groupId, newDivision(stateMachine)); + final ReadStreamManagement management = new ReadStreamManagement(server); + + final RaftClientRequest raftClientRequest = RaftClientRequest.newBuilder() + .setClientId(clientId) + .setServerId(serverId) + .setGroupId(groupId) + .setCallId(1L) + .setMessage(Message.valueOf(ByteString.copyFromUtf8("query"))) + .setType(RaftClientRequest.readRequestType()) + .build(); + final ByteBuffer header = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(raftClientRequest); + final ByteBuf headerBuf = Unpooled.wrappedBuffer(header); + final DataStreamRequestByteBuf request = new DataStreamRequestByteBuf( + clientId, + Type.STREAM_HEADER, + raftClientRequest.getCallId(), + 0L, + Collections.singletonList(StandardWriteOption.FLUSH), + headerBuf); + + try { + eventLoop.submit(() -> assertTrue(management.process(request, ctx))).sync(); + + assertTrue(queryDone.await(10, TimeUnit.SECONDS)); + assertEquals(0, headerBuf.refCnt()); + assertFalse(queryInEventLoop.get(), "read-only state machine query should not run on Netty event loop"); + } finally { + embeddedChannel.finishAndReleaseAll(); + management.shutdown(); + eventLoopGroup.shutdownGracefully(0, 100, TimeUnit.MILLISECONDS).sync(); + } + } + @Test void readCleansChannelMapOnEarlyException() throws Exception { // Scenario: STREAM_DATA arrives without prior STREAM_HEADER, so readImpl fails early. From 7a37daf9a7e04d0ccf1bcd43519f4dc8f8d96243 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 18 Jun 2026 01:14:17 +0800 Subject: [PATCH 04/13] upd Signed-off-by: peterxcli --- .../main/java/org/apache/ratis/netty/NettyDataStreamUtils.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java b/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java index bd30339b33..17308bbec6 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java @@ -261,7 +261,7 @@ static DataStreamReplyByteBuf decodeDataStreamReplyByteBuf(ByteBuf buf) { .map(header -> checkHeader(header, buf)) .map(header -> DataStreamReplyByteBuf.newBuilder() .setDataStreamReplyHeader(header) - .setBuffer(decodeData(buf, header, ByteBuf::retainedSlice)) + .setBuf(decodeData(buf, header, ByteBuf::retainedSlice)) .build()) .orElse(null); } From 3a06f7701e1d3a5f20cee32fe1a19e66f934e72b Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 19 Jun 2026 03:06:34 +0800 Subject: [PATCH 05/13] handle the bytebuf release for bytebuf reply Signed-off-by: peterxcli --- .../ratis/client/impl/DataStreamClientImpl.java | 5 +++-- .../datastream/impl/DataStreamReplyByteBuf.java | 6 ++++++ .../org/apache/ratis/protocol/DataStreamReply.java | 13 +------------ .../ratis/netty/client/NettyClientStreamRpc.java | 2 +- .../ratis/client/impl/TestDataStreamClientImpl.java | 3 ++- .../ratis/datastream/DataStreamClusterTests.java | 4 ++-- 6 files changed, 15 insertions(+), 18 deletions(-) diff --git a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java index 1f89973159..0c7fd3e471 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java @@ -25,6 +25,7 @@ import org.apache.ratis.client.RaftClient; import org.apache.ratis.client.api.DataStreamInput; import org.apache.ratis.conf.RaftProperties; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.datastream.impl.DataStreamPacketByteBuffer; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; @@ -272,7 +273,7 @@ private void receive(DataStreamReply reply) { final CompletableFuture pending; synchronized (this) { if (closed) { - reply.release(); + DataStreamReplyByteBuf.release(reply); return; } pending = pendingReads.poll(); @@ -344,7 +345,7 @@ public synchronized CompletableFuture readAsync() { public synchronized void close() { closed = true; for (DataStreamReply reply; (reply = replies.poll()) != null;) { - reply.release(); + DataStreamReplyByteBuf.release(reply); } failReads(new AlreadyClosedException(clientId + ": stream already closed, request=" + header)); } diff --git a/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamReplyByteBuf.java b/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamReplyByteBuf.java index 2477b089b7..6ce707dc17 100644 --- a/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamReplyByteBuf.java +++ b/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamReplyByteBuf.java @@ -118,4 +118,10 @@ static ByteBuffer copy(ByteBuf buf) { buf.readBytes(bytes); return ByteBuffer.wrap(bytes); } + + public static void release(DataStreamReply reply) { + if (reply instanceof DataStreamReplyByteBuf) { + ((DataStreamReplyByteBuf) reply).release(); + } + } } diff --git a/ratis-common/src/main/java/org/apache/ratis/protocol/DataStreamReply.java b/ratis-common/src/main/java/org/apache/ratis/protocol/DataStreamReply.java index cf1f826867..3d4f53d378 100644 --- a/ratis-common/src/main/java/org/apache/ratis/protocol/DataStreamReply.java +++ b/ratis-common/src/main/java/org/apache/ratis/protocol/DataStreamReply.java @@ -22,7 +22,7 @@ import java.util.Collection; -public interface DataStreamReply extends DataStreamPacket, AutoCloseable { +public interface DataStreamReply extends DataStreamPacket { boolean isSuccess(); @@ -30,15 +30,4 @@ public interface DataStreamReply extends DataStreamPacket, AutoCloseable { /** @return the commit information when the reply is created. */ Collection getCommitInfos(); - - /** - * Release resources owned by this reply. - */ - default void release() { - } - - @Override - default void close() { - release(); - } } diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java index 352118afab..7eb2d9527a 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java @@ -348,7 +348,7 @@ synchronized boolean receiveReply(DataStreamReply reply) { replyConsumer.accept(replyToComplete); } catch (Throwable t) { if (replyToComplete == reply) { - reply.release(); + DataStreamReplyByteBuf.release(reply); } completeExceptionally(t); return true; diff --git a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java index 498572c0df..6726c05b79 100644 --- a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java +++ b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java @@ -21,6 +21,7 @@ import org.apache.ratis.client.DataStreamClientRpc; import org.apache.ratis.client.api.DataStreamInput; import org.apache.ratis.conf.RaftProperties; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; @@ -116,7 +117,7 @@ public void testReceiveSkipsCancelledPendingRead() throws Exception { Assertions.assertTrue(active.isDone()); Assertions.assertSame(reply, active.getNow(null)); - reply.release(); + DataStreamReplyByteBuf.release(reply); } } } diff --git a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java index 01972e7d67..a066a8f1a4 100644 --- a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java +++ b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java @@ -165,7 +165,7 @@ void runTestStreamReadOnly(CLUSTER cluster, RaftPeer leader, RaftPeer primarySer Assertions.assertTrue(clientReply.isSuccess()); Assertions.assertEquals(primaryServer.getId(), clientReply.getServerId()); } finally { - reply.release(); + DataStreamReplyByteBuf.release(reply); } final ExecutionException eof = Assertions.assertThrows(ExecutionException.class, @@ -184,7 +184,7 @@ private static ByteString toByteString(DataStreamReply reply) { } throw new AssertionError("Unexpected reply " + reply); } finally { - reply.release(); + DataStreamReplyByteBuf.release(reply); } } From bbbd445130dfb04959ee8884d2cd7620825a5e45 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 19 Jun 2026 03:44:53 +0800 Subject: [PATCH 06/13] minor Signed-off-by: peterxcli --- .../ratis/netty/NettyDataStreamUtils.java | 18 +----------------- .../ratis/netty/client/NettyClientReplies.java | 3 +-- .../netty/client/NettyClientStreamRpc.java | 17 +++++++---------- .../netty/server/ReadStreamManagement.java | 16 ++-------------- 4 files changed, 11 insertions(+), 43 deletions(-) diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java b/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java index b3b8dd907e..9f8cbecc8b 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java @@ -246,28 +246,12 @@ static DataStreamReplyByteBuf decodeDataStreamReplyByteBuf(ByteBuf buf) { static DataStreamReplyByteBuffer toDataStreamReplyByteBuffer(DataStreamReplyByteBuf reply) { try { - return DataStreamReplyByteBuffer.newBuilder() - .setDataStreamPacket(reply) - .setBuffer(copy(reply.slice())) - .setSuccess(reply.isSuccess()) - .setBytesWritten(reply.getBytesWritten()) - .setCommitInfos(reply.getCommitInfos()) - .build(); + return reply.copy(); } finally { reply.release(); } } - static DataStreamReplyByteBuf decodeDataStreamReplyByteBuf(ByteBuf buf) { - return Optional.ofNullable(decodeDataStreamReplyHeader(buf)) - .map(header -> checkHeader(header, buf)) - .map(header -> DataStreamReplyByteBuf.newBuilder() - .setDataStreamReplyHeader(header) - .setBuf(decodeData(buf, header, ByteBuf::retainedSlice)) - .build()) - .orElse(null); - } - static DataStreamReplyHeader decodeDataStreamReplyHeader(ByteBuf buf) { if (DataStreamPacketHeader.getSizeOfHeaderLen() > buf.readableBytes()) { return null; diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientReplies.java b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientReplies.java index 121ef3c0fc..0aa6bc7721 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientReplies.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientReplies.java @@ -18,7 +18,6 @@ package org.apache.ratis.netty.client; -import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; import org.apache.ratis.protocol.ClientInvocationId; import org.apache.ratis.protocol.DataStreamPacket; @@ -65,7 +64,7 @@ ReplyEntry submitRequest(RequestEntry requestEntry, boolean isClose, Completable return map.computeIfAbsent(requestEntry, r -> new ReplyEntry(isClose, f)); } - void receiveReply(DataStreamReplyByteBuf reply) { + void receiveReply(DataStreamReply reply) { final RequestEntry requestEntry = new RequestEntry(reply); final ReplyEntry replyEntry = map.remove(requestEntry); LOG.debug("remove: {}; replyEntry: {}; reply: {}", requestEntry, replyEntry, reply); diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java index 7eb2d9527a..80fb64678b 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java @@ -340,15 +340,15 @@ static class ReadOnlyStreamingReply { } synchronized boolean receiveReply(DataStreamReply reply) { - NettyClientReplies.ReplyEntry.cancel(timeoutFuture); + NettyUtils.cancel(timeoutFuture); final boolean terminal = !reply.isSuccess() || terminalEntry.equals(new NettyClientReplies.RequestEntry(reply)); - final DataStreamReply replyToComplete = terminal && reply instanceof DataStreamReplyByteBuf ? + final DataStreamReply replyToComplete = reply instanceof DataStreamReplyByteBuf ? NettyDataStreamUtils.toDataStreamReplyByteBuffer((DataStreamReplyByteBuf) reply) : reply; try { replyConsumer.accept(replyToComplete); } catch (Throwable t) { if (replyToComplete == reply) { - DataStreamReplyByteBuf.release(reply); + DataStreamReplyByteBuf.release(replyToComplete); } completeExceptionally(t); return true; @@ -363,7 +363,7 @@ synchronized boolean receiveReply(DataStreamReply reply) { } synchronized void completeExceptionally(Throwable t) { - NettyClientReplies.ReplyEntry.cancel(timeoutFuture); + NettyUtils.cancel(timeoutFuture); replyFuture.completeExceptionally(t); } @@ -450,13 +450,10 @@ private void process(DataStreamReplyByteBuf reply) { final DataStreamReply replyToReceive; try { - replyToReceive = reply instanceof DataStreamReplyByteBuf ? - NettyDataStreamUtils.toDataStreamReplyByteBuffer((DataStreamReplyByteBuf) reply) : reply; + replyToReceive = NettyDataStreamUtils.toDataStreamReplyByteBuffer(reply); } catch (Throwable cause) { - try (DataStreamReply replyToClose = reply) { - LOG.warn("{} : channelRead error for {}:", name, replyToClose.getClass().getSimpleName(), cause); - replyMap.completeExceptionally(cause); - } + LOG.warn("{} : channelRead error for {}:", name, reply.getClass().getSimpleName(), cause); + replyMap.completeExceptionally(cause); return; } diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java b/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java index 94413cb29e..6543d2693b 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java @@ -64,21 +64,9 @@ static class ReadStream implements WritableByteChannel { ReadStream(RaftClientRequest request, long streamId, ChannelHandlerContext ctx) { this.clientId = request.getClientId(); this.streamId = streamId; + this.request = request; this.ctx = ctx; - - final RaftClientReply reply = RaftClientReply.newBuilder() - .setRequest(request) - .setSuccess() - .build(); - this.terminalReply = DataStreamReplyByteBuffer.newBuilder() - .setClientId(clientId) - .setType(Type.STREAM_HEADER) - .setStreamId(streamId) - .setStreamOffset(0) - .setBuffer(toRaftClientReplyProto(reply).toByteString().asReadOnlyByteBuffer()) - .setSuccess(true) - .setBytesWritten(0) - .build(); + this.terminalReply = newTerminalReply(); } @Override From e8527fc28461049b5317757d04e0da28f39ec3fc Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 19 Jun 2026 04:05:25 +0800 Subject: [PATCH 07/13] Observer API Signed-off-by: peterxcli --- .../ratis/client/DataStreamClientRpc.java | 5 +- .../client/impl/DataStreamClientImpl.java | 3 +- .../ratis/datastream/DataStreamObserver.java | 30 ++++++++ .../netty/client/NettyClientStreamRpc.java | 36 +++++---- .../client/impl/TestDataStreamClientImpl.java | 24 +++--- ...tNettyClientStreamRpcReconnectBackoff.java | 74 +++++++++++++++++++ 6 files changed, 146 insertions(+), 26 deletions(-) create mode 100644 ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java diff --git a/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java b/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java index fd5bd9538e..11210de9ee 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java @@ -18,13 +18,14 @@ package org.apache.ratis.client; +import org.apache.ratis.datastream.DataStreamObserver; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.protocol.DataStreamReply; import org.apache.ratis.protocol.DataStreamRequest; import org.apache.ratis.util.JavaUtils; import java.io.Closeable; import java.util.concurrent.CompletableFuture; -import java.util.function.Consumer; /** * A client interface for sending stream requests. @@ -40,7 +41,7 @@ default CompletableFuture streamAsync(DataStreamRequest request /** Async call to send a request and receive multiple replies for the request. */ default CompletableFuture streamAsync( - DataStreamRequest request, Consumer replyConsumer) { + DataStreamRequest request, DataStreamObserver replyHandler) { throw new UnsupportedOperationException(getClass() + " does not support " + JavaUtils.getCurrentStackTraceElement().getMethodName()); } diff --git a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java index 0c7fd3e471..30ec9a6b6f 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java @@ -258,7 +258,8 @@ private DataStreamInputImpl(RaftClientRequest request) { final ByteBuffer buffer = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(header); final DataStreamRequestHeader h = new DataStreamRequestHeader(header.getClientId(), Type.STREAM_HEADER, header.getCallId(), 0, buffer.remaining(), StandardWriteOption.FLUSH, StandardWriteOption.CLOSE); - this.replyFuture = dataStreamClientRpc.streamAsync(new DataStreamRequestByteBuffer(h, buffer), this::receive); + this.replyFuture = dataStreamClientRpc.streamAsync(new DataStreamRequestByteBuffer(h, buffer), + reply -> receive(reply.copy())); replyFuture.whenComplete((reply, exception) -> { if (exception != null) { failReads(exception); diff --git a/ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java b/ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java new file mode 100644 index 0000000000..82529177cb --- /dev/null +++ b/ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java @@ -0,0 +1,30 @@ +/* + * 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.datastream; + +/** An interface similar to gRPC {@link org.apache.ratis.thirdparty.io.grpc.stub.StreamObserver}. */ +@FunctionalInterface +public interface DataStreamObserver { + void onNext(V value); + + default void onError(Throwable t) { + } + + default void onCompleted() { + } +} diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java index 80fb64678b..e1179ceed7 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java @@ -21,6 +21,7 @@ import org.apache.ratis.client.DataStreamClientRpc; import org.apache.ratis.client.RaftClientConfigKeys; import org.apache.ratis.conf.RaftProperties; +import org.apache.ratis.datastream.DataStreamObserver; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuf; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; @@ -81,7 +82,6 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; @@ -328,33 +328,38 @@ synchronized boolean shouldFlush(int countMin, SizeInBytes bytesMin, DataStreamR static class ReadOnlyStreamingReply { private final NettyClientReplies.RequestEntry terminalEntry; private final CompletableFuture replyFuture; - private final Consumer replyConsumer; + private final DataStreamObserver replyHandler; private Supplier> timeoutScheduler; private ScheduledFuture timeoutFuture; ReadOnlyStreamingReply(DataStreamRequest request, CompletableFuture replyFuture, - Consumer replyConsumer) { + DataStreamObserver replyHandler) { this.terminalEntry = new NettyClientReplies.RequestEntry(request); this.replyFuture = replyFuture; - this.replyConsumer = replyConsumer; + this.replyHandler = replyHandler; } - synchronized boolean receiveReply(DataStreamReply reply) { + synchronized boolean receiveReply(DataStreamReplyByteBuf reply) { NettyUtils.cancel(timeoutFuture); final boolean terminal = !reply.isSuccess() || terminalEntry.equals(new NettyClientReplies.RequestEntry(reply)); - final DataStreamReply replyToComplete = reply instanceof DataStreamReplyByteBuf ? - NettyDataStreamUtils.toDataStreamReplyByteBuffer((DataStreamReplyByteBuf) reply) : reply; + final DataStreamReply replyToComplete; try { - replyConsumer.accept(replyToComplete); + replyToComplete = terminal? reply.copy(): null; + replyHandler.onNext(reply); } catch (Throwable t) { - if (replyToComplete == reply) { - DataStreamReplyByteBuf.release(replyToComplete); - } completeExceptionally(t); return true; } if (terminal) { + try { + if (reply.isSuccess()) { + replyHandler.onCompleted(); + } + } catch (Throwable t) { + completeExceptionally(t); + return true; + } replyFuture.complete(replyToComplete); return true; } @@ -364,6 +369,11 @@ synchronized boolean receiveReply(DataStreamReply reply) { synchronized void completeExceptionally(Throwable t) { NettyUtils.cancel(timeoutFuture); + try { + replyHandler.onError(t); + } catch (Throwable e) { + LOG.warn("Failed to notify read-only stream observer", e); + } replyFuture.completeExceptionally(t); } @@ -599,11 +609,11 @@ public CompletableFuture streamAsync(DataStreamRequest request) @Override public CompletableFuture streamAsync( - DataStreamRequest request, Consumer replyConsumer) { + DataStreamRequest request, DataStreamObserver replyHandler) { final CompletableFuture f = new CompletableFuture<>(); final ClientInvocationId clientInvocationId = ClientInvocationId.valueOf(request.getClientId(), request.getStreamId()); - final ReadOnlyStreamingReply replyEntry = new ReadOnlyStreamingReply(request, f, replyConsumer); + final ReadOnlyStreamingReply replyEntry = new ReadOnlyStreamingReply(request, f, replyHandler); if (readOnlyStreamingReplies.putIfAbsent(clientInvocationId, replyEntry) != null) { f.completeExceptionally(new AlreadyClosedException(this + ": A read-only stream already exists for " + clientInvocationId)); diff --git a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java index 6726c05b79..0fce01ca72 100644 --- a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java +++ b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java @@ -21,6 +21,7 @@ import org.apache.ratis.client.DataStreamClientRpc; import org.apache.ratis.client.api.DataStreamInput; import org.apache.ratis.conf.RaftProperties; +import org.apache.ratis.datastream.DataStreamObserver; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; @@ -32,13 +33,13 @@ import org.apache.ratis.protocol.RaftClientRequest; import org.apache.ratis.protocol.RaftGroupId; import org.apache.ratis.protocol.RaftPeer; +import org.apache.ratis.thirdparty.io.netty.buffer.Unpooled; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import java.nio.ByteBuffer; import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; public class TestDataStreamClientImpl { private static RaftPeer newPeer(String id) { @@ -47,18 +48,18 @@ private static RaftPeer newPeer(String id) { private static class RecordingDataStreamClientRpc implements DataStreamClientRpc { private final AtomicReference request = new AtomicReference<>(); - private final AtomicReference> replyConsumer = new AtomicReference<>(); + private final AtomicReference> replyHandler = new AtomicReference<>(); @Override public CompletableFuture streamAsync( - DataStreamRequest dataStreamRequest, Consumer replyConsumer) { + DataStreamRequest dataStreamRequest, DataStreamObserver replyHandler) { try { final ByteBuffer buffer = ((DataStreamRequestByteBuffer) dataStreamRequest).slice(); request.set(ClientProtoUtils.toRaftClientRequest(RaftClientRequestProto.parseFrom(buffer))); } catch (Exception e) { throw new IllegalStateException(e); } - this.replyConsumer.set(replyConsumer); + this.replyHandler.set(replyHandler); return new CompletableFuture<>(); } @@ -66,8 +67,10 @@ RaftClientRequest getRequest() { return request.get(); } - void receive(DataStreamReply reply) { - replyConsumer.get().accept(reply); + void receive(DataStreamReplyByteBuf reply) { + try (DataStreamReplyByteBuf r = reply) { + replyHandler.get().onNext(r); + } } @Override @@ -105,19 +108,20 @@ public void testReceiveSkipsCancelledPendingRead() throws Exception { final CompletableFuture active = input.readAsync(); cancelled.cancel(false); - final DataStreamReply reply = DataStreamReplyByteBuffer.newBuilder() + final DataStreamReplyByteBuf reply = DataStreamReplyByteBuf.newBuilder() .setClientId(ClientId.randomId()) .setType(Type.STREAM_DATA) .setStreamId(1) .setStreamOffset(0) - .setBuffer(ByteBuffer.allocate(0)) + .setBuf(Unpooled.EMPTY_BUFFER) .setSuccess(true) .build(); dataStreamClientRpc.receive(reply); Assertions.assertTrue(active.isDone()); - Assertions.assertSame(reply, active.getNow(null)); - DataStreamReplyByteBuf.release(reply); + final DataStreamReply received = active.getNow(null); + Assertions.assertInstanceOf(DataStreamReplyByteBuffer.class, received); + Assertions.assertEquals(Type.STREAM_DATA, received.getType()); } } } diff --git a/ratis-test/src/test/java/org/apache/ratis/netty/client/TestNettyClientStreamRpcReconnectBackoff.java b/ratis-test/src/test/java/org/apache/ratis/netty/client/TestNettyClientStreamRpcReconnectBackoff.java index 304c488fc1..45cb1b0f85 100644 --- a/ratis-test/src/test/java/org/apache/ratis/netty/client/TestNettyClientStreamRpcReconnectBackoff.java +++ b/ratis-test/src/test/java/org/apache/ratis/netty/client/TestNettyClientStreamRpcReconnectBackoff.java @@ -18,17 +18,28 @@ package org.apache.ratis.netty.client; import org.apache.ratis.conf.RaftProperties; +import org.apache.ratis.datastream.DataStreamObserver; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.netty.NettyConfigKeys; +import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; +import org.apache.ratis.protocol.ClientId; +import org.apache.ratis.protocol.DataStreamReply; +import org.apache.ratis.protocol.DataStreamRequestHeader; import org.apache.ratis.protocol.RaftPeer; import org.apache.ratis.retry.ExponentialBackoffRetry; import org.apache.ratis.retry.RetryPolicy; +import org.apache.ratis.thirdparty.io.netty.buffer.Unpooled; import org.apache.ratis.util.TimeDuration; import org.junit.jupiter.api.Test; import java.net.InetSocketAddress; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; public class TestNettyClientStreamRpcReconnectBackoff { @@ -76,4 +87,67 @@ private static void assertSleepInRange(RetryPolicy policy, int attempt, TimeDura assertTrue(actual >= expected / 2, "delay too small: " + actual); assertTrue(actual <= expected + expected / 2, "delay too large: " + actual); } + + @Test + public void testReadOnlyStreamingReplyCompletesObserverOnTerminalReply() { + final ClientId clientId = ClientId.randomId(); + final long streamId = 1; + final DataStreamRequestHeader request = new DataStreamRequestHeader(clientId, Type.STREAM_HEADER, streamId, 0, 0); + final CompletableFuture replyFuture = new CompletableFuture<>(); + final AtomicReference observed = new AtomicReference<>(); + final AtomicBoolean completed = new AtomicBoolean(); + final NettyClientStreamRpc.ReadOnlyStreamingReply replies = new NettyClientStreamRpc.ReadOnlyStreamingReply( + request, replyFuture, new DataStreamObserver() { + @Override + public void onNext(DataStreamReplyByteBuf value) { + observed.set(value); + } + + @Override + public void onCompleted() { + completed.set(true); + } + }); + + final DataStreamReplyByteBuf terminalReply = DataStreamReplyByteBuf.newBuilder() + .setClientId(clientId) + .setType(Type.STREAM_HEADER) + .setStreamId(streamId) + .setStreamOffset(0) + .setBuf(Unpooled.EMPTY_BUFFER) + .setSuccess(true) + .build(); + try (DataStreamReplyByteBuf reply = terminalReply) { + assertTrue(replies.receiveReply(reply)); + } + + assertSame(terminalReply, observed.get()); + assertTrue(completed.get()); + assertTrue(replyFuture.isDone()); + } + + @Test + public void testReadOnlyStreamingReplyNotifiesObserverOnError() { + final ClientId clientId = ClientId.randomId(); + final DataStreamRequestHeader request = new DataStreamRequestHeader(clientId, Type.STREAM_HEADER, 1, 0, 0); + final CompletableFuture replyFuture = new CompletableFuture<>(); + final AtomicReference observed = new AtomicReference<>(); + final NettyClientStreamRpc.ReadOnlyStreamingReply replies = new NettyClientStreamRpc.ReadOnlyStreamingReply( + request, replyFuture, new DataStreamObserver() { + @Override + public void onNext(DataStreamReplyByteBuf value) { + } + + @Override + public void onError(Throwable t) { + observed.set(t); + } + }); + + final Throwable cause = new IllegalStateException("test"); + replies.completeExceptionally(cause); + + assertSame(cause, observed.get()); + assertTrue(replyFuture.isCompletedExceptionally()); + } } From 3cb6794bcb9fcd35dfd8bd67a885a70cc9de7a56 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 19 Jun 2026 04:10:16 +0800 Subject: [PATCH 08/13] refactor Signed-off-by: peterxcli --- .../netty/server/ReadStreamManagement.java | 33 ++++++++----------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java b/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java index 6543d2693b..d5583c9636 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/server/ReadStreamManagement.java @@ -55,7 +55,6 @@ public class ReadStreamManagement { static class ReadStream implements WritableByteChannel { private final ClientId clientId; private final long streamId; - private final RaftClientRequest request; private final ChannelHandlerContext ctx; private final CompletableFuture closed = new CompletableFuture<>(); private final DataStreamReplyByteBuffer terminalReply; @@ -64,9 +63,20 @@ static class ReadStream implements WritableByteChannel { ReadStream(RaftClientRequest request, long streamId, ChannelHandlerContext ctx) { this.clientId = request.getClientId(); this.streamId = streamId; - this.request = request; this.ctx = ctx; - this.terminalReply = newTerminalReply(); + final RaftClientReply reply = RaftClientReply.newBuilder() + .setRequest(request) + .setSuccess() + .build(); + this.terminalReply = DataStreamReplyByteBuffer.newBuilder() + .setClientId(clientId) + .setType(Type.STREAM_HEADER) + .setStreamId(streamId) + .setStreamOffset(0) + .setBuffer(toRaftClientReplyProto(reply).toByteString().asReadOnlyByteBuffer()) + .setSuccess(true) + .setBytesWritten(0) + .build(); } @Override @@ -119,23 +129,6 @@ private synchronized DataStreamReplyByteBuffer newReply(ByteBuffer buffer) { .setBytesWritten(buffer.remaining()) .build(); } - - private synchronized DataStreamReplyByteBuffer newTerminalReply() { - final RaftClientReply reply = RaftClientReply.newBuilder() - .setRequest(request) - .setSuccess() - .build(); - final ByteBuffer buffer = toRaftClientReplyProto(reply).toByteString().asReadOnlyByteBuffer(); - return DataStreamReplyByteBuffer.newBuilder() - .setClientId(clientId) - .setType(Type.STREAM_HEADER) - .setStreamId(streamId) - .setStreamOffset(0) - .setBuffer(buffer) - .setSuccess(true) - .setBytesWritten(0) - .build(); - } } private final RaftServer server; From 7b15a13173bf77eba272c447a8caac54539ce41f Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 19 Jun 2026 22:25:17 +0800 Subject: [PATCH 09/13] DataStreamClientImpl idempotent close Co-Authored-By: amaliujia Signed-off-by: peterxcli --- .../org/apache/ratis/client/impl/DataStreamClientImpl.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java index 30ec9a6b6f..3d96c567c4 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java @@ -344,6 +344,9 @@ public synchronized CompletableFuture readAsync() { @Override public synchronized void close() { + if (closed) { + return; + } closed = true; for (DataStreamReply reply; (reply = replies.poll()) != null;) { DataStreamReplyByteBuf.release(reply); From 23dacc30d8c414f5ba5b22743e4d15d41a6cd790 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 20 Jun 2026 04:43:28 +0800 Subject: [PATCH 10/13] refactor Signed-off-by: peterxcli --- .../client/impl/DataStreamClientImpl.java | 18 ++++- .../netty/client/NettyClientStreamRpc.java | 14 ++-- .../datastream/DataStreamClusterTests.java | 12 --- ...tNettyClientStreamRpcReconnectBackoff.java | 74 ------------------- 4 files changed, 22 insertions(+), 96 deletions(-) diff --git a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java index 3d96c567c4..659c614547 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java @@ -25,6 +25,7 @@ import org.apache.ratis.client.RaftClient; import org.apache.ratis.client.api.DataStreamInput; import org.apache.ratis.conf.RaftProperties; +import org.apache.ratis.datastream.DataStreamObserver; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.datastream.impl.DataStreamPacketByteBuffer; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; @@ -259,7 +260,22 @@ private DataStreamInputImpl(RaftClientRequest request) { final DataStreamRequestHeader h = new DataStreamRequestHeader(header.getClientId(), Type.STREAM_HEADER, header.getCallId(), 0, buffer.remaining(), StandardWriteOption.FLUSH, StandardWriteOption.CLOSE); this.replyFuture = dataStreamClientRpc.streamAsync(new DataStreamRequestByteBuffer(h, buffer), - reply -> receive(reply.copy())); + new DataStreamObserver() { + @Override + public void onNext(DataStreamReplyByteBuf reply) { + receive(reply.copy()); + } + + @Override + public void onError(Throwable t) { + failReads(t); + } + + @Override + public void onCompleted() { + markEndOfStream(); + } + }); replyFuture.whenComplete((reply, exception) -> { if (exception != null) { failReads(exception); diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java index e1179ceed7..ccb2a3e8f9 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java @@ -345,21 +345,17 @@ synchronized boolean receiveReply(DataStreamReplyByteBuf reply) { final DataStreamReply replyToComplete; try { replyToComplete = terminal? reply.copy(): null; - replyHandler.onNext(reply); + if (terminal && reply.isSuccess()) { + replyHandler.onCompleted(); + } else { + replyHandler.onNext(reply); + } } catch (Throwable t) { completeExceptionally(t); return true; } if (terminal) { - try { - if (reply.isSuccess()) { - replyHandler.onCompleted(); - } - } catch (Throwable t) { - completeExceptionally(t); - return true; - } replyFuture.complete(replyToComplete); return true; } diff --git a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java index a066a8f1a4..67ac49c719 100644 --- a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java +++ b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java @@ -24,7 +24,6 @@ import org.apache.ratis.protocol.RoutingTable; import org.apache.ratis.server.impl.MiniRaftCluster; import org.apache.ratis.client.RaftClient; -import org.apache.ratis.client.impl.ClientProtoUtils; import org.apache.ratis.client.impl.DataStreamClientImpl.DataStreamOutputImpl; import org.apache.ratis.datastream.DataStreamTestUtils.MultiDataStreamStateMachine; import org.apache.ratis.datastream.DataStreamTestUtils.SingleDataStream; @@ -157,17 +156,6 @@ void runTestStreamReadOnly(CLUSTER cluster, RaftPeer leader, RaftPeer primarySer Assertions.assertEquals(chunk, toByteString(data)); } - final DataStreamReply reply = in.readAsync().join(); - DataStreamTestUtils.assertSuccessReply(Type.STREAM_HEADER, 0, reply); - - try { - final RaftClientReply clientReply = ClientProtoUtils.getRaftClientReply(reply); - Assertions.assertTrue(clientReply.isSuccess()); - Assertions.assertEquals(primaryServer.getId(), clientReply.getServerId()); - } finally { - DataStreamReplyByteBuf.release(reply); - } - final ExecutionException eof = Assertions.assertThrows(ExecutionException.class, () -> in.readAsync().get(1, TimeUnit.SECONDS)); Assertions.assertInstanceOf(EOFException.class, eof.getCause()); diff --git a/ratis-test/src/test/java/org/apache/ratis/netty/client/TestNettyClientStreamRpcReconnectBackoff.java b/ratis-test/src/test/java/org/apache/ratis/netty/client/TestNettyClientStreamRpcReconnectBackoff.java index 45cb1b0f85..304c488fc1 100644 --- a/ratis-test/src/test/java/org/apache/ratis/netty/client/TestNettyClientStreamRpcReconnectBackoff.java +++ b/ratis-test/src/test/java/org/apache/ratis/netty/client/TestNettyClientStreamRpcReconnectBackoff.java @@ -18,28 +18,17 @@ package org.apache.ratis.netty.client; import org.apache.ratis.conf.RaftProperties; -import org.apache.ratis.datastream.DataStreamObserver; -import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.netty.NettyConfigKeys; -import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; -import org.apache.ratis.protocol.ClientId; -import org.apache.ratis.protocol.DataStreamReply; -import org.apache.ratis.protocol.DataStreamRequestHeader; import org.apache.ratis.protocol.RaftPeer; import org.apache.ratis.retry.ExponentialBackoffRetry; import org.apache.ratis.retry.RetryPolicy; -import org.apache.ratis.thirdparty.io.netty.buffer.Unpooled; import org.apache.ratis.util.TimeDuration; import org.junit.jupiter.api.Test; import java.net.InetSocketAddress; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicReference; import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; public class TestNettyClientStreamRpcReconnectBackoff { @@ -87,67 +76,4 @@ private static void assertSleepInRange(RetryPolicy policy, int attempt, TimeDura assertTrue(actual >= expected / 2, "delay too small: " + actual); assertTrue(actual <= expected + expected / 2, "delay too large: " + actual); } - - @Test - public void testReadOnlyStreamingReplyCompletesObserverOnTerminalReply() { - final ClientId clientId = ClientId.randomId(); - final long streamId = 1; - final DataStreamRequestHeader request = new DataStreamRequestHeader(clientId, Type.STREAM_HEADER, streamId, 0, 0); - final CompletableFuture replyFuture = new CompletableFuture<>(); - final AtomicReference observed = new AtomicReference<>(); - final AtomicBoolean completed = new AtomicBoolean(); - final NettyClientStreamRpc.ReadOnlyStreamingReply replies = new NettyClientStreamRpc.ReadOnlyStreamingReply( - request, replyFuture, new DataStreamObserver() { - @Override - public void onNext(DataStreamReplyByteBuf value) { - observed.set(value); - } - - @Override - public void onCompleted() { - completed.set(true); - } - }); - - final DataStreamReplyByteBuf terminalReply = DataStreamReplyByteBuf.newBuilder() - .setClientId(clientId) - .setType(Type.STREAM_HEADER) - .setStreamId(streamId) - .setStreamOffset(0) - .setBuf(Unpooled.EMPTY_BUFFER) - .setSuccess(true) - .build(); - try (DataStreamReplyByteBuf reply = terminalReply) { - assertTrue(replies.receiveReply(reply)); - } - - assertSame(terminalReply, observed.get()); - assertTrue(completed.get()); - assertTrue(replyFuture.isDone()); - } - - @Test - public void testReadOnlyStreamingReplyNotifiesObserverOnError() { - final ClientId clientId = ClientId.randomId(); - final DataStreamRequestHeader request = new DataStreamRequestHeader(clientId, Type.STREAM_HEADER, 1, 0, 0); - final CompletableFuture replyFuture = new CompletableFuture<>(); - final AtomicReference observed = new AtomicReference<>(); - final NettyClientStreamRpc.ReadOnlyStreamingReply replies = new NettyClientStreamRpc.ReadOnlyStreamingReply( - request, replyFuture, new DataStreamObserver() { - @Override - public void onNext(DataStreamReplyByteBuf value) { - } - - @Override - public void onError(Throwable t) { - observed.set(t); - } - }); - - final Throwable cause = new IllegalStateException("test"); - replies.completeExceptionally(cause); - - assertSame(cause, observed.get()); - assertTrue(replyFuture.isCompletedExceptionally()); - } } From 6ec28a913a2054d11ca9b5822a3c654d90910f2f Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 20 Jun 2026 05:22:07 +0800 Subject: [PATCH 11/13] refactor Signed-off-by: peterxcli --- .../netty/client/NettyClientStreamRpc.java | 11 +-- .../client/impl/TestDataStreamClientImpl.java | 56 ++++++++++++- .../datastream/DataStreamClusterTests.java | 83 +++++++++++++------ 3 files changed, 112 insertions(+), 38 deletions(-) diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java index ccb2a3e8f9..57ec0979da 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java @@ -454,20 +454,11 @@ private void process(DataStreamReplyByteBuf reply) { return; } - final DataStreamReply replyToReceive; try { - replyToReceive = NettyDataStreamUtils.toDataStreamReplyByteBuffer(reply); + replyMap.receiveReply(NettyDataStreamUtils.toDataStreamReplyByteBuffer(reply)); } catch (Throwable cause) { LOG.warn("{} : channelRead error for {}:", name, reply.getClass().getSimpleName(), cause); replyMap.completeExceptionally(cause); - return; - } - - try { - replyMap.receiveReply(replyToReceive); - } catch (Throwable cause) { - LOG.warn("{} : channelRead error for {}", name, reply, cause); - replyMap.completeExceptionally(cause); } } diff --git a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java index 0fce01ca72..2bb209a500 100644 --- a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java +++ b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java @@ -37,8 +37,10 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.io.EOFException; import java.nio.ByteBuffer; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicReference; public class TestDataStreamClientImpl { @@ -49,6 +51,7 @@ private static RaftPeer newPeer(String id) { private static class RecordingDataStreamClientRpc implements DataStreamClientRpc { private final AtomicReference request = new AtomicReference<>(); private final AtomicReference> replyHandler = new AtomicReference<>(); + private final AtomicReference> replyFuture = new AtomicReference<>(); @Override public CompletableFuture streamAsync( @@ -60,7 +63,9 @@ public CompletableFuture streamAsync( throw new IllegalStateException(e); } this.replyHandler.set(replyHandler); - return new CompletableFuture<>(); + final CompletableFuture future = new CompletableFuture<>(); + replyFuture.set(future); + return future; } RaftClientRequest getRequest() { @@ -73,6 +78,16 @@ void receive(DataStreamReplyByteBuf reply) { } } + void complete() { + replyHandler.get().onCompleted(); + replyFuture.get().complete(null); + } + + void completeExceptionally(Throwable cause) { + replyHandler.get().onError(cause); + replyFuture.get().completeExceptionally(cause); + } + @Override public void close() { } @@ -124,4 +139,43 @@ public void testReceiveSkipsCancelledPendingRead() throws Exception { Assertions.assertEquals(Type.STREAM_DATA, received.getType()); } } + + @Test + public void testReadOnlyInputCompletesPendingReadOnCompleted() throws Exception { + final RaftPeer follower = newPeer("follower"); + final RecordingDataStreamClientRpc dataStreamClientRpc = new RecordingDataStreamClientRpc(); + + try (DataStreamClient dataStreamClient = newDataStreamClient(follower, dataStreamClientRpc); + DataStreamInput input = dataStreamClient.streamReadOnly(ByteBuffer.wrap(new byte[] {1}))) { + final CompletableFuture pending = input.readAsync(); + + dataStreamClientRpc.complete(); + + assertFutureCause(pending, EOFException.class); + assertFutureCause(input.readAsync(), EOFException.class); + } + } + + @Test + public void testReadOnlyInputNotifiesPendingReadOnError() throws Exception { + final RaftPeer follower = newPeer("follower"); + final RecordingDataStreamClientRpc dataStreamClientRpc = new RecordingDataStreamClientRpc(); + + try (DataStreamClient dataStreamClient = newDataStreamClient(follower, dataStreamClientRpc); + DataStreamInput input = dataStreamClient.streamReadOnly(ByteBuffer.wrap(new byte[] {1}))) { + final CompletableFuture pending = input.readAsync(); + final Throwable cause = new IllegalStateException("test"); + + dataStreamClientRpc.completeExceptionally(cause); + + Assertions.assertSame(cause, assertFutureCause(pending, IllegalStateException.class)); + Assertions.assertSame(cause, assertFutureCause(input.readAsync(), IllegalStateException.class)); + } + } + + private static Throwable assertFutureCause( + CompletableFuture future, Class expectedCauseClass) { + final ExecutionException exception = Assertions.assertThrows(ExecutionException.class, future::get); + return Assertions.assertInstanceOf(expectedCauseClass, exception.getCause()); + } } diff --git a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java index 67ac49c719..02ad13e3ec 100644 --- a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java +++ b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java @@ -18,23 +18,28 @@ package org.apache.ratis.datastream; import org.apache.ratis.BaseTest; -import org.apache.ratis.client.api.DataStreamInput; +import org.apache.ratis.client.DataStreamClient; import org.apache.ratis.io.StandardWriteOption; import org.apache.ratis.protocol.RaftPeer; import org.apache.ratis.protocol.RoutingTable; import org.apache.ratis.server.impl.MiniRaftCluster; import org.apache.ratis.client.RaftClient; +import org.apache.ratis.client.impl.ClientProtoUtils; import org.apache.ratis.client.impl.DataStreamClientImpl.DataStreamOutputImpl; import org.apache.ratis.datastream.DataStreamTestUtils.MultiDataStreamStateMachine; import org.apache.ratis.datastream.DataStreamTestUtils.SingleDataStream; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; -import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; +import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; import org.apache.ratis.proto.RaftProtos.ReplicationLevel; import org.apache.ratis.protocol.DataStreamReply; +import org.apache.ratis.protocol.DataStreamRequest; +import org.apache.ratis.protocol.DataStreamRequestHeader; +import org.apache.ratis.protocol.Message; import org.apache.ratis.protocol.RaftClientReply; import org.apache.ratis.protocol.RaftClientRequest; import org.apache.ratis.retry.RetryPolicies; +import org.apache.ratis.rpc.CallId; import org.apache.ratis.server.RaftServer; import org.apache.ratis.server.RaftServerConfigKeys; import org.apache.ratis.thirdparty.com.google.protobuf.ByteString; @@ -45,14 +50,13 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; -import java.io.EOFException; import java.io.File; import java.nio.ByteBuffer; import java.nio.channels.FileChannel; import java.nio.file.StandardOpenOption; import java.util.Collection; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; @@ -147,33 +151,58 @@ void runTestStreamReadOnlyWithNonLeaderPrimary(CLUSTER cluster) throws Exception void runTestStreamReadOnly(CLUSTER cluster, RaftPeer leader, RaftPeer primaryServer) throws Exception { final ByteString query = ByteString.copyFromUtf8("stream-read-only"); try (RaftClient client = cluster.createClient(leader.getId(), cluster.getGroup(), RetryPolicies.noRetry(), - primaryServer); - DataStreamInput in = client.getDataStreamApi().streamReadOnly(query.asReadOnlyByteBuffer())) { - for (int i = 0; i < MultiDataStreamStateMachine.READ_ONLY_STREAM_CHUNKS; i++) { - final ByteString chunk = MultiDataStreamStateMachine.getReadOnlyStreamChunk(query, i); - final DataStreamReply data = in.readAsync().join(); - DataStreamTestUtils.assertSuccessReply(Type.STREAM_DATA, chunk.size(), data); - Assertions.assertEquals(chunk, toByteString(data)); + primaryServer)) { + final AtomicInteger replyCount = new AtomicInteger(); + final CompletableFuture completed = new CompletableFuture<>(); + final DataStreamClient dataStreamClient = (DataStreamClient) client.getDataStreamApi(); + final CompletableFuture terminalReplyFuture = dataStreamClient.getClientRpc().streamAsync( + newReadOnlyStreamRequest(client, primaryServer, query), new DataStreamObserver() { + @Override + public void onNext(DataStreamReplyByteBuf reply) { + final int i = replyCount.getAndIncrement(); + Assertions.assertTrue(i < MultiDataStreamStateMachine.READ_ONLY_STREAM_CHUNKS); + + final ByteString chunk = MultiDataStreamStateMachine.getReadOnlyStreamChunk(query, i); + DataStreamTestUtils.assertSuccessReply(Type.STREAM_DATA, chunk.size(), reply); + Assertions.assertEquals(chunk, ByteString.copyFrom(reply.slice().nioBuffer())); + } + + @Override + public void onError(Throwable t) { + completed.completeExceptionally(t); + } + + @Override + public void onCompleted() { + completed.complete(null); + } + }); + + completed.get(1, TimeUnit.SECONDS); + Assertions.assertEquals(MultiDataStreamStateMachine.READ_ONLY_STREAM_CHUNKS, replyCount.get()); + final DataStreamReply terminalReply = terminalReplyFuture.get(1, TimeUnit.SECONDS); + try { + DataStreamTestUtils.assertSuccessReply(Type.STREAM_HEADER, 0, terminalReply); + } finally { + DataStreamReplyByteBuf.release(terminalReply); } - - final ExecutionException eof = Assertions.assertThrows(ExecutionException.class, - () -> in.readAsync().get(1, TimeUnit.SECONDS)); - Assertions.assertInstanceOf(EOFException.class, eof.getCause()); } } - private static ByteString toByteString(DataStreamReply reply) { - try { - if (reply instanceof DataStreamReplyByteBuffer) { - final ByteBuffer buffer = ((DataStreamReplyByteBuffer) reply).slice(); - return ByteString.copyFrom(buffer); - } else if (reply instanceof DataStreamReplyByteBuf) { - return ByteString.copyFrom(((DataStreamReplyByteBuf) reply).slice().nioBuffer()); - } - throw new AssertionError("Unexpected reply " + reply); - } finally { - DataStreamReplyByteBuf.release(reply); - } + private static DataStreamRequest newReadOnlyStreamRequest(RaftClient client, RaftPeer primaryServer, + ByteString query) { + final RaftClientRequest request = RaftClientRequest.newBuilder() + .setClientId(client.getId()) + .setServerId(primaryServer.getId()) + .setGroupId(client.getGroupId()) + .setCallId(CallId.getAndIncrement()) + .setMessage(Message.valueOf(query)) + .setType(RaftClientRequest.readRequestType()) + .build(); + final ByteBuffer buffer = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(request); + final DataStreamRequestHeader header = new DataStreamRequestHeader(request.getClientId(), Type.STREAM_HEADER, + request.getCallId(), 0, buffer.remaining(), StandardWriteOption.FLUSH, StandardWriteOption.CLOSE); + return new DataStreamRequestByteBuffer(header, buffer); } void runTestInvalidPrimaryInRoutingTable(CLUSTER cluster) throws Exception { From d693fbd4cf5bab57348d46e10158eff1c50b9858 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 20 Jun 2026 05:36:43 +0800 Subject: [PATCH 12/13] refactor test Signed-off-by: peterxcli --- .../client/impl/TestDataStreamClientImpl.java | 13 +-- .../server/TestDataStreamManagement.java | 107 +++++++++--------- 2 files changed, 53 insertions(+), 67 deletions(-) diff --git a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java index 2bb209a500..9cd9217283 100644 --- a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java +++ b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java @@ -100,18 +100,6 @@ private static DataStreamClient newDataStreamClient( ClientId.randomId(), RaftGroupId.randomId(), dataStreamServer, dataStreamClientRpc, properties); } - @Test - public void testStreamReadOnlyUsesPrimaryDataStreamServer() throws Exception { - final RaftPeer follower = newPeer("follower"); - final RecordingDataStreamClientRpc dataStreamClientRpc = new RecordingDataStreamClientRpc(); - - try (DataStreamClient dataStreamClient = newDataStreamClient(follower, dataStreamClientRpc); - DataStreamInput input = dataStreamClient.streamReadOnly(ByteBuffer.wrap(new byte[] {1}))) { - Assertions.assertNotNull(input); - Assertions.assertEquals(follower.getId(), dataStreamClientRpc.getRequest().getServerId()); - } - } - @Test public void testReceiveSkipsCancelledPendingRead() throws Exception { final RaftPeer follower = newPeer("follower"); @@ -122,6 +110,7 @@ public void testReceiveSkipsCancelledPendingRead() throws Exception { final CompletableFuture cancelled = input.readAsync(); final CompletableFuture active = input.readAsync(); cancelled.cancel(false); + Assertions.assertEquals(follower.getId(), dataStreamClientRpc.getRequest().getServerId()); final DataStreamReplyByteBuf reply = DataStreamReplyByteBuf.newBuilder() .setClientId(ClientId.randomId()) diff --git a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java index 1605d18fcf..419f21e405 100644 --- a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java +++ b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java @@ -87,37 +87,14 @@ 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)); - final ReadStreamManagement management = new ReadStreamManagement(server); + final ReadStreamManagement management = newReadStreamManagement(serverId, groupId, dataApi); final EmbeddedChannel embeddedChannel = new EmbeddedChannel(new ChannelInboundHandlerAdapter()); - final RaftClientRequest raftClientRequest = RaftClientRequest.newBuilder() - .setClientId(clientId) - .setServerId(serverId) - .setGroupId(groupId) - .setCallId(1L) - .setMessage(Message.valueOf(query)) - .setType(RaftClientRequest.readRequestType()) - .build(); - final ByteBuffer header = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(raftClientRequest); - final ByteBuf headerBuf = Unpooled.wrappedBuffer(header); - final DataStreamRequestByteBuf request = new DataStreamRequestByteBuf( - clientId, - Type.STREAM_HEADER, - raftClientRequest.getCallId(), - 0L, - Collections.singletonList(StandardWriteOption.FLUSH), - headerBuf); + final ReadOnlyRequest readOnlyRequest = newReadOnlyRequest(clientId, serverId, groupId, 1L, query); try { - assertTrue(management.process(request, embeddedChannel.pipeline().firstContext())); - assertEquals(0, headerBuf.refCnt()); + assertTrue(management.process(readOnlyRequest.request, embeddedChannel.pipeline().firstContext())); + assertEquals(0, readOnlyRequest.headerBuf.refCnt()); final WritableByteChannel stream = streamRef.get(); assertNotNull(stream); @@ -162,38 +139,15 @@ public void query(Message request, WritableByteChannel stream) { queryDone.countDown(); } }; - final StateMachine stateMachine = new BaseStateMachine() { - @Override - public DataApi data() { - return dataApi; - } - }; - final RaftServer server = newRaftServer(serverId, new RaftProperties(), groupId, newDivision(stateMachine)); - final ReadStreamManagement management = new ReadStreamManagement(server); - - final RaftClientRequest raftClientRequest = RaftClientRequest.newBuilder() - .setClientId(clientId) - .setServerId(serverId) - .setGroupId(groupId) - .setCallId(1L) - .setMessage(Message.valueOf(ByteString.copyFromUtf8("query"))) - .setType(RaftClientRequest.readRequestType()) - .build(); - final ByteBuffer header = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(raftClientRequest); - final ByteBuf headerBuf = Unpooled.wrappedBuffer(header); - final DataStreamRequestByteBuf request = new DataStreamRequestByteBuf( - clientId, - Type.STREAM_HEADER, - raftClientRequest.getCallId(), - 0L, - Collections.singletonList(StandardWriteOption.FLUSH), - headerBuf); + final ReadStreamManagement management = newReadStreamManagement(serverId, groupId, dataApi); + final ReadOnlyRequest readOnlyRequest = newReadOnlyRequest( + clientId, serverId, groupId, 1L, ByteString.copyFromUtf8("query")); try { - eventLoop.submit(() -> assertTrue(management.process(request, ctx))).sync(); + eventLoop.submit(() -> assertTrue(management.process(readOnlyRequest.request, ctx))).sync(); assertTrue(queryDone.await(10, TimeUnit.SECONDS)); - assertEquals(0, headerBuf.refCnt()); + assertEquals(0, readOnlyRequest.headerBuf.refCnt()); assertFalse(queryInEventLoop.get(), "read-only state machine query should not run on Netty event loop"); } finally { embeddedChannel.finishAndReleaseAll(); @@ -242,6 +196,49 @@ void readCleansChannelMapOnEarlyException() throws Exception { } } + private static class ReadOnlyRequest { + private final DataStreamRequestByteBuf request; + private final ByteBuf headerBuf; + + ReadOnlyRequest(DataStreamRequestByteBuf request, ByteBuf headerBuf) { + this.request = request; + this.headerBuf = headerBuf; + } + } + + private static ReadStreamManagement newReadStreamManagement( + RaftPeerId serverId, RaftGroupId groupId, DataApi dataApi) { + final StateMachine stateMachine = new BaseStateMachine() { + @Override + public DataApi data() { + return dataApi; + } + }; + final RaftServer server = newRaftServer(serverId, new RaftProperties(), groupId, newDivision(stateMachine)); + return new ReadStreamManagement(server); + } + + private static ReadOnlyRequest newReadOnlyRequest(ClientId clientId, RaftPeerId serverId, RaftGroupId groupId, + long callId, ByteString query) { + final RaftClientRequest raftClientRequest = RaftClientRequest.newBuilder() + .setClientId(clientId) + .setServerId(serverId) + .setGroupId(groupId) + .setCallId(callId) + .setMessage(Message.valueOf(query)) + .setType(RaftClientRequest.readRequestType()) + .build(); + final ByteBuffer header = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(raftClientRequest); + final ByteBuf headerBuf = Unpooled.wrappedBuffer(header); + return new ReadOnlyRequest(new DataStreamRequestByteBuf( + clientId, + Type.STREAM_HEADER, + raftClientRequest.getCallId(), + 0L, + Collections.singletonList(StandardWriteOption.FLUSH), + headerBuf), headerBuf); + } + private static void assertSuccessReply(Type expectedType, long expectedBytesWritten, DataStreamReply reply) { assertEquals(expectedType, reply.getType()); assertTrue(reply.isSuccess()); From 490662411bb526ad937e4d442fb22272bb6ecee4 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 21 Jun 2026 01:36:21 +0800 Subject: [PATCH 13/13] Add RC for readAsync reply and extract inner class Co-Authored-By: szetszwo Signed-off-by: peterxcli --- .../dev-support/findbugsExcludeFile.xml | 2 +- .../ratis/client/DataStreamClientRpc.java | 14 +- .../ratis/client/api/DataStreamApi.java | 8 +- .../ratis/client/api/DataStreamInput.java | 9 +- .../client/impl/DataStreamClientImpl.java | 155 +----------------- .../client/impl/DataStreamInputImpl.java | 141 ++++++++++++++++ .../ratis/datastream/DataStreamObserver.java | 19 ++- .../impl/DataStreamPacketByteBuf.java | 5 +- .../ratis/netty/client/ClientReadStream.java | 80 +++++++++ .../netty/client/NettyClientStreamRpc.java | 123 ++++---------- .../client/impl/TestDataStreamClientImpl.java | 34 ++-- .../datastream/DataStreamClusterTests.java | 81 +++------ .../server/TestDataStreamManagement.java | 33 +++- 13 files changed, 376 insertions(+), 328 deletions(-) create mode 100644 ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamInputImpl.java create mode 100644 ratis-netty/src/main/java/org/apache/ratis/netty/client/ClientReadStream.java diff --git a/ratis-client/dev-support/findbugsExcludeFile.xml b/ratis-client/dev-support/findbugsExcludeFile.xml index d7e97b5893..ddeb10b40c 100644 --- a/ratis-client/dev-support/findbugsExcludeFile.xml +++ b/ratis-client/dev-support/findbugsExcludeFile.xml @@ -28,7 +28,7 @@ - + diff --git a/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java b/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java index 11210de9ee..6f94709b0d 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java @@ -19,10 +19,10 @@ package org.apache.ratis.client; import org.apache.ratis.datastream.DataStreamObserver; -import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; import org.apache.ratis.protocol.DataStreamReply; import org.apache.ratis.protocol.DataStreamRequest; import org.apache.ratis.util.JavaUtils; +import org.apache.ratis.util.ReferenceCountedObject; import java.io.Closeable; import java.util.concurrent.CompletableFuture; @@ -39,9 +39,15 @@ default CompletableFuture streamAsync(DataStreamRequest request + JavaUtils.getCurrentStackTraceElement().getMethodName()); } - /** Async call to send a request and receive multiple replies for the request. */ - default CompletableFuture streamAsync( - DataStreamRequest request, DataStreamObserver replyHandler) { + /** + * Async call to send a request to receive a stream of intermediate replies and a final reply. + * + * @param request the request + * @param replyHandler to handle intermediate replies + * @return a future the final reply + */ + default CompletableFuture streamAsync(DataStreamRequest request, + DataStreamObserver> replyHandler) { throw new UnsupportedOperationException(getClass() + " does not support " + JavaUtils.getCurrentStackTraceElement().getMethodName()); } diff --git a/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamApi.java b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamApi.java index 800337d8f5..7152d35ba2 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamApi.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamApi.java @@ -51,13 +51,11 @@ default DataStreamOutput stream() { /** Create a stream by providing a customized header message and route table. */ DataStreamOutput stream(ByteBuffer headerMessage, RoutingTable routingTable); - /** - * Create a stream to read data for readonly requests. - */ + /** Create a stream to read data. */ default DataStreamInput streamReadOnly() { return streamReadOnly(null); } - /** Create a stream by providing a customized header message for readonly requests. */ - DataStreamInput streamReadOnly(ByteBuffer message); + /** Create a stream by providing a customized header message for reading data. */ + DataStreamInput streamReadOnly(ByteBuffer headerMessage); } diff --git a/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java index 0356a6a97f..e50033405f 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java @@ -18,6 +18,7 @@ package org.apache.ratis.client.api; import org.apache.ratis.protocol.DataStreamReply; +import org.apache.ratis.util.ReferenceCountedObject; import java.io.Closeable; import java.util.concurrent.CompletableFuture; @@ -28,10 +29,10 @@ public interface DataStreamInput extends Closeable { /** * Read the next chunk in the stream asynchronously. - * The caller owns the returned {@link DataStreamReply} and should call - * {@link DataStreamReply#release()} after consuming it. + * The caller owns the returned {@link DataStreamReply} which is a {@link ReferenceCountedObject}. + * It must call {@link ReferenceCountedObject#release()} after consuming it. * - * @return a future of the reply. + * @return a future of the reference-counted reply. */ - CompletableFuture readAsync(); + CompletableFuture> readAsync(); } diff --git a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java index 659c614547..bc8334bf98 100644 --- a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java +++ b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java @@ -25,11 +25,8 @@ import org.apache.ratis.client.RaftClient; import org.apache.ratis.client.api.DataStreamInput; import org.apache.ratis.conf.RaftProperties; -import org.apache.ratis.datastream.DataStreamObserver; -import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; 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.FilePositionCount; import org.apache.ratis.io.StandardWriteOption; import org.apache.ratis.io.WriteOption; @@ -46,7 +43,6 @@ import org.apache.ratis.protocol.RoutingTable; import org.apache.ratis.protocol.exceptions.AlreadyClosedException; import org.apache.ratis.rpc.CallId; -import org.apache.ratis.thirdparty.com.google.protobuf.ByteString; import org.apache.ratis.thirdparty.io.netty.buffer.ByteBuf; import org.apache.ratis.util.IOUtils; import org.apache.ratis.util.JavaUtils; @@ -56,16 +52,12 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.EOFException; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.WritableByteChannel; -import java.util.ArrayDeque; import java.util.Arrays; import java.util.Collections; import java.util.Objects; -import java.util.Optional; -import java.util.Queue; import java.util.concurrent.CompletableFuture; /** @@ -245,132 +237,6 @@ private CompletableFuture sendForward(DataStreamReply writeRepl } } - public final class DataStreamInputImpl implements DataStreamInput { - private final RaftClientRequest header; - private final CompletableFuture replyFuture; - private final Queue replies = new ArrayDeque<>(); - private final Queue> pendingReads = new ArrayDeque<>(); - private Throwable readException; - private boolean endOfStream; - private boolean closed; - - private DataStreamInputImpl(RaftClientRequest request) { - this.header = request; - final ByteBuffer buffer = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(header); - final DataStreamRequestHeader h = new DataStreamRequestHeader(header.getClientId(), Type.STREAM_HEADER, - header.getCallId(), 0, buffer.remaining(), StandardWriteOption.FLUSH, StandardWriteOption.CLOSE); - this.replyFuture = dataStreamClientRpc.streamAsync(new DataStreamRequestByteBuffer(h, buffer), - new DataStreamObserver() { - @Override - public void onNext(DataStreamReplyByteBuf reply) { - receive(reply.copy()); - } - - @Override - public void onError(Throwable t) { - failReads(t); - } - - @Override - public void onCompleted() { - markEndOfStream(); - } - }); - replyFuture.whenComplete((reply, exception) -> { - if (exception != null) { - failReads(exception); - } else { - markEndOfStream(); - } - }); - } - - private void receive(DataStreamReply reply) { - for (;;) { - final CompletableFuture pending; - synchronized (this) { - if (closed) { - DataStreamReplyByteBuf.release(reply); - return; - } - pending = pendingReads.poll(); - if (pending == null) { - replies.add(reply); - return; - } - } - if (pending.complete(reply)) { - return; - } - } - } - - private void failReads(Throwable t) { - for (;;) { - final CompletableFuture pending; - synchronized (this) { - readException = t; - pending = pendingReads.poll(); - if (pending == null) { - return; - } - } - pending.completeExceptionally(t); - } - } - - private EOFException newEndOfStreamException() { - return new EOFException(clientId + ": end of stream, request=" + header); - } - - private void markEndOfStream() { - for (;;) { - final CompletableFuture pending; - synchronized (this) { - endOfStream = true; - pending = pendingReads.poll(); - if (pending == null) { - return; - } - } - pending.completeExceptionally(newEndOfStreamException()); - } - } - - @Override - public synchronized CompletableFuture readAsync() { - if (closed) { - return JavaUtils.completeExceptionally(new AlreadyClosedException( - clientId + ": stream already closed, request=" + header)); - } - final DataStreamReply reply = replies.poll(); - if (reply != null) { - return CompletableFuture.completedFuture(reply); - } - if (readException != null) { - return JavaUtils.completeExceptionally(readException); - } - if (endOfStream) { - return JavaUtils.completeExceptionally(newEndOfStreamException()); - } - final CompletableFuture f = new CompletableFuture<>(); - pendingReads.add(f); - return f; - } - - @Override - public synchronized void close() { - if (closed) { - return; - } - closed = true; - for (DataStreamReply reply; (reply = replies.poll()) != null;) { - DataStreamReplyByteBuf.release(reply); - } - failReads(new AlreadyClosedException(clientId + ": stream already closed, request=" + header)); - } - } - @Override public DataStreamClientRpc getClientRpc() { return dataStreamClientRpc; @@ -400,6 +266,12 @@ public DataStreamOutputRpc stream(ByteBuffer headerMessage, RoutingTable routing .build()); } + @Override + public DataStreamInput streamReadOnly(ByteBuffer headerMessage) { + return new DataStreamInputImpl(dataStreamClientRpc, + newBuilder(headerMessage).setType(RaftClientRequest.readRequestType()).build()); + } + private RaftClientRequest.Builder newBuilder(ByteBuffer headerMessage) { final Message message = headerMessage == null ? null : Message.valueOf(headerMessage); return RaftClientRequest.newBuilder() @@ -410,21 +282,6 @@ private RaftClientRequest.Builder newBuilder(ByteBuffer headerMessage) { .setMessage(message); } - @Override - public DataStreamInput streamReadOnly(ByteBuffer headerMessage) { - final Message message = - Optional.ofNullable(headerMessage).map(ByteString::copyFrom).map(Message::valueOf).orElse(null); - final RaftClientRequest request = RaftClientRequest.newBuilder() - .setClientId(clientId) - .setServerId(dataStreamServer.getId()) - .setGroupId(groupId) - .setCallId(CallId.getAndIncrement()) - .setMessage(message) - .setType(RaftClientRequest.readRequestType()) - .build(); - return new DataStreamInputImpl(request); - } - @Override public void close() throws IOException { dataStreamClientRpc.close(); diff --git a/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamInputImpl.java b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamInputImpl.java new file mode 100644 index 0000000000..c3e4f3420d --- /dev/null +++ b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamInputImpl.java @@ -0,0 +1,141 @@ +/* + * 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.client.impl; + +import org.apache.ratis.client.DataStreamClientRpc; +import org.apache.ratis.client.api.DataStreamInput; +import org.apache.ratis.datastream.DataStreamObserver; +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.protocol.ClientId; +import org.apache.ratis.protocol.DataStreamReply; +import org.apache.ratis.protocol.DataStreamRequestHeader; +import org.apache.ratis.protocol.RaftClientRequest; +import org.apache.ratis.protocol.exceptions.AlreadyClosedException; +import org.apache.ratis.util.JavaUtils; +import org.apache.ratis.util.ReferenceCountedObject; + +import java.io.EOFException; +import java.nio.ByteBuffer; +import java.util.LinkedList; +import java.util.Objects; +import java.util.Queue; +import java.util.concurrent.CompletableFuture; + +final class DataStreamInputImpl implements DataStreamInput, + DataStreamObserver> { + private final RaftClientRequest header; + private final ClientId clientId; + private final Queue> replies = new LinkedList<>(); + private final Queue>> pendingReads = new LinkedList<>(); + + /* + * null : the stream is open. + * AlreadyClosedException: the stream is closed. + * Other exception : the stream is failed. + */ + private Throwable readException; + + DataStreamInputImpl(DataStreamClientRpc dataStreamClientRpc, RaftClientRequest request) { + this.header = request; + this.clientId = request.getClientId(); + final ByteBuffer buffer = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(header); + final DataStreamRequestHeader h = new DataStreamRequestHeader(clientId, Type.STREAM_HEADER, + header.getCallId(), 0, buffer.remaining(), StandardWriteOption.FLUSH, StandardWriteOption.CLOSE); + dataStreamClientRpc.streamAsync(new DataStreamRequestByteBuffer(h, buffer), this) + .whenComplete(asWhenCompleteBiConsumer()); + } + + @Override + public synchronized void onNext(ReferenceCountedObject reply) { + if (readException != null) { + return; + } + + reply.retain(); + for (CompletableFuture> pending; + (pending = pendingReads.poll()) != null; ) { + if (pending.complete(reply)) { + return; + } + } + replies.add(reply); + } + + @Override + public synchronized void onError(Throwable throwable) { + // An error case, release the replies + releaseReplies(); + if (readException == null) { + readException = throwable; + failPendingReads(); + } + } + + @Override + public synchronized void onCompleted() { + // Not an error case, do not release the replies + if (readException == null) { + // No more onNext(), the pending reads cannot be completed. + readException = new EOFException(clientId + ": end of stream, request=" + header); + failPendingReads(); + } + } + + private void releaseReplies() { + for (ReferenceCountedObject reply; (reply = replies.poll()) != null; ) { + reply.release(); + } + } + + private void failPendingReads() { + Objects.requireNonNull(readException, "readException == null"); + for (CompletableFuture> p; (p = pendingReads.poll()) != null; ) { + p.completeExceptionally(readException); + } + } + + @Override + public synchronized CompletableFuture> readAsync() { + final ReferenceCountedObject reply = replies.poll(); + if (reply != null) { + return CompletableFuture.completedFuture(reply); + } + if (readException != null) { + return JavaUtils.completeExceptionally(readException); + } + final CompletableFuture> f = + new CompletableFuture<>(); + pendingReads.add(f); + return f; + } + + @Override + public synchronized void close() { + // When close() is called the first time, we set + // (1) release all replies + // (2) set readException to AlreadyClosedException + // (3) complete all pendingReads, and + releaseReplies(); + if (readException == null) { + readException = new AlreadyClosedException(clientId + ": stream already closed, request=" + header); + failPendingReads(); + } + } +} diff --git a/ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java b/ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java index 82529177cb..92f62b7518 100644 --- a/ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java +++ b/ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java @@ -17,14 +17,23 @@ */ package org.apache.ratis.datastream; -/** An interface similar to gRPC {@link org.apache.ratis.thirdparty.io.grpc.stub.StreamObserver}. */ -@FunctionalInterface +import java.util.function.BiConsumer; + +/** An interface similar to gRPC {@link org.apache.ratis.thirdparty.io.grpc.stub.StreamObserver}. */ public interface DataStreamObserver { void onNext(V value); - default void onError(Throwable t) { - } + void onCompleted(); + + void onError(Throwable throwable); - default void onCompleted() { + default BiConsumer asWhenCompleteBiConsumer() { + return (v, throwable) -> { + if (throwable != null) { + onError(throwable); + } else { + onCompleted(); + } + }; } } diff --git a/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamPacketByteBuf.java b/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamPacketByteBuf.java index 76400ff803..e6ceeb93df 100644 --- a/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamPacketByteBuf.java +++ b/ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamPacketByteBuf.java @@ -30,7 +30,7 @@ *

* This class is immutable. */ -class DataStreamPacketByteBuf extends DataStreamPacketImpl { +public class DataStreamPacketByteBuf extends DataStreamPacketImpl { private final AtomicReference buf; DataStreamPacketByteBuf(ClientId clientId, Type type, long streamId, long streamOffset, ByteBuf buf) { @@ -48,7 +48,8 @@ final ByteBuf getBuf() { @Override public final long getDataLength() { - return getBuf().readableBytes(); + final ByteBuf got = buf.get(); + return got != null ? got.readableBytes() : -1; } public final ByteBuf slice() { diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/client/ClientReadStream.java b/ratis-netty/src/main/java/org/apache/ratis/netty/client/ClientReadStream.java new file mode 100644 index 0000000000..c9096ae7be --- /dev/null +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/client/ClientReadStream.java @@ -0,0 +1,80 @@ +/* + * 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.netty.client; + +import org.apache.ratis.datastream.DataStreamObserver; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; +import org.apache.ratis.protocol.DataStreamReply; +import org.apache.ratis.protocol.DataStreamRequest; +import org.apache.ratis.thirdparty.io.netty.util.concurrent.ScheduledFuture; +import org.apache.ratis.util.NettyUtils; +import org.apache.ratis.util.Preconditions; +import org.apache.ratis.util.ReferenceCountedObject; + +import java.util.concurrent.CompletableFuture; +import java.util.function.Supplier; + +class ClientReadStream { + private final NettyClientReplies.RequestEntry terminalEntry; + private final CompletableFuture replyFuture; + private final DataStreamObserver> replyHandler; + private Supplier> timeoutScheduler; + private ScheduledFuture timeoutFuture; + + ClientReadStream(DataStreamRequest request, CompletableFuture replyFuture, + DataStreamObserver> replyHandler) { + this.terminalEntry = new NettyClientReplies.RequestEntry(request); + this.replyFuture = replyFuture; + this.replyHandler = replyHandler; + } + + synchronized boolean receiveReply(ReferenceCountedObject ref) { + NettyUtils.cancel(timeoutFuture); + try { + replyHandler.onNext(ref); + } catch (Throwable t) { + completeExceptionally(t); + return true; + } + + final DataStreamReplyByteBuf reply = Preconditions.assertInstanceOf(ref.get(), DataStreamReplyByteBuf.class); + final boolean terminal = !reply.isSuccess() || terminalEntry.equals(new NettyClientReplies.RequestEntry(reply)); + if (terminal) { + replyFuture.complete(reply.copy()); + return true; + } + scheduleTimeout(); + return false; + } + + synchronized void completeExceptionally(Throwable t) { + NettyUtils.cancel(timeoutFuture); + replyFuture.completeExceptionally(t); + } + + synchronized void scheduleTimeout(Supplier> scheduleMethod) { + timeoutScheduler = scheduleMethod; + scheduleTimeout(); + } + + private void scheduleTimeout() { + if (!replyFuture.isDone() && timeoutScheduler != null) { + timeoutFuture = timeoutScheduler.get(); + } + } +} diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java index 57ec0979da..0853150009 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientStreamRpc.java @@ -19,9 +19,10 @@ package org.apache.ratis.netty.client; import org.apache.ratis.client.DataStreamClientRpc; +import org.apache.ratis.datastream.DataStreamObserver; import org.apache.ratis.client.RaftClientConfigKeys; import org.apache.ratis.conf.RaftProperties; -import org.apache.ratis.datastream.DataStreamObserver; +import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuf; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; @@ -57,7 +58,6 @@ import org.apache.ratis.thirdparty.io.netty.handler.codec.ByteToMessageDecoder; import org.apache.ratis.thirdparty.io.netty.handler.codec.MessageToMessageEncoder; import org.apache.ratis.thirdparty.io.netty.handler.ssl.SslContext; -import org.apache.ratis.thirdparty.io.netty.util.concurrent.ScheduledFuture; import org.apache.ratis.util.JavaUtils; import org.apache.ratis.util.MemoizedFunction; import org.apache.ratis.util.MemoizedSupplier; @@ -66,6 +66,7 @@ import org.apache.ratis.util.ReferenceCountedObject; import org.apache.ratis.util.SizeInBytes; import org.apache.ratis.util.TimeDuration; +import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -325,72 +326,11 @@ synchronized boolean shouldFlush(int countMin, SizeInBytes bytesMin, DataStreamR } } - static class ReadOnlyStreamingReply { - private final NettyClientReplies.RequestEntry terminalEntry; - private final CompletableFuture replyFuture; - private final DataStreamObserver replyHandler; - private Supplier> timeoutScheduler; - private ScheduledFuture timeoutFuture; - - ReadOnlyStreamingReply(DataStreamRequest request, CompletableFuture replyFuture, - DataStreamObserver replyHandler) { - this.terminalEntry = new NettyClientReplies.RequestEntry(request); - this.replyFuture = replyFuture; - this.replyHandler = replyHandler; - } - - synchronized boolean receiveReply(DataStreamReplyByteBuf reply) { - NettyUtils.cancel(timeoutFuture); - final boolean terminal = !reply.isSuccess() || terminalEntry.equals(new NettyClientReplies.RequestEntry(reply)); - final DataStreamReply replyToComplete; - try { - replyToComplete = terminal? reply.copy(): null; - if (terminal && reply.isSuccess()) { - replyHandler.onCompleted(); - } else { - replyHandler.onNext(reply); - } - } catch (Throwable t) { - completeExceptionally(t); - return true; - } - - if (terminal) { - replyFuture.complete(replyToComplete); - return true; - } - scheduleTimeout(); - return false; - } - - synchronized void completeExceptionally(Throwable t) { - NettyUtils.cancel(timeoutFuture); - try { - replyHandler.onError(t); - } catch (Throwable e) { - LOG.warn("Failed to notify read-only stream observer", e); - } - replyFuture.completeExceptionally(t); - } - - synchronized void scheduleTimeout(Supplier> scheduleMethod) { - timeoutScheduler = scheduleMethod; - scheduleTimeout(); - } - - private void scheduleTimeout() { - if (!replyFuture.isDone() && timeoutScheduler != null) { - timeoutFuture = timeoutScheduler.get(); - } - } - } - private final String name; private final Connection connection; private final NettyClientReplies replies = new NettyClientReplies(); - private final ConcurrentMap readOnlyStreamingReplies - = new ConcurrentHashMap<>(); + private final ConcurrentMap readStreams = new ConcurrentHashMap<>(); private final TimeDuration requestTimeout; private final TimeDuration closeTimeout; @@ -424,38 +364,48 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { LOG.error("{}: unexpected message {}", name, msg.getClass()); return; } - try (DataStreamReplyByteBuf reply = (DataStreamReplyByteBuf) msg) { - process(reply); + final DataStreamReplyByteBuf reply = (DataStreamReplyByteBuf) msg; + final ReferenceCountedObject ref = ReferenceCountedObject.newBuilder() + .setValue((DataStreamReplyByteBuf) msg) + .setReleaseMethod(r -> { + if (r != null) { + Preconditions.assertSame(reply, r, "reply"); + reply.release(); + } + }).build(); + try (UncheckedAutoCloseableSupplier ignored = ref.retainAndReleaseOnClose()) { + process(ref); } } - private void process(DataStreamReplyByteBuf reply) { - LOG.debug("{}: read {}", name, reply); - final ClientInvocationId clientInvocationId = ClientInvocationId.valueOf( - reply.getClientId(), reply.getStreamId()); - final ReadOnlyStreamingReply readOnlyStreamingReply = readOnlyStreamingReplies.get(clientInvocationId); - if (readOnlyStreamingReply != null) { + private void process(ReferenceCountedObject ref) { + final DataStreamReplyByteBuf replyByteBuf = (DataStreamReplyByteBuf) ref.get(); + LOG.debug("{}: read {}", name, replyByteBuf); + final ClientInvocationId clientInvocationId = ClientInvocationId.valueOf(replyByteBuf); + final ClientReadStream clientReadStream = readStreams.get(clientInvocationId); + if (clientReadStream != null) { try { - if (readOnlyStreamingReply.receiveReply(reply)) { - readOnlyStreamingReplies.remove(clientInvocationId, readOnlyStreamingReply); + if (clientReadStream.receiveReply(ref)) { + readStreams.remove(clientInvocationId, clientReadStream); } } catch (Throwable cause) { LOG.warn("{} : channelRead error:", name, cause); - readOnlyStreamingReplies.remove(clientInvocationId, readOnlyStreamingReply); - readOnlyStreamingReply.completeExceptionally(cause); + readStreams.remove(clientInvocationId, clientReadStream); + clientReadStream.completeExceptionally(cause); } return; } + // just copy it for write requests + final DataStreamReplyByteBuffer reply = replyByteBuf.copy(); final NettyClientReplies.ReplyMap replyMap = replies.getReplyMap(clientInvocationId); if (replyMap == null) { LOG.error("{}: {} replyMap not found for reply: {}", name, clientInvocationId, reply); - reply.release(); return; } try { - replyMap.receiveReply(NettyDataStreamUtils.toDataStreamReplyByteBuffer(reply)); + replyMap.receiveReply(reply); } catch (Throwable cause) { LOG.warn("{} : channelRead error for {}:", name, reply.getClass().getSimpleName(), cause); replyMap.completeExceptionally(cause); @@ -595,13 +545,12 @@ public CompletableFuture streamAsync(DataStreamRequest request) } @Override - public CompletableFuture streamAsync( - DataStreamRequest request, DataStreamObserver replyHandler) { + public CompletableFuture streamAsync(DataStreamRequest request, + DataStreamObserver> replyHandler) { final CompletableFuture f = new CompletableFuture<>(); - final ClientInvocationId clientInvocationId = ClientInvocationId.valueOf(request.getClientId(), - request.getStreamId()); - final ReadOnlyStreamingReply replyEntry = new ReadOnlyStreamingReply(request, f, replyHandler); - if (readOnlyStreamingReplies.putIfAbsent(clientInvocationId, replyEntry) != null) { + final ClientInvocationId clientInvocationId = ClientInvocationId.valueOf(request); + final ClientReadStream replyEntry = new ClientReadStream(request, f, replyHandler); + if (readStreams.putIfAbsent(clientInvocationId, replyEntry) != null) { f.completeExceptionally(new AlreadyClosedException(this + ": A read-only stream already exists for " + clientInvocationId)); return f; @@ -613,7 +562,7 @@ public CompletableFuture streamAsync( synchronized (replyEntry) { channel = connection.getChannelUninterruptibly(); if (channel == null) { - readOnlyStreamingReplies.remove(clientInvocationId, replyEntry); + readStreams.remove(clientInvocationId, replyEntry); f.completeExceptionally(new AlreadyClosedException(this + ": Failed to send " + request)); return f; } @@ -623,7 +572,7 @@ public CompletableFuture streamAsync( } channelFuture.addListener(future -> { if (!future.isSuccess()) { - readOnlyStreamingReplies.remove(clientInvocationId, replyEntry); + readStreams.remove(clientInvocationId, replyEntry); final IOException e = new IOException(this + ": Failed to send " + request + " to " + channel.remoteAddress(), future.cause()); replyEntry.completeExceptionally(e); @@ -633,7 +582,7 @@ public CompletableFuture streamAsync( replyEntry.scheduleTimeout(() -> channel.eventLoop().schedule(() -> { if (!f.isDone()) { - readOnlyStreamingReplies.remove(clientInvocationId, replyEntry); + readStreams.remove(clientInvocationId, replyEntry); replyEntry.completeExceptionally(new TimeoutIOException( "Timeout " + requestTimeout + ": Failed to send " + request + " via channel " + channel)); } diff --git a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java index 9cd9217283..b493e49bf8 100644 --- a/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java +++ b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java @@ -23,7 +23,6 @@ import org.apache.ratis.conf.RaftProperties; import org.apache.ratis.datastream.DataStreamObserver; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; -import org.apache.ratis.datastream.impl.DataStreamReplyByteBuffer; import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; import org.apache.ratis.proto.RaftProtos.RaftClientRequestProto; @@ -34,6 +33,7 @@ import org.apache.ratis.protocol.RaftGroupId; import org.apache.ratis.protocol.RaftPeer; import org.apache.ratis.thirdparty.io.netty.buffer.Unpooled; +import org.apache.ratis.util.ReferenceCountedObject; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -50,12 +50,13 @@ private static RaftPeer newPeer(String id) { private static class RecordingDataStreamClientRpc implements DataStreamClientRpc { private final AtomicReference request = new AtomicReference<>(); - private final AtomicReference> replyHandler = new AtomicReference<>(); + private final AtomicReference>> replyHandler = new AtomicReference<>(); private final AtomicReference> replyFuture = new AtomicReference<>(); @Override public CompletableFuture streamAsync( - DataStreamRequest dataStreamRequest, DataStreamObserver replyHandler) { + DataStreamRequest dataStreamRequest, + DataStreamObserver> replyHandler) { try { final ByteBuffer buffer = ((DataStreamRequestByteBuffer) dataStreamRequest).slice(); request.set(ClientProtoUtils.toRaftClientRequest(RaftClientRequestProto.parseFrom(buffer))); @@ -73,8 +74,15 @@ RaftClientRequest getRequest() { } void receive(DataStreamReplyByteBuf reply) { - try (DataStreamReplyByteBuf r = reply) { - replyHandler.get().onNext(r); + final ReferenceCountedObject ref = ReferenceCountedObject.newBuilder() + .setValue(reply) + .setReleaseMethod(DataStreamReplyByteBuf::release) + .build(); + ref.retain(); + try { + replyHandler.get().onNext(ref); + } finally { + ref.release(); } } @@ -107,8 +115,8 @@ public void testReceiveSkipsCancelledPendingRead() throws Exception { try (DataStreamClient dataStreamClient = newDataStreamClient(follower, dataStreamClientRpc); DataStreamInput input = dataStreamClient.streamReadOnly(ByteBuffer.wrap(new byte[] {1}))) { - final CompletableFuture cancelled = input.readAsync(); - final CompletableFuture active = input.readAsync(); + final CompletableFuture> cancelled = input.readAsync(); + final CompletableFuture> active = input.readAsync(); cancelled.cancel(false); Assertions.assertEquals(follower.getId(), dataStreamClientRpc.getRequest().getServerId()); @@ -123,9 +131,11 @@ public void testReceiveSkipsCancelledPendingRead() throws Exception { dataStreamClientRpc.receive(reply); Assertions.assertTrue(active.isDone()); - final DataStreamReply received = active.getNow(null); - Assertions.assertInstanceOf(DataStreamReplyByteBuffer.class, received); - Assertions.assertEquals(Type.STREAM_DATA, received.getType()); + final ReferenceCountedObject received = active.getNow(null); + Assertions.assertNotNull(received); + final DataStreamReplyByteBuf data = Assertions.assertInstanceOf(DataStreamReplyByteBuf.class, received.get()); + Assertions.assertEquals(Type.STREAM_DATA, data.getType()); + received.release(); } } @@ -136,7 +146,7 @@ public void testReadOnlyInputCompletesPendingReadOnCompleted() throws Exception try (DataStreamClient dataStreamClient = newDataStreamClient(follower, dataStreamClientRpc); DataStreamInput input = dataStreamClient.streamReadOnly(ByteBuffer.wrap(new byte[] {1}))) { - final CompletableFuture pending = input.readAsync(); + final CompletableFuture> pending = input.readAsync(); dataStreamClientRpc.complete(); @@ -152,7 +162,7 @@ public void testReadOnlyInputNotifiesPendingReadOnError() throws Exception { try (DataStreamClient dataStreamClient = newDataStreamClient(follower, dataStreamClientRpc); DataStreamInput input = dataStreamClient.streamReadOnly(ByteBuffer.wrap(new byte[] {1}))) { - final CompletableFuture pending = input.readAsync(); + final CompletableFuture> pending = input.readAsync(); final Throwable cause = new IllegalStateException("test"); dataStreamClientRpc.completeExceptionally(cause); diff --git a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java index 02ad13e3ec..95e9bef49c 100644 --- a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java +++ b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamClusterTests.java @@ -18,7 +18,7 @@ package org.apache.ratis.datastream; import org.apache.ratis.BaseTest; -import org.apache.ratis.client.DataStreamClient; +import org.apache.ratis.client.api.DataStreamInput; import org.apache.ratis.io.StandardWriteOption; import org.apache.ratis.protocol.RaftPeer; import org.apache.ratis.protocol.RoutingTable; @@ -29,7 +29,6 @@ import org.apache.ratis.datastream.DataStreamTestUtils.MultiDataStreamStateMachine; import org.apache.ratis.datastream.DataStreamTestUtils.SingleDataStream; import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf; -import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer; import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; import org.apache.ratis.proto.RaftProtos.ReplicationLevel; import org.apache.ratis.protocol.DataStreamReply; @@ -39,24 +38,24 @@ import org.apache.ratis.protocol.RaftClientReply; import org.apache.ratis.protocol.RaftClientRequest; import org.apache.ratis.retry.RetryPolicies; -import org.apache.ratis.rpc.CallId; import org.apache.ratis.server.RaftServer; import org.apache.ratis.server.RaftServerConfigKeys; import org.apache.ratis.thirdparty.com.google.protobuf.ByteString; import org.apache.ratis.util.CollectionUtils; import org.apache.ratis.util.FileUtils; +import org.apache.ratis.util.ReferenceCountedObject; import org.apache.ratis.util.Timestamp; import org.apache.ratis.util.function.CheckedConsumer; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.io.EOFException; import java.io.File; -import java.nio.ByteBuffer; import java.nio.channels.FileChannel; import java.nio.file.StandardOpenOption; import java.util.Collection; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; @@ -150,59 +149,33 @@ void runTestStreamReadOnlyWithNonLeaderPrimary(CLUSTER cluster) throws Exception void runTestStreamReadOnly(CLUSTER cluster, RaftPeer leader, RaftPeer primaryServer) throws Exception { final ByteString query = ByteString.copyFromUtf8("stream-read-only"); - try (RaftClient client = cluster.createClient(leader.getId(), cluster.getGroup(), RetryPolicies.noRetry(), - primaryServer)) { - final AtomicInteger replyCount = new AtomicInteger(); - final CompletableFuture completed = new CompletableFuture<>(); - final DataStreamClient dataStreamClient = (DataStreamClient) client.getDataStreamApi(); - final CompletableFuture terminalReplyFuture = dataStreamClient.getClientRpc().streamAsync( - newReadOnlyStreamRequest(client, primaryServer, query), new DataStreamObserver() { - @Override - public void onNext(DataStreamReplyByteBuf reply) { - final int i = replyCount.getAndIncrement(); - Assertions.assertTrue(i < MultiDataStreamStateMachine.READ_ONLY_STREAM_CHUNKS); - - final ByteString chunk = MultiDataStreamStateMachine.getReadOnlyStreamChunk(query, i); - DataStreamTestUtils.assertSuccessReply(Type.STREAM_DATA, chunk.size(), reply); - Assertions.assertEquals(chunk, ByteString.copyFrom(reply.slice().nioBuffer())); - } - - @Override - public void onError(Throwable t) { - completed.completeExceptionally(t); - } - - @Override - public void onCompleted() { - completed.complete(null); - } - }); - - completed.get(1, TimeUnit.SECONDS); - Assertions.assertEquals(MultiDataStreamStateMachine.READ_ONLY_STREAM_CHUNKS, replyCount.get()); - final DataStreamReply terminalReply = terminalReplyFuture.get(1, TimeUnit.SECONDS); + try (RaftClient client = cluster.createClient(leader.getId(), cluster.getGroup(), RetryPolicies.noRetry(), primaryServer); + DataStreamInput in = client.getDataStreamApi().streamReadOnly(query.asReadOnlyByteBuffer())) { + for (int i = 0; i < MultiDataStreamStateMachine.READ_ONLY_STREAM_CHUNKS; i++) { + final ByteString chunk = MultiDataStreamStateMachine.getReadOnlyStreamChunk(query, i); + final ReferenceCountedObject ref = in.readAsync().join(); + final DataStreamReplyByteBuf data = (DataStreamReplyByteBuf) ref.get(); + DataStreamTestUtils.assertSuccessReply(Type.STREAM_DATA, chunk.size(), data); + Assertions.assertEquals(chunk, ByteString.copyFrom(data.slice().nioBuffer())); + ref.release(); + } + + final ReferenceCountedObject ref = in.readAsync().join(); try { - DataStreamTestUtils.assertSuccessReply(Type.STREAM_HEADER, 0, terminalReply); + final DataStreamReplyByteBuf reply = (DataStreamReplyByteBuf) ref.get(); + DataStreamTestUtils.assertSuccessReply(Type.STREAM_HEADER, 0, reply); + + final RaftClientReply clientReply = ClientProtoUtils.getRaftClientReply(reply); + Assertions.assertTrue(clientReply.isSuccess()); + Assertions.assertEquals(primaryServer.getId(), clientReply.getServerId()); } finally { - DataStreamReplyByteBuf.release(terminalReply); + ref.release(); } - } - } - private static DataStreamRequest newReadOnlyStreamRequest(RaftClient client, RaftPeer primaryServer, - ByteString query) { - final RaftClientRequest request = RaftClientRequest.newBuilder() - .setClientId(client.getId()) - .setServerId(primaryServer.getId()) - .setGroupId(client.getGroupId()) - .setCallId(CallId.getAndIncrement()) - .setMessage(Message.valueOf(query)) - .setType(RaftClientRequest.readRequestType()) - .build(); - final ByteBuffer buffer = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(request); - final DataStreamRequestHeader header = new DataStreamRequestHeader(request.getClientId(), Type.STREAM_HEADER, - request.getCallId(), 0, buffer.remaining(), StandardWriteOption.FLUSH, StandardWriteOption.CLOSE); - return new DataStreamRequestByteBuffer(header, buffer); + final ExecutionException eof = Assertions.assertThrows(ExecutionException.class, + () -> in.readAsync().get(1, TimeUnit.SECONDS)); + Assertions.assertInstanceOf(EOFException.class, eof.getCause()); + } } void runTestInvalidPrimaryInRoutingTable(CLUSTER cluster) throws Exception { diff --git a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java index 419f21e405..188f119fca 100644 --- a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java +++ b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java @@ -139,15 +139,38 @@ public void query(Message request, WritableByteChannel stream) { queryDone.countDown(); } }; - final ReadStreamManagement management = newReadStreamManagement(serverId, groupId, dataApi); - final ReadOnlyRequest readOnlyRequest = newReadOnlyRequest( - clientId, serverId, groupId, 1L, ByteString.copyFromUtf8("query")); + final StateMachine stateMachine = new BaseStateMachine() { + @Override + public DataApi data() { + return dataApi; + } + }; + final RaftServer server = newRaftServer(serverId, new RaftProperties(), groupId, newDivision(stateMachine)); + final ReadStreamManagement management = new ReadStreamManagement(server); + + final RaftClientRequest raftClientRequest = RaftClientRequest.newBuilder() + .setClientId(clientId) + .setServerId(serverId) + .setGroupId(groupId) + .setCallId(1L) + .setMessage(Message.valueOf(ByteString.copyFromUtf8("query"))) + .setType(RaftClientRequest.readRequestType()) + .build(); + final ByteBuffer header = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(raftClientRequest); + final ByteBuf headerBuf = Unpooled.wrappedBuffer(header); + final DataStreamRequestByteBuf request = new DataStreamRequestByteBuf( + clientId, + Type.STREAM_HEADER, + raftClientRequest.getCallId(), + 0L, + Collections.singletonList(StandardWriteOption.FLUSH), + headerBuf); try { - eventLoop.submit(() -> assertTrue(management.process(readOnlyRequest.request, ctx))).sync(); + eventLoop.submit(() -> assertTrue(management.process(request, ctx))).sync(); assertTrue(queryDone.await(10, TimeUnit.SECONDS)); - assertEquals(0, readOnlyRequest.headerBuf.refCnt()); + assertEquals(0, headerBuf.refCnt()); assertFalse(queryInEventLoop.get(), "read-only state machine query should not run on Netty event loop"); } finally { embeddedChannel.finishAndReleaseAll();