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..b1bc462702df 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,20 +79,24 @@ 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(); - - 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 63f6c41b5a71..f547a1eeadff 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 @@ private BeamFnApi.ProcessBundleDescriptor loadDescriptor(String id) { 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 f12340552fd0..9fd1f0f0d68b 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 BeamFnApi.InstructionResponse.Builder trySplit(InstructionRequest request 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 79e34fe5765f..0640cc0d289d 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 int hashCode() { } private final ConcurrentMap multiplexerCache; + private final ConcurrentMap channelCache; private final Function channelFactory; private final OutboundObserverFactory outboundObserverFactory; @@ -85,6 +86,7 @@ public BeamFnDataGrpcClient( this.channelFactory = channelFactory; this.outboundObserverFactory = outboundObserverFactory; this.multiplexerCache = new ConcurrentHashMap<>(); + this.channelCache = new ConcurrentHashMap<>(); } @Override @@ -128,6 +130,10 @@ public void close() { } } multiplexerCache.clear(); + for (ManagedChannel channel : channelCache.values()) { + channel.shutdown(); + } + channelCache.clear(); } @Override @@ -142,10 +148,11 @@ private BeamFnDataGrpcMultiplexer getMultiplexer( return multiplexerCache.computeIfAbsent( key, k -> { + ManagedChannel channel = + channelCache.computeIfAbsent(k, unused -> 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( 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 c9adb9d7d142..1a0c4f3fb912 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 @@ public class BeamFnStateGrpcClientCache { private static final Logger LOG = LoggerFactory.getLogger(BeamFnStateGrpcClientCache.class); - private final Map cache; + private final Map cache; private final ManagedChannelFactory channelFactory; private final OutboundObserverFactory outboundObserverFactory; private final IdGenerator idGenerator; @@ -71,7 +73,7 @@ public synchronized BeamFnStateClient forApiServiceDescriptor( // 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 synchronized BeamFnStateClient forApiServiceDescriptor( return rval; } + /** Shuts down the underlying channels of all cached state clients. */ + public void close() { + List 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 @@ private GrpcStateClient(ApiServiceDescriptor apiServiceDescriptor) { synchronized (lock) { if (errorDuringConstruction) { outboundObserver.onCompleted(); + channel.shutdown(); } } } @@ -160,6 +175,7 @@ private void closeAndCleanUp(RuntimeException cause) { errorDuringConstruction = true; } else { outboundObserver.onCompleted(); + channel.shutdown(); } } }