This is an automated email from the ASF dual-hosted git repository.
szetszwo pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/ratis.git
The following commit(s) were added to refs/heads/master by this push:
new abfe70f34 RATIS-2546. Add input stream to DataStreamApi for read
operations in Client (#1481)
abfe70f34 is described below
commit abfe70f34cc63b66049ed13d880232972c2f7e3d
Author: Peter Lee <[email protected]>
AuthorDate: Wed Jun 24 00:49:26 2026 +0800
RATIS-2546. Add input stream to DataStreamApi for read operations in Client
(#1481)
---
.../apache/ratis/client/DataStreamClientRpc.java | 15 +++
.../org/apache/ratis/client/api/DataStreamApi.java | 8 ++
.../DataStreamInput.java} | 25 ++--
.../ratis/client/impl/DataStreamClientImpl.java | 7 +
.../ratis/client/impl/DataStreamInputImpl.java | 143 ++++++++++++++++++++
.../ratis/datastream/DataStreamObserver.java | 39 ++++++
.../datastream/impl/DataStreamPacketByteBuf.java | 5 +-
.../datastream/impl/DataStreamReplyByteBuf.java | 15 +++
.../java/org/apache/ratis/util/Preconditions.java | 3 +-
.../apache/ratis/netty/NettyDataStreamUtils.java | 24 ++++
.../ratis/netty/client/ClientReadStream.java | 80 +++++++++++
.../ratis/netty/client/NettyClientReplies.java | 3 +-
.../ratis/netty/client/NettyClientStreamRpc.java | 86 +++++++++++-
.../client/impl/TestDataStreamClientImpl.java | 146 +++++++++++++++++++++
.../ratis/datastream/DataStreamClusterTests.java | 73 +++++++++++
.../netty/server/TestDataStreamManagement.java | 129 ++++++++++++++----
16 files changed, 754 insertions(+), 47 deletions(-)
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 a9bcd9d58..6f94709b0 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,9 +18,11 @@
package org.apache.ratis.client;
+import org.apache.ratis.datastream.DataStreamObserver;
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;
@@ -36,4 +38,17 @@ public interface DataStreamClientRpc extends Closeable {
throw new UnsupportedOperationException(getClass() + " does not support "
+ JavaUtils.getCurrentStackTraceElement().getMethodName());
}
+
+ /**
+ * 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<DataStreamReply> streamAsync(DataStreamRequest
request,
+ DataStreamObserver<ReferenceCountedObject<DataStreamReply>>
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 9e5e2438c..7152d35ba 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,12 @@ public interface DataStreamApi extends Closeable {
/** Create a stream by providing a customized header message and route
table. */
DataStreamOutput stream(ByteBuffer headerMessage, RoutingTable routingTable);
+
+ /** Create a stream to read data. */
+ default DataStreamInput streamReadOnly() {
+ return streamReadOnly(null);
+ }
+
+ /** 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/DataStreamClientRpc.java
b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java
similarity index 56%
copy from
ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java
copy to
ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java
index a9bcd9d58..e50033405 100644
---
a/ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java
+++
b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java
@@ -15,25 +15,24 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-
-package org.apache.ratis.client;
+package org.apache.ratis.client.api;
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;
/**
- * A client interface for sending stream requests.
- * The underlying implementation is pluggable, depending on the {@link
org.apache.ratis.datastream.DataStreamType}.
- * The implementations of this interface define how the requests are
transported to the server.
+ * An asynchronous input stream supporting zero buffer copying.
*/
-public interface DataStreamClientRpc extends Closeable {
- /** Async call to send a request. */
- default CompletableFuture<DataStreamReply> streamAsync(DataStreamRequest
request) {
- throw new UnsupportedOperationException(getClass() + " does not support "
- + JavaUtils.getCurrentStackTraceElement().getMethodName());
- }
+public interface DataStreamInput extends Closeable {
+ /**
+ * Read the next chunk in the stream asynchronously.
+ * 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 reference-counted reply.
+ */
+ CompletableFuture<ReferenceCountedObject<DataStreamReply>> 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 3c055009f..bc8334bf9 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,6 +23,7 @@ import org.apache.ratis.client.DataStreamClient;
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;
@@ -265,6 +266,12 @@ public class DataStreamClientImpl implements
DataStreamClient {
.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()
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 000000000..f1e38df88
--- /dev/null
+++
b/ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamInputImpl.java
@@ -0,0 +1,143 @@
+/*
+ * 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.Preconditions;
+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<ReferenceCountedObject<DataStreamReply>> {
+ private final RaftClientRequest header;
+ private final ClientId clientId;
+ private final Queue<ReferenceCountedObject<DataStreamReply>> replies = new
LinkedList<>();
+ private final
Queue<CompletableFuture<ReferenceCountedObject<DataStreamReply>>> pendingReads
= new LinkedList<>();
+
+ /*
+ * null : the stream is open.
+ * EOFException : the stream is ended, i.e. no more data.
+ * AlreadyClosedException: the stream is closed by the caller.
+ * 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<DataStreamReply>
reply) {
+ if (readException != null) {
+ return;
+ }
+
+ reply.retain();
+ final CompletableFuture<ReferenceCountedObject<DataStreamReply>> pending =
pendingReads.poll();
+ if (pending != null) {
+ final boolean completed = pending.complete(reply);
+ Preconditions.assertTrue(completed);
+ return;
+ }
+ replies.add(reply);
+ }
+
+ @Override
+ public synchronized void onError(Throwable throwable) {
+ Objects.requireNonNull(throwable, "throwable == null");
+ // 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<DataStreamReply> reply; (reply =
replies.poll()) != null; ) {
+ reply.release();
+ }
+ }
+
+ private void failPendingReads() {
+ Objects.requireNonNull(readException, "readException == null");
+ for (CompletableFuture<ReferenceCountedObject<DataStreamReply>> p; (p =
pendingReads.poll()) != null; ) {
+ p.completeExceptionally(readException);
+ }
+ }
+
+ @Override
+ public synchronized
CompletableFuture<ReferenceCountedObject<DataStreamReply>> readAsync() {
+ final ReferenceCountedObject<DataStreamReply> reply = replies.poll();
+ if (reply != null) {
+ return CompletableFuture.completedFuture(reply);
+ }
+ if (readException != null) {
+ return JavaUtils.completeExceptionally(readException);
+ }
+ final CompletableFuture<ReferenceCountedObject<DataStreamReply>> 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
new file mode 100644
index 000000000..92f62b751
--- /dev/null
+++
b/ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java
@@ -0,0 +1,39 @@
+/*
+ * 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;
+
+import java.util.function.BiConsumer;
+
+/** An interface similar to gRPC {@link
org.apache.ratis.thirdparty.io.grpc.stub.StreamObserver}. */
+public interface DataStreamObserver<V> {
+ void onNext(V value);
+
+ void onCompleted();
+
+ void onError(Throwable throwable);
+
+ default <T> BiConsumer<T, Throwable> 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 76400ff80..e6ceeb93d 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 @@ import java.util.concurrent.atomic.AtomicReference;
* <p>
* This class is immutable.
*/
-class DataStreamPacketByteBuf extends DataStreamPacketImpl {
+public class DataStreamPacketByteBuf extends DataStreamPacketImpl {
private final AtomicReference<ByteBuf> buf;
DataStreamPacketByteBuf(ClientId clientId, Type type, long streamId, long
streamOffset, ByteBuf buf) {
@@ -48,7 +48,8 @@ class DataStreamPacketByteBuf extends DataStreamPacketImpl {
@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-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 2477b089b..ab2849b19 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
@@ -21,10 +21,13 @@ import org.apache.ratis.proto.RaftProtos.CommitInfoProto;
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.util.Preconditions;
+import org.apache.ratis.util.ReferenceCountedObject;
import java.nio.ByteBuffer;
import java.util.Collection;
import java.util.Collections;
+import java.util.Objects;
/**
* Implements {@link DataStreamReply} with {@link ByteBuf}.
@@ -64,6 +67,18 @@ public final class DataStreamReplyByteBuf extends
DataStreamPacketByteBuf implem
return new Builder();
}
+ public static ReferenceCountedObject<DataStreamReply>
asReferenceCounted(DataStreamReplyByteBuf reply) {
+ Objects.requireNonNull(reply, "reply == null");
+ return ReferenceCountedObject.<DataStreamReply>newBuilder()
+ .setValue(reply)
+ .setReleaseMethod(r -> {
+ if (r != null) {
+ Preconditions.assertSame(reply, r, "reply").release();
+ }
+ })
+ .build();
+ }
+
private final boolean success;
private final long bytesWritten;
private final Collection<CommitInfoProto> commitInfos;
diff --git
a/ratis-common/src/main/java/org/apache/ratis/util/Preconditions.java
b/ratis-common/src/main/java/org/apache/ratis/util/Preconditions.java
index f37b250f5..b2498dac1 100644
--- a/ratis-common/src/main/java/org/apache/ratis/util/Preconditions.java
+++ b/ratis-common/src/main/java/org/apache/ratis/util/Preconditions.java
@@ -88,9 +88,10 @@ public interface Preconditions {
() -> name + ": expected == " + expected + " but computed == " +
computed);
}
- static void assertSame(Object expected, Object computed, String name) {
+ static <T> T assertSame(T expected, Object computed, String name) {
assertTrue(expected == computed,
() -> name + ": expected == " + expected + " but computed == " +
computed);
+ return expected;
}
static void assertEquals(Object expected, Object computed, String name) {
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 2c5015b12..9f8cbecc8 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
@@ -169,6 +169,22 @@ public interface NettyDataStreamUtils {
encodeDataStreamReply(reply, reply.slice(), out, allocator);
}
+ static void encodeDataStreamReplyByteBuf(DataStreamReplyByteBuf reply,
Consumer<ByteBuf> 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))
@@ -228,6 +244,14 @@ public interface NettyDataStreamUtils {
.orElse(null);
}
+ static DataStreamReplyByteBuffer
toDataStreamReplyByteBuffer(DataStreamReplyByteBuf reply) {
+ try {
+ return reply.copy();
+ } finally {
+ reply.release();
+ }
+ }
+
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/ClientReadStream.java
b/ratis-netty/src/main/java/org/apache/ratis/netty/client/ClientReadStream.java
new file mode 100644
index 000000000..c9096ae7b
--- /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<DataStreamReply> replyFuture;
+ private final DataStreamObserver<ReferenceCountedObject<DataStreamReply>>
replyHandler;
+ private Supplier<ScheduledFuture<?>> timeoutScheduler;
+ private ScheduledFuture<?> timeoutFuture;
+
+ ClientReadStream(DataStreamRequest request,
CompletableFuture<DataStreamReply> replyFuture,
+ DataStreamObserver<ReferenceCountedObject<DataStreamReply>>
replyHandler) {
+ this.terminalEntry = new NettyClientReplies.RequestEntry(request);
+ this.replyFuture = replyFuture;
+ this.replyHandler = replyHandler;
+ }
+
+ synchronized boolean receiveReply(ReferenceCountedObject<DataStreamReply>
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<ScheduledFuture<?>>
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/NettyClientReplies.java
b/ratis-netty/src/main/java/org/apache/ratis/netty/client/NettyClientReplies.java
index 121ef3c0f..0aa6bc772 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 @@ public class NettyClientReplies {
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 ccb0b15da..4f4f0d1f0 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,8 +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.impl.DataStreamReplyByteBuffer;
import org.apache.ratis.datastream.impl.DataStreamRequestByteBuf;
import org.apache.ratis.datastream.impl.DataStreamRequestByteBuffer;
import org.apache.ratis.datastream.impl.DataStreamReplyByteBuf;
@@ -64,6 +66,7 @@ import org.apache.ratis.util.Preconditions;
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;
@@ -74,6 +77,8 @@ import java.util.Collections;
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;
@@ -325,6 +330,7 @@ public class NettyClientStreamRpc implements
DataStreamClientRpc {
private final Connection connection;
private final NettyClientReplies replies = new NettyClientReplies();
+ private final ConcurrentMap<ClientInvocationId, ClientReadStream>
readStreams = new ConcurrentHashMap<>();
private final TimeDuration requestTimeout;
private final TimeDuration closeTimeout;
@@ -358,15 +364,33 @@ public class NettyClientStreamRpc implements
DataStreamClientRpc {
LOG.error("{}: unexpected message {}", name, msg.getClass());
return;
}
- try (DataStreamReplyByteBuf reply = (DataStreamReplyByteBuf) msg) {
- process(reply);
+ final DataStreamReplyByteBuf reply = (DataStreamReplyByteBuf) msg;
+ final ReferenceCountedObject<DataStreamReply> ref =
DataStreamReplyByteBuf.asReferenceCounted(reply);
+ try (UncheckedAutoCloseableSupplier<DataStreamReply> ignored =
ref.retainAndReleaseOnClose()) {
+ process(ref);
}
}
- private void process(DataStreamReplyByteBuf reply) {
- LOG.debug("{}: read {}", name, reply);
- final ClientInvocationId clientInvocationId =
ClientInvocationId.valueOf(
- reply.getClientId(), reply.getStreamId());
+ private void process(ReferenceCountedObject<DataStreamReply> 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 (clientReadStream.receiveReply(ref)) {
+ readStreams.remove(clientInvocationId, clientReadStream);
+ }
+ } catch (Throwable cause) {
+ LOG.warn("{} : channelRead error:", name, 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);
@@ -376,7 +400,7 @@ public class NettyClientStreamRpc implements
DataStreamClientRpc {
try {
replyMap.receiveReply(reply);
} catch (Throwable cause) {
- LOG.warn("{} : channelRead error for {}", name, reply, cause);
+ LOG.warn("{} : channelRead error for {}:", name,
reply.getClass().getSimpleName(), cause);
replyMap.completeExceptionally(cause);
}
}
@@ -513,6 +537,54 @@ public class NettyClientStreamRpc implements
DataStreamClientRpc {
return f;
}
+ @Override
+ public CompletableFuture<DataStreamReply> streamAsync(DataStreamRequest
request,
+ DataStreamObserver<ReferenceCountedObject<DataStreamReply>>
replyHandler) {
+ final CompletableFuture<DataStreamReply> f = new CompletableFuture<>();
+ final ClientInvocationId clientInvocationId =
ClientInvocationId.valueOf(request);
+ final ClientReadStream readStream = new ClientReadStream(request, f,
replyHandler);
+ if (readStreams.putIfAbsent(clientInvocationId, readStream) != null) {
+ f.completeExceptionally(new AlreadyClosedException(this + ": Read stream
already exists for "
+ + clientInvocationId + ", request=" + request));
+ return f;
+ }
+
+ final ChannelFuture channelFuture;
+ final Channel channel;
+ LOG.debug("{}: write read-only stream begin {}", this, request);
+ synchronized (readStream) {
+ channel = connection.getChannelUninterruptibly();
+ if (channel == null) {
+ readStreams.remove(clientInvocationId, readStream);
+ f.completeExceptionally(new AlreadyClosedException(this + ": Failed to
getChannel for " + request));
+ return f;
+ }
+ final Function<DataStreamRequest, ChannelFuture> writeMethod =
outstandingRequests.shouldFlush(
+ flushRequestCountMin, flushRequestBytesMin, request)?
channel::writeAndFlush: channel::write;
+ channelFuture = writeMethod.apply(request);
+ }
+ channelFuture.addListener(future -> {
+ if (!future.isSuccess()) {
+ readStreams.remove(clientInvocationId, readStream);
+ final IOException e = new IOException(this + ": Failed to send " +
request + " to " + channel.remoteAddress(),
+ future.cause());
+ readStream.completeExceptionally(e);
+ LOG.error("Channel write failed", e);
+ } else {
+ LOG.debug("{}: write read-only stream after {}", this, request);
+
+ readStream.scheduleTimeout(() -> channel.eventLoop().schedule(() -> {
+ if (!f.isDone()) {
+ readStreams.remove(clientInvocationId, readStream);
+ readStream.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
new file mode 100644
index 000000000..523c49f05
--- /dev/null
+++
b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java
@@ -0,0 +1,146 @@
+/*
+ * 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.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.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.RaftClientRequest;
+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;
+
+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 {
+ private static RaftPeer newPeer(String id) {
+ return RaftPeer.newBuilder().setId(id).build();
+ }
+
+ private static class RecordingDataStreamClientRpc implements
DataStreamClientRpc {
+ private final AtomicReference<RaftClientRequest> request = new
AtomicReference<>();
+ private final
AtomicReference<DataStreamObserver<ReferenceCountedObject<DataStreamReply>>>
replyHandler = new AtomicReference<>();
+ private final AtomicReference<CompletableFuture<DataStreamReply>>
replyFuture = new AtomicReference<>();
+
+ @Override
+ public CompletableFuture<DataStreamReply> streamAsync(
+ DataStreamRequest dataStreamRequest,
+ DataStreamObserver<ReferenceCountedObject<DataStreamReply>>
replyHandler) {
+ try {
+ final ByteBuffer buffer = ((DataStreamRequestByteBuffer)
dataStreamRequest).slice();
+
request.set(ClientProtoUtils.toRaftClientRequest(RaftClientRequestProto.parseFrom(buffer)));
+ } catch (Exception e) {
+ throw new IllegalStateException(e);
+ }
+ this.replyHandler.set(replyHandler);
+ final CompletableFuture<DataStreamReply> future = new
CompletableFuture<>();
+ replyFuture.set(future);
+ return future;
+ }
+
+ RaftClientRequest getRequest() {
+ return request.get();
+ }
+
+ void receive(DataStreamReplyByteBuf reply) {
+ final ReferenceCountedObject<DataStreamReply> ref =
DataStreamReplyByteBuf.asReferenceCounted(reply);
+ ref.retain();
+ try {
+ replyHandler.get().onNext(ref);
+ } finally {
+ ref.release();
+ }
+ }
+
+ 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() {
+ }
+ }
+
+ 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 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<ReferenceCountedObject<DataStreamReply>> 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<ReferenceCountedObject<DataStreamReply>> 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 <T extends Throwable> Throwable assertFutureCause(
+ CompletableFuture<?> future, Class<T> 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 dabc93dda..95e9bef49 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,33 +18,46 @@
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.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.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.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 +84,17 @@ public abstract class DataStreamClusterTests<CLUSTER extends
MiniRaftCluster> ex
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 +129,55 @@ public abstract class DataStreamClusterTests<CLUSTER
extends MiniRaftCluster> ex
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 ReferenceCountedObject<DataStreamReply> 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<DataStreamReply> ref =
in.readAsync().join();
+ try {
+ 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 {
+ ref.release();
+ }
+
+ 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 {
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 516f8c171..188f119fc 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.buffer.Unpooled;
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.ArrayList;
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;
@@ -82,6 +87,58 @@ class TestDataStreamManagement {
streamRef.set(stream);
}
};
+ final ReadStreamManagement management = newReadStreamManagement(serverId,
groupId, dataApi);
+ final EmbeddedChannel embeddedChannel = new EmbeddedChannel(new
ChannelInboundHandlerAdapter());
+
+ final ReadOnlyRequest readOnlyRequest = newReadOnlyRequest(clientId,
serverId, groupId, 1L, query);
+
+ try {
+ assertTrue(management.process(readOnlyRequest.request,
embeddedChannel.pipeline().firstContext()));
+ assertEquals(0, readOnlyRequest.headerBuf.refCnt());
+
+ final WritableByteChannel stream = streamRef.get();
+ assertNotNull(stream);
+ stream.write(response.asReadOnlyByteBuffer());
+ stream.close();
+
+ final List<DataStreamReply> replies = new ArrayList<>();
+ JavaUtils.attempt(() -> {
+ for (Object outbound; (outbound = embeddedChannel.readOutbound()) !=
null;) {
+ replies.add((DataStreamReply) outbound);
+ }
+ assertEquals(2, replies.size());
+ }, 10, TimeDuration.valueOf(100, TimeUnit.MILLISECONDS), "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();
+ }
+ }
+
+ @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() {
@@ -90,14 +147,13 @@ class TestDataStreamManagement {
};
final RaftServer server = newRaftServer(serverId, new RaftProperties(),
groupId, newDivision(stateMachine));
final ReadStreamManagement management = new ReadStreamManagement(server);
- final EmbeddedChannel embeddedChannel = new EmbeddedChannel(new
ChannelInboundHandlerAdapter());
final RaftClientRequest raftClientRequest = RaftClientRequest.newBuilder()
.setClientId(clientId)
.setServerId(serverId)
.setGroupId(groupId)
.setCallId(1L)
- .setMessage(Message.valueOf(query))
+ .setMessage(Message.valueOf(ByteString.copyFromUtf8("query")))
.setType(RaftClientRequest.readRequestType())
.build();
final ByteBuffer header =
ClientProtoUtils.toRaftClientRequestProtoByteBuffer(raftClientRequest);
@@ -111,29 +167,15 @@ class TestDataStreamManagement {
headerBuf);
try {
- assertTrue(management.process(request,
embeddedChannel.pipeline().firstContext()));
- assertEquals(0, headerBuf.refCnt());
+ eventLoop.submit(() -> assertTrue(management.process(request,
ctx))).sync();
- final WritableByteChannel stream = streamRef.get();
- assertNotNull(stream);
- stream.write(response.asReadOnlyByteBuffer());
- stream.close();
-
- final List<DataStreamReply> replies = new ArrayList<>();
- JavaUtils.attempt(() -> {
- for (Object outbound; (outbound = embeddedChannel.readOutbound()) !=
null;) {
- replies.add((DataStreamReply) outbound);
- }
- assertEquals(2, replies.size());
- }, 10, TimeDuration.valueOf(100, TimeUnit.MILLISECONDS), "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());
+ 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();
}
}
@@ -177,6 +219,49 @@ class TestDataStreamManagement {
}
}
+ 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());
