This is an automated email from the ASF dual-hosted git repository.
shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 0072521920c [Java] Fix GRPC channel leak (#39663)
0072521920c is described below
commit 0072521920c2d4219f3191e2ffd748ee4b84babb
Author: Shunping Huang <[email protected]>
AuthorDate: Fri Oct 16 08:41:32 2026 -0400
[Java] Fix GRPC channel leak (#39663)
---
.../beam/fn/harness/ExternalWorkerService.java | 33 +++++++++++++---------
.../java/org/apache/beam/fn/harness/FnHarness.java | 1 +
.../fn/harness/control/ProcessBundleHandler.java | 3 ++
.../beam/fn/harness/data/BeamFnDataGrpcClient.java | 11 ++++++--
.../harness/state/BeamFnStateGrpcClientCache.java | 20 +++++++++++--
5 files changed, 50 insertions(+), 18 deletions(-)
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ExternalWorkerService.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ExternalWorkerService.java
index 466d8ef3193..b1bc462702d 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ExternalWorkerService.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ExternalWorkerService.java
@@ -39,6 +39,7 @@ import org.apache.beam.sdk.options.PortablePipelineOptions;
import org.apache.beam.sdk.util.Sleeper;
import org.apache.beam.sdk.util.construction.Environments;
import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
+import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
@@ -78,20 +79,24 @@ public class ExternalWorkerService extends
BeamFnExternalWorkerPoolImplBase impl
.withInterceptors(
ImmutableList.of(AddHarnessIdInterceptor.create(request.getWorkerId())));
- ProvisionServiceGrpc.ProvisionServiceBlockingStub provisionStub =
- ProvisionServiceGrpc.newBlockingStub(
- channelFactory.forDescriptor(request.getProvisionEndpoint()));
- ProvisionApi.ProvisionInfo provisionInfo =
- provisionStub
-
.getProvisionInfo(ProvisionApi.GetProvisionInfoRequest.newBuilder().build())
- .getInfo();
-
- runnerCapabilites =
Sets.newHashSet(provisionInfo.getRunnerCapabilitiesList());
- if (provisionInfo.hasControlEndpoint()) {
- controlEndpoint = provisionInfo.getControlEndpoint();
- }
- if (provisionInfo.hasLoggingEndpoint()) {
- loggingEndpoint = provisionInfo.getLoggingEndpoint();
+ ManagedChannel channel =
channelFactory.forDescriptor(request.getProvisionEndpoint());
+ try {
+ ProvisionServiceGrpc.ProvisionServiceBlockingStub provisionStub =
+ ProvisionServiceGrpc.newBlockingStub(channel);
+ ProvisionApi.ProvisionInfo provisionInfo =
+ provisionStub
+
.getProvisionInfo(ProvisionApi.GetProvisionInfoRequest.newBuilder().build())
+ .getInfo();
+
+ runnerCapabilites =
Sets.newHashSet(provisionInfo.getRunnerCapabilitiesList());
+ if (provisionInfo.hasControlEndpoint()) {
+ controlEndpoint = provisionInfo.getControlEndpoint();
+ }
+ if (provisionInfo.hasLoggingEndpoint()) {
+ loggingEndpoint = provisionInfo.getLoggingEndpoint();
+ }
+ } finally {
+ channel.shutdown();
}
}
// Lambda closured variables must be final.
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
index 63f6c41b5a7..f547a1eeadf 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
@@ -437,6 +437,7 @@ public class FnHarness {
beamFnStatusClient.close();
}
processBundleHandler.shutdown();
+ channel.shutdown();
} catch (Exception e) {
LOG.error("Shutting down harness due to exception", e);
e.printStackTrace();
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
index f12340552fd..9fd1f0f0d68 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
@@ -768,6 +768,9 @@ public class ProcessBundleHandler {
public void shutdown() throws Exception {
bundleProcessorCache.shutdown();
beamFnDataClient.close();
+ if (beamFnStateGrpcClientCache != null) {
+ beamFnStateGrpcClientCache.close();
+ }
}
@VisibleForTesting
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
index 79e34fe5765..0640cc0d289 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
@@ -76,6 +76,7 @@ public class BeamFnDataGrpcClient implements BeamFnDataClient
{
}
private final ConcurrentMap<MultiplexerKey, BeamFnDataGrpcMultiplexer>
multiplexerCache;
+ private final ConcurrentMap<MultiplexerKey, ManagedChannel> channelCache;
private final Function<Endpoints.ApiServiceDescriptor, ManagedChannel>
channelFactory;
private final OutboundObserverFactory outboundObserverFactory;
@@ -85,6 +86,7 @@ public class BeamFnDataGrpcClient implements BeamFnDataClient
{
this.channelFactory = channelFactory;
this.outboundObserverFactory = outboundObserverFactory;
this.multiplexerCache = new ConcurrentHashMap<>();
+ this.channelCache = new ConcurrentHashMap<>();
}
@Override
@@ -128,6 +130,10 @@ public class BeamFnDataGrpcClient implements
BeamFnDataClient {
}
}
multiplexerCache.clear();
+ for (ManagedChannel channel : channelCache.values()) {
+ channel.shutdown();
+ }
+ channelCache.clear();
}
@Override
@@ -142,10 +148,11 @@ public class BeamFnDataGrpcClient implements
BeamFnDataClient {
return multiplexerCache.computeIfAbsent(
key,
k -> {
+ ManagedChannel channel =
+ channelCache.computeIfAbsent(k, unused ->
channelFactory.apply(apiServiceDescriptor));
OutboundObserverFactory.BasicFactory<Elements, Elements>
baseOutboundObserverFactory =
inboundObserver -> {
- BeamFnDataGrpc.BeamFnDataStub stub =
-
BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor));
+ BeamFnDataGrpc.BeamFnDataStub stub =
BeamFnDataGrpc.newStub(channel);
if (dataStreamId != null && !dataStreamId.isEmpty()) {
Metadata headers = new Metadata();
headers.put(
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
index c9adb9d7d14..1a0c4f3fb91 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
@@ -18,7 +18,9 @@
package org.apache.beam.fn.harness.state;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
@@ -44,7 +46,7 @@ import org.slf4j.LoggerFactory;
public class BeamFnStateGrpcClientCache {
private static final Logger LOG =
LoggerFactory.getLogger(BeamFnStateGrpcClientCache.class);
- private final Map<ApiServiceDescriptor, BeamFnStateClient> cache;
+ private final Map<ApiServiceDescriptor, GrpcStateClient> cache;
private final ManagedChannelFactory channelFactory;
private final OutboundObserverFactory outboundObserverFactory;
private final IdGenerator idGenerator;
@@ -71,7 +73,7 @@ public class BeamFnStateGrpcClientCache {
// We specifically are synchronized so that we only create one
GrpcStateClient at a time
// preventing a race where multiple GrpcStateClient objects might be
constructed at the same
// for the same ApiServiceDescriptor.
- BeamFnStateClient rval;
+ GrpcStateClient rval;
synchronized (cache) {
rval = cache.get(apiServiceDescriptor);
}
@@ -87,6 +89,18 @@ public class BeamFnStateGrpcClientCache {
return rval;
}
+ /** Shuts down the underlying channels of all cached state clients. */
+ public void close() {
+ List<GrpcStateClient> clients;
+ synchronized (cache) {
+ clients = new ArrayList<>(cache.values());
+ cache.clear();
+ }
+ for (GrpcStateClient client : clients) {
+ client.channel.shutdown();
+ }
+ }
+
/** A {@link BeamFnStateClient} for a given {@link ApiServiceDescriptor}. */
private class GrpcStateClient implements BeamFnStateClient {
private final Object lock = new Object();
@@ -111,6 +125,7 @@ public class BeamFnStateGrpcClientCache {
synchronized (lock) {
if (errorDuringConstruction) {
outboundObserver.onCompleted();
+ channel.shutdown();
}
}
}
@@ -160,6 +175,7 @@ public class BeamFnStateGrpcClientCache {
errorDuringConstruction = true;
} else {
outboundObserver.onCompleted();
+ channel.shutdown();
}
}
}
