[WIP] RATIS-2546. Add input stream to DataStreamApi for read operations in Client#1
Conversation
Signed-off-by: peterxcli <peterxcli@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a read-only streaming API (streamReadOnly) to the Ratis client, allowing clients to read data asynchronously via a stream. This includes adding the DataStreamInput interface and its implementation DataStreamInputImpl, supporting multiple replies for a single request, and introducing DataStreamReplyByteBuf to handle Netty ByteBuf replies. The review feedback highlights two critical memory leak issues related to unreleased Netty ByteBuf resources: one in DataStreamClientImpl.receive when futures are cancelled or replies are received after the stream is closed/ended, and another in NettyClientStreamRpc.channelRead when handling exceptions in the read-only streaming path.
| private void receive(DataStreamReply reply) { | ||
| final CompletableFuture<DataStreamReply> pending; | ||
| synchronized (this) { | ||
| if (closed) { | ||
| reply.release(); | ||
| return; | ||
| } | ||
| pending = pendingReads.poll(); | ||
| if (pending == null) { | ||
| replies.add(reply); | ||
| return; | ||
| } | ||
| } | ||
| pending.complete(reply); | ||
| } |
There was a problem hiding this comment.
Potential memory leak of DataStreamReply (which holds Netty ByteBuf resources) in two scenarios:
- If the client cancels or times out the future returned by
readAsync(),pending.complete(reply)will returnfalsebecause the future is already completed. The reply is never released. - If replies are received after the stream is closed, failed, or ended, they are added to the
repliesqueue and might be leaked if the client does not callclose()immediately.
We should release the reply immediately if the stream is closed/failed/ended, or if the pending future was already completed.
private void receive(DataStreamReply reply) {
final CompletableFuture<DataStreamReply> pending;
synchronized (this) {
if (closed || readException != null || endOfStream) {
reply.release();
return;
}
pending = pendingReads.poll();
if (pending == null) {
replies.add(reply);
return;
}
}
if (!pending.complete(reply)) {
reply.release();
}
}| } catch (Throwable cause) { | ||
| LOG.warn("{} : channelRead error:", name, cause); | ||
| readOnlyStreamingReplies.remove(clientInvocationId, readOnlyStreamingReply); | ||
| readOnlyStreamingReply.completeExceptionally(cause); | ||
| } |
There was a problem hiding this comment.
In the read-only path, if readOnlyStreamingReply.receiveReply(reply) throws an exception, the reply is not released in the catch block, leading to a Netty ByteBuf leak. This is inconsistent with the non-read-only path below (line 453) where reply.release() is correctly called in the catch block.
| } catch (Throwable cause) { | |
| LOG.warn("{} : channelRead error:", name, cause); | |
| readOnlyStreamingReplies.remove(clientInvocationId, readOnlyStreamingReply); | |
| readOnlyStreamingReply.completeExceptionally(cause); | |
| } | |
| } catch (Throwable cause) { | |
| LOG.warn("{} : channelRead error:", name, cause); | |
| readOnlyStreamingReplies.remove(clientInvocationId, readOnlyStreamingReply); | |
| readOnlyStreamingReply.completeExceptionally(cause); | |
| reply.release(); | |
| } |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 68c738c4cb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| closed = true; | ||
| for (DataStreamReply reply; (reply = replies.poll()) != null;) { | ||
| reply.release(); | ||
| } | ||
| failReads(new AlreadyClosedException(clientId + ": stream already closed, request=" + header)); |
There was a problem hiding this comment.
DataStreamInputImpl.close() only marks the local stream closed and fails pending reads, but it never signals the underlying read-only RPC to terminate. The NettyClientStreamRpc entry stays in readOnlyStreamingReplies until a terminal reply or timeout (streamAsync(..., Consumer) cleanup path), so callers that close inputs early still keep per-stream state/timers and can continue receiving/discarding server data until timeout. This can accumulate avoidable memory/network pressure when many read-only streams are abandoned early.
Useful? React with 👍 / 👎.
1 participant
This pull request introduces support for read-only data streaming in the Ratis client and adds a new
DataStreamReplyByteBufimplementation for handling replies using Netty'sByteBuf. It also refactors and extends several interfaces and utility classes to support asynchronous read operations and resource management for streaming replies. The most important changes are summarized below.Read-only DataStream support
DataStreamInputinterface, which provides asynchronous, zero-copy reading of stream replies and exposes a future for the finalRaftClientReply(ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamInput.java).DataStreamApiand its implementation to create read-only streams, includingstreamReadOnly()andstreamReadOnly(ByteBuffer message)(ratis-client/src/main/java/org/apache/ratis/client/api/DataStreamApi.java,ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java). [1] [2]DataStreamInputImplto manage the lifecycle, asynchronous reading, and resource cleanup of replies for read-only streams (ratis-client/src/main/java/org/apache/ratis/client/impl/DataStreamClientImpl.java).ByteBuf-based DataStreamReply implementation
DataStreamReplyByteBuf, a new implementation ofDataStreamReplyusing Netty'sByteBuf, with builder pattern, resource management, and commit info support (ratis-common/src/main/java/org/apache/ratis/datastream/impl/DataStreamReplyByteBuf.java).DataStreamReplyinterface to include arelease()method for resource management, with a default no-op implementation (ratis-common/src/main/java/org/apache/ratis/protocol/DataStreamReply.java).Protocol and utility enhancements
ClientProtoUtils.getRaftClientReply()to support extractingRaftClientReplyfrom bothDataStreamReplyByteBufferand the newDataStreamReplyByteBuf(ratis-client/src/main/java/org/apache/ratis/client/impl/ClientProtoUtils.java).DataStreamReplyByteBuf, including conversion to/fromDataStreamReplyByteBufferand header processing (ratis-netty/src/main/java/org/apache/ratis/netty/NettyDataStreamUtils.java).API changes
DataStreamClientRpcwith a new overloadedstreamAsyncmethod that accepts aConsumer<DataStreamReply>, enabling async receipt of multiple replies for a request (ratis-client/src/main/java/org/apache/ratis/client/DataStreamClientRpc.java).These changes collectively enhance the Ratis client with efficient, asynchronous, and resource-safe read-only streaming capabilities and better integration with Netty's buffer management.
What is the link to the Apache JIRA
https://issues.apache.org/jira/browse/RATIS-2546
How was this patch tested?
(Please explain how this patch was tested. Ex: unit tests, manual tests)
(If this patch involves UI changes, please attach a screen-shot; otherwise, remove this)