[Java] Fix GRPC channel leak

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.

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();

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

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java

@@ -76,6 +76,7 @@ public int hashCode() {
   }
 
   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 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<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(

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<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 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<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 @@ 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();
         }
       }
     }