Skip to content

Fix gRPC stream observer leak on ProcessBundleHandler shutdown #39390

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Abacn merged 2 commits into from
Jul 21, 2026
+21 −1
Merged

Fix gRPC stream observer leak on ProcessBundleHandler shutdown #39390

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cc7c24f
Fix gRPC stream observer leak on ProcessBundleHandler shutdown
Abacn Jul 20, 2026
eacf9b9
Update sdks/java/harness/src/main/java/org/apache/beam/fn/harness/dat…
Abacn Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -767,6 +767,7 @@ public BeamFnApi.InstructionResponse.Builder trySplit(InstructionRequest request
/** Shutdown the bundles, running the tearDown() functions. */
public void shutdown() throws Exception {
bundleProcessorCache.shutdown();
beamFnDataClient.close();

@Abacn Abacn Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

BeamFnDataGrpcClient is created in FnHarness.main:

beam/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java

Line 335 in 39c60f0

BeamFnDataGrpcClient beamFnDataMultiplexer =

and ownership passed to ProcessBundleHandler:

beam/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java

Line 371 in 39c60f0

beamFnDataMultiplexer,

Therefore we should close it here.

When FnHarness.main is invoked multiple times (instead of as a standalone worker like on Dataflow runner), previously unclosed BeamFnDataGrpcClient leak gRPC streams

}
Comment on lines 768 to 771

@gemini-code-assist gemini-code-assist Bot Jul 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If bundleProcessorCache.shutdown() throws an exception, beamFnDataClient.close() will not be called, potentially leaking resources. Wrapping the shutdown calls in a try-finally block ensures that beamFnDataClient.close() is always executed.

  public void shutdown() throws Exception {
    try {
      bundleProcessorCache.shutdown();
    } finally {
      beamFnDataClient.close();
    }
  }

@Abacn Abacn Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this may be overly defensive. close is always best effort. Currently bundleProcessorCache.shutdown() just invalidates a cache container


@VisibleForTesting
Expand Down
9 changes: 8 additions & 1 deletion sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
*/
package org.apache.beam.fn.harness.data;

import java.io.Closeable;
import java.io.IOException;
import java.util.List;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements;
import org.apache.beam.model.pipeline.v1.Endpoints;
Expand All @@ -30,7 +32,7 @@
* provide a receiver of outbound elements. Callers can register themselves as receivers for inbound
* elements or can get a handle for a receiver of outbound elements.
*/
public interface BeamFnDataClient {
public interface BeamFnDataClient extends Closeable {
/**
* Registers a receiver for the provided instruction id.
*
Expand Down Expand Up @@ -72,4 +74,9 @@ void unregisterReceiver(
/** Get the outbound observer for the specified apiServiceDescriptor and dataStreamId. */
StreamObserver<Elements> getOutboundObserver(
Endpoints.ApiServiceDescriptor apiServiceDescriptor, String dataStreamId);

@Override
default void close() throws IOException {
// Default to no-op
}
}
12 changes: 12 additions & 0 deletions sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,18 @@ public void poisonInstructionId(String instructionId) {
}
}

@Override
public void close() {
for (BeamFnDataGrpcMultiplexer client : multiplexerCache.values()) {
try {
client.close();
} catch (Exception e) {
LOG.warn("Failed to close multiplexer", e);
}
}
Comment thread
Abacn marked this conversation as resolved.
multiplexerCache.clear();
}

@Override
public StreamObserver<Elements> getOutboundObserver(
ApiServiceDescriptor apiServiceDescriptor, String dataStreamId) {
Expand Down
Loading