From 4bb3e969790721ba05120d13c39ecf0b46e4c20f Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Thu, 6 Aug 2026 15:29:02 -0400 Subject: [PATCH] Fix GRPC channel leak --- .../fn/data/BeamFnDataGrpcMultiplexer.java | 15 ++++++++++++++ .../fn/harness/ExternalWorkerService.java | 20 ++++++++++++------- .../org/apache/beam/fn/harness/FnHarness.java | 3 +++ .../fn/harness/data/BeamFnDataGrpcClient.java | 6 +++--- 4 files changed, 34 insertions(+), 10 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java index 0b9d6adab4f0..d8dd956ef5c1 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java @@ -28,6 +28,7 @@ import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.pipeline.v1.Endpoints; import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Status; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; @@ -57,6 +58,7 @@ public class BeamFnDataGrpcMultiplexer implements AutoCloseable { private final Endpoints.@Nullable ApiServiceDescriptor apiServiceDescriptor; private final StreamObserver inboundObserver; private final StreamObserver outboundObserver; + private final @Nullable ManagedChannel channel; private final ConcurrentHashMap< /*instructionId=*/ String, CompletableFuture>> receivers; @@ -73,7 +75,17 @@ public BeamFnDataGrpcMultiplexer( OutboundObserverFactory outboundObserverFactory, OutboundObserverFactory.BasicFactory baseOutboundObserverFactory) { + this(apiServiceDescriptor, outboundObserverFactory, baseOutboundObserverFactory, null); + } + + public BeamFnDataGrpcMultiplexer( + Endpoints.@Nullable ApiServiceDescriptor apiServiceDescriptor, + OutboundObserverFactory outboundObserverFactory, + OutboundObserverFactory.BasicFactory + baseOutboundObserverFactory, + @Nullable ManagedChannel channel) { this.apiServiceDescriptor = apiServiceDescriptor; + this.channel = channel; this.receivers = new ConcurrentHashMap<>(); this.poisonedInstructionIds = CacheBuilder.newBuilder().expireAfterWrite(POISONED_INSTRUCTION_ID_CACHE_TIMEOUT).build(); @@ -194,6 +206,9 @@ public void close() throws Exception { outboundObserver.onError( Status.CANCELLED.withDescription("Multiplexer hanging up").asException()); inboundObserver.onCompleted(); + if (channel != null) { + channel.shutdown(); + } if (exception != null) { throw exception; } 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 466d8ef3193b..5abb44c89b74 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.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,13 +79,18 @@ public void startWorker( .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(); + ManagedChannel channel = channelFactory.forDescriptor(request.getProvisionEndpoint()); + ProvisionApi.ProvisionInfo provisionInfo; + try { + ProvisionServiceGrpc.ProvisionServiceBlockingStub provisionStub = + ProvisionServiceGrpc.newBlockingStub(channel); + provisionInfo = + provisionStub + .getProvisionInfo(ProvisionApi.GetProvisionInfoRequest.newBuilder().build()) + .getInfo(); + } finally { + channel.shutdown(); + } runnerCapabilites = Sets.newHashSet(provisionInfo.getRunnerCapabilitiesList()); if (provisionInfo.hasControlEndpoint()) { 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 7bda0e18cad8..bff9ceabb52b 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 @@ -453,6 +453,9 @@ private BeamFnApi.ProcessBundleDescriptor loadDescriptor(String id) { beamFnStatusClient.close(); } processBundleHandler.shutdown(); + if (channel != null) { + 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/data/BeamFnDataGrpcClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java index 79e34fe5765f..9af20dd65320 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 @@ -142,10 +142,10 @@ private BeamFnDataGrpcMultiplexer getMultiplexer( return multiplexerCache.computeIfAbsent( key, k -> { + ManagedChannel channel = channelFactory.apply(apiServiceDescriptor); OutboundObserverFactory.BasicFactory 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( @@ -156,7 +156,7 @@ private BeamFnDataGrpcMultiplexer getMultiplexer( return stub.data(inboundObserver); }; return new BeamFnDataGrpcMultiplexer( - apiServiceDescriptor, outboundObserverFactory, baseOutboundObserverFactory); + apiServiceDescriptor, outboundObserverFactory, baseOutboundObserverFactory, channel); }); } }