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..ddeb10b40c 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..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 @@ -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 @@ default CompletableFuture streamAsync(DataStreamRequest request 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 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 9e5e2438cb..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 @@ -50,4 +50,12 @@ 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. */ + 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/api/DataStreamInput.java b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java new file mode 100644 index 0000000000..e50033405f --- /dev/null +++ b/ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java @@ -0,0 +1,38 @@ +/* + * 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 org.apache.ratis.util.ReferenceCountedObject; + +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} which is a {@link ReferenceCountedObject}. + * It must call {@link ReferenceCountedObject#release()} after consuming it. + * + * @return a future of the reference-counted reply. + */ + 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 3c055009f1..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 @@ -23,6 +23,7 @@ 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 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() 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/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/DataStreamObserver.java b/ratis-common/src/main/java/org/apache/ratis/datastream/DataStreamObserver.java new file mode 100644 index 0000000000..92f62b7518 --- /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 { + void onNext(V value); + + void onCompleted(); + + void onError(Throwable throwable); + + 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-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 459aee363c..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 @@ -30,4 +30,4 @@ public interface DataStreamReply extends DataStreamPacket { /** @return the commit information when the reply is created. */ Collection getCommitInfos(); -} \ No newline at end of file +} 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 2c5015b123..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 @@ -169,6 +169,22 @@ static void encodeDataStreamReply(DataStreamReplyByteBuf reply, Consumer encodeDataStreamReply(reply, reply.slice(), out, allocator); } + 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)) @@ -228,6 +244,14 @@ static DataStreamReplyByteBuf decodeDataStreamReplyByteBuf(ByteBuf buf) { .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 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/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 ccb0b15da0..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,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.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.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 @@ synchronized boolean shouldFlush(int countMin, SizeInBytes bytesMin, DataStreamR private final Connection connection; private final NettyClientReplies replies = new NettyClientReplies(); + private final ConcurrentMap readStreams = new ConcurrentHashMap<>(); private final TimeDuration requestTimeout; private final TimeDuration closeTimeout; @@ -358,15 +364,40 @@ 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()); + 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 (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 +407,7 @@ private void process(DataStreamReplyByteBuf reply) { 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 +544,54 @@ public CompletableFuture streamAsync(DataStreamRequest request) return f; } + @Override + public CompletableFuture streamAsync(DataStreamRequest request, + DataStreamObserver> replyHandler) { + final CompletableFuture f = new CompletableFuture<>(); + 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; + } + + final ChannelFuture channelFuture; + final Channel channel; + LOG.debug("{}: write read-only stream begin {}", this, request); + synchronized (replyEntry) { + channel = connection.getChannelUninterruptibly(); + if (channel == null) { + readStreams.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()) { + readStreams.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()) { + readStreams.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-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 5336760a0b..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 @@ -64,20 +64,19 @@ static class ReadStream implements WritableByteChannel { this.clientId = request.getClientId(); this.streamId = streamId; this.ctx = ctx; - final RaftClientReply reply = RaftClientReply.newBuilder() - .setRequest(request) - .setSuccess() - .build(); + .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(); + .setClientId(clientId) + .setType(Type.STREAM_HEADER) + .setStreamId(streamId) + .setStreamOffset(0) + .setBuffer(toRaftClientReplyProto(reply).toByteString().asReadOnlyByteBuffer()) + .setSuccess(true) + .setBytesWritten(0) + .build(); } @Override 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..b493e49bf8 --- /dev/null +++ b/ratis-test/src/test/java/org/apache/ratis/client/impl/TestDataStreamClientImpl.java @@ -0,0 +1,180 @@ +/* + * 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 request = new AtomicReference<>(); + private final AtomicReference>> replyHandler = new AtomicReference<>(); + private final AtomicReference> replyFuture = new AtomicReference<>(); + + @Override + public CompletableFuture streamAsync( + 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.replyHandler.set(replyHandler); + final CompletableFuture future = new CompletableFuture<>(); + replyFuture.set(future); + return future; + } + + RaftClientRequest getRequest() { + return request.get(); + } + + void receive(DataStreamReplyByteBuf reply) { + final ReferenceCountedObject ref = ReferenceCountedObject.newBuilder() + .setValue(reply) + .setReleaseMethod(DataStreamReplyByteBuf::release) + .build(); + 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 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); + Assertions.assertEquals(follower.getId(), dataStreamClientRpc.getRequest().getServerId()); + + final DataStreamReplyByteBuf reply = DataStreamReplyByteBuf.newBuilder() + .setClientId(ClientId.randomId()) + .setType(Type.STREAM_DATA) + .setStreamId(1) + .setStreamOffset(0) + .setBuf(Unpooled.EMPTY_BUFFER) + .setSuccess(true) + .build(); + dataStreamClientRpc.receive(reply); + + Assertions.assertTrue(active.isDone()); + 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(); + } + } + + @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 dabc93dda2..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,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 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 +129,55 @@ 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 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 { + 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 516f8c1714..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 @@ -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; @@ -82,6 +87,58 @@ public void query(Message request, WritableByteChannel stream) { 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 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 @@ public DataApi data() { }; 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 @@ public DataApi data() { 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 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 @@ 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());