[FnApi Java] Close named data stream multiplexers once no bundle is u…#39255
[FnApi Java] Close named data stream multiplexers once no bundle is u…#39255krishnamanchikalapudi wants to merge 2 commits into
Conversation
…sing them Named data streams introduced in apache#38863 provide bundle isolation, but the SDK harness cached one multiplexer (and its underlying gRPC stream) per (endpoint, dataStreamId) forever. A runner assigning fresh data stream ids over time would leak gRPC streams and memory in the harness. This reference counts named data stream usage per bundle in ProcessBundleHandler and closes the multiplexers for a named data stream once no bundle is retaining it. The default (unnamed) data stream keeps its lifetime-of-the-harness behavior. Also adds a fork-friendly GitHub-hosted CI workflow that builds and tests the Java SDK harness since upstream PreCommit workflows require self-hosted runners. Fixes apache#39001 (Java SDK harness portion)
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a resource leak in the Java SDK harness where named data stream multiplexers were cached indefinitely. By introducing reference counting, the harness now correctly tracks bundle usage and closes multiplexers once they are no longer required, significantly improving memory efficiency and resource cleanup. The default unnamed data stream maintains its existing behavior of persisting for the lifetime of the harness. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Ignored Files
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![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 addresses a resource leak in the SDK harness by introducing reference counting for named data streams, ensuring they are closed once no longer in use by any bundle. The changes include adding retention and release hooks in ProcessBundleHandler and implementing the lifecycle management in BeamFnDataGrpcClient, supported by new unit tests. The review feedback suggests simplifying the multiplexer retrieval logic in BeamFnDataGrpcClient to avoid redundant computeIfAbsent calls inside a synchronized block, and addresses potential thread leaks in the new tests by ensuring that ManagedChannel instances are properly shut down in finally blocks.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| synchronized (dataStreamLifecycleLock) { | ||
| return multiplexerCache.computeIfAbsent( | ||
| key, | ||
| k -> { | ||
| OutboundObserverFactory.BasicFactory<Elements, Elements> baseOutboundObserverFactory = | ||
| inboundObserver -> { | ||
| BeamFnDataGrpc.BeamFnDataStub stub = | ||
| BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor)); | ||
| if (dataStreamId != null && !dataStreamId.isEmpty()) { | ||
| Metadata headers = new Metadata(); | ||
| headers.put( | ||
| Metadata.Key.of("data_stream_id", Metadata.ASCII_STRING_MARSHALLER), | ||
| dataStreamId); | ||
| stub = | ||
| stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers)); | ||
| } | ||
| return stub.data(inboundObserver); | ||
| }; | ||
| return new BeamFnDataGrpcMultiplexer( | ||
| apiServiceDescriptor, outboundObserverFactory, baseOutboundObserverFactory); | ||
| }); | ||
| } |
There was a problem hiding this comment.
Since we are already inside the synchronized (dataStreamLifecycleLock) block, using computeIfAbsent on multiplexerCache is redundant. We can simplify this by performing a simple get and put check inside the synchronized block, which also avoids allocating the lambda/closure on every cache miss.
synchronized (dataStreamLifecycleLock) {
BeamFnDataGrpcMultiplexer multiplexer = multiplexerCache.get(key);
if (multiplexer == null) {
OutboundObserverFactory.BasicFactory<Elements, Elements> baseOutboundObserverFactory =
inboundObserver -> {
BeamFnDataGrpc.BeamFnDataStub stub =
BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor));
if (dataStreamId != null && !dataStreamId.isEmpty()) {
Metadata headers = new Metadata();
headers.put(
Metadata.Key.of("data_stream_id", Metadata.ASCII_STRING_MARSHALLER),
dataStreamId);
stub =
stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers));
}
return stub.data(inboundObserver);
};
multiplexer = new BeamFnDataGrpcMultiplexer(
apiServiceDescriptor, outboundObserverFactory, baseOutboundObserverFactory);
multiplexerCache.put(key, multiplexer);
}
return multiplexer;
}| try { | ||
| ManagedChannel channel = | ||
| InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); | ||
|
|
||
| BeamFnDataGrpcClient clientFactory = | ||
| new BeamFnDataGrpcClient( | ||
| (Endpoints.ApiServiceDescriptor descriptor) -> channel, | ||
| OutboundObserverFactory.trivial()); | ||
|
|
||
| // Two bundles concurrently use the same named data stream. | ||
| clientFactory.retainDataStream(dataStreamId); | ||
| clientFactory.retainDataStream(dataStreamId); | ||
| clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId); | ||
| assertTrue(firstConnection.await(5, TimeUnit.SECONDS)); | ||
| assertEquals(1, connectionCount.get()); | ||
|
|
||
| // Releasing one of the two usages should not close the stream. | ||
| clientFactory.releaseDataStream(dataStreamId); | ||
| assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS)); | ||
|
|
||
| // Releasing the last usage should close the stream. | ||
| clientFactory.releaseDataStream(dataStreamId); | ||
| assertTrue(streamTerminated.await(5, TimeUnit.SECONDS)); | ||
|
|
||
| // A subsequent usage of the same named data stream establishes a new stream. | ||
| clientFactory.retainDataStream(dataStreamId); | ||
| clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId); | ||
| assertTrue(secondConnection.await(5, TimeUnit.SECONDS)); | ||
| assertEquals(2, connectionCount.get()); | ||
| clientFactory.releaseDataStream(dataStreamId); | ||
| } finally { | ||
| server.shutdownNow(); | ||
| } |
There was a problem hiding this comment.
The ManagedChannel created for the test is not shut down in the finally block, which can leak threads and other resources in the test JVM. It should be shut down along with the server.
ManagedChannel channel =
InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
try {
BeamFnDataGrpcClient clientFactory =
new BeamFnDataGrpcClient(
(Endpoints.ApiServiceDescriptor descriptor) -> channel,
OutboundObserverFactory.trivial());
// Two bundles concurrently use the same named data stream.
clientFactory.retainDataStream(dataStreamId);
clientFactory.retainDataStream(dataStreamId);
clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId);
assertTrue(firstConnection.await(5, TimeUnit.SECONDS));
assertEquals(1, connectionCount.get());
// Releasing one of the two usages should not close the stream.
clientFactory.releaseDataStream(dataStreamId);
assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS));
// Releasing the last usage should close the stream.
clientFactory.releaseDataStream(dataStreamId);
assertTrue(streamTerminated.await(5, TimeUnit.SECONDS));
// A subsequent usage of the same named data stream establishes a new stream.
clientFactory.retainDataStream(dataStreamId);
clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId);
assertTrue(secondConnection.await(5, TimeUnit.SECONDS));
assertEquals(2, connectionCount.get());
clientFactory.releaseDataStream(dataStreamId);
} finally {
channel.shutdownNow();
server.shutdownNow();
}| try { | ||
| ManagedChannel channel = | ||
| InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); | ||
|
|
||
| BeamFnDataGrpcClient clientFactory = | ||
| new BeamFnDataGrpcClient( | ||
| (Endpoints.ApiServiceDescriptor descriptor) -> channel, | ||
| OutboundObserverFactory.trivial()); | ||
|
|
||
| clientFactory.retainDataStream(""); | ||
| StreamObserver<BeamFnApi.Elements> outboundObserver = | ||
| clientFactory.getOutboundObserver(apiServiceDescriptor, ""); | ||
| clientFactory.releaseDataStream(""); | ||
|
|
||
| // The default data stream is kept open and reused even after release. | ||
| assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS)); | ||
| assertSame(outboundObserver, clientFactory.getOutboundObserver(apiServiceDescriptor, "")); | ||
| } finally { | ||
| server.shutdownNow(); | ||
| } |
There was a problem hiding this comment.
The ManagedChannel created for the test is not shut down in the finally block, which can leak threads and other resources in the test JVM. It should be shut down along with the server.
| try { | |
| ManagedChannel channel = | |
| InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); | |
| BeamFnDataGrpcClient clientFactory = | |
| new BeamFnDataGrpcClient( | |
| (Endpoints.ApiServiceDescriptor descriptor) -> channel, | |
| OutboundObserverFactory.trivial()); | |
| clientFactory.retainDataStream(""); | |
| StreamObserver<BeamFnApi.Elements> outboundObserver = | |
| clientFactory.getOutboundObserver(apiServiceDescriptor, ""); | |
| clientFactory.releaseDataStream(""); | |
| // The default data stream is kept open and reused even after release. | |
| assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS)); | |
| assertSame(outboundObserver, clientFactory.getOutboundObserver(apiServiceDescriptor, "")); | |
| } finally { | |
| server.shutdownNow(); | |
| } | |
| ManagedChannel channel = | |
| InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); | |
| try { | |
| BeamFnDataGrpcClient clientFactory = | |
| new BeamFnDataGrpcClient( | |
| (Endpoints.ApiServiceDescriptor descriptor) -> channel, | |
| OutboundObserverFactory.trivial()); | |
| clientFactory.retainDataStream(""); | |
| StreamObserver<BeamFnApi.Elements> outboundObserver = | |
| clientFactory.getOutboundObserver(apiServiceDescriptor, ""); | |
| clientFactory.releaseDataStream(""); | |
| // The default data stream is kept open and reused even after release. | |
| assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS)); | |
| assertSame(outboundObserver, clientFactory.getOutboundObserver(apiServiceDescriptor, "")); | |
| } finally { | |
| channel.shutdownNow(); | |
| server.shutdownNow(); | |
| } |
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
1 participant
…sing them
Named data streams introduced in #38863 provide bundle isolation, but the SDK harness cached one multiplexer (and its underlying gRPC stream) per (endpoint, dataStreamId) forever. A runner assigning fresh data stream ids over time would leak gRPC streams and memory in the harness.
This reference counts named data stream usage per bundle in ProcessBundleHandler and closes the multiplexers for a named data stream once no bundle is retaining it. The default (unnamed) data stream keeps its lifetime-of-the-harness behavior.
Also adds a fork-friendly GitHub-hosted CI workflow that builds and tests the Java SDK harness since upstream PreCommit workflows require self-hosted runners.
Fixes #39001 (Java SDK harness portion)
Please add a meaningful description for your change here
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.