[FnApi Java] Close named data stream multiplexers once no bundle is u…

.github/workflows/beam_PreCommit_Java.yml

@@ -182,6 +182,7 @@ jobs:
       - name: Setup environment
         uses: ./.github/actions/setup-environment-action
         with:
+          java-version: 11
           python-version: default
           disable-cache: true
       - name: run Java PreCommit script

.github/workflows/fork_ci_java_harness.yml

@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# Fork-friendly CI for the Java SDK harness. The upstream Beam PreCommit
+# workflows run on self-hosted runners which are unavailable on forks, so
+# this workflow provides build, test, and style verification for changes to
+# the Java SDK harness using GitHub-hosted runners.
+name: Fork CI Java Harness
+
+on:
+  pull_request:
+    branches: ['master', 'release-*']
+    paths:
+      - 'model/**'
+      - 'sdks/java/core/**'
+      - 'sdks/java/harness/**'
+      - '.github/workflows/fork_ci_java_harness.yml'
+  workflow_dispatch:
+
+permissions: read-all
+
+concurrency:
+  group: '${{ github.workflow }} @ ${{ github.head_ref || github.ref }}'
+  cancel-in-progress: true
+
+jobs:
+  java_sdk_harness_tests:
+    name: Java SDK harness build, tests and style checks
+    runs-on: ubuntu-latest
+    timeout-minutes: 120
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v7
+      - name: Setup environment
+        uses: ./.github/actions/setup-environment-action
+        with:
+          java-version: default
+      - name: Run Java SDK harness tests
+        run: |
+          ./gradlew :sdks:java:harness:test \
+            -PdisableSpotlessApply \
+            --max-workers=2 \
+            --continue
+      - name: Run Java SDK harness style checks
+        run: |
+          ./gradlew :sdks:java:harness:spotlessCheck \
+            :sdks:java:harness:checkstyleMain \
+            :sdks:java:harness:checkstyleTest \
+            -PdisableSpotlessApply \
+            --max-workers=2 \
+            --continue
+      - name: Upload test reports
+        if: always()
+        uses: actions/upload-artifact@v7
+        with:
+          name: java-harness-test-reports
+          path: sdks/java/harness/build/reports
+          if-no-files-found: ignore

.test-infra/validate-runner/build.gradle

@@ -24,7 +24,7 @@ description = "Apache Beam :: Validate :: Runner"
 repositories {
     mavenCentral()
     maven {
-        url "https://repo.jenkins-ci.org/releases/"
+        url "https://repo.jenkins-ci.org/public/"
     }
     maven {
         url "https://packages.confluent.io/maven/"

CHANGES.md

@@ -91,6 +91,7 @@
 ## Bugfixes
 
 * Fixed unbounded checkpoint state growth for splittable DoFns that self-checkpoint on the portable Flink runner (Java) ([#27648](https://github.com/apache/beam/issues/27648)).
+* Fixed named data stream multiplexers and their underlying gRPC streams leaking in the SDK harness by closing them once no bundle is using them (Java) ([#39001](https://github.com/apache/beam/issues/39001)).
 * Improved Java pipeline performance by avoiding repeated `DoFn` type descriptor resolution when creating cached invokers ([#39309](https://github.com/apache/beam/issues/39309)).
 * (Python) Fixed a memory leak in Python SDK caused by storing exceptions with potentially large stack frames in a cache ([#39406](https://github.com/apache/beam/issues/39406)).
 

runners/flink/flink_runner.gradle

@@ -204,6 +204,9 @@ test {
   // systemProperty "org.slf4j.simpleLogger.log.org.apache.beam.runners.flink.translation.wrappers.streaming", "debug"
   jvmArgs "-XX:-UseGCOverheadLimit"
   jvmArgs += flinkTestJvmArgs()
+  if (JavaVersion.current().isJava9Compatible()) {
+    jvmArgs "--add-opens=java.base/java.util=ALL-UNNAMED"
+  }
   if (System.getProperty("beamSurefireArgline")) {
     jvmArgs System.getProperty("beamSurefireArgline")
   }

sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java

@@ -154,7 +154,7 @@ public GcsUtilV1 create(PipelineOptions options) {
               gcsOptions.getEnableBucketWriteMetricCounter()
                   ? gcsOptions.getGcsWriteCounterPrefix()
                   : null),
-          gcsOptions.getGoogleCloudStorageReadOptions());
+          gcsOptions);
     }
   }
 
@@ -240,7 +240,8 @@ public boolean shouldRetry(IOException e) {
         uploadBufferSizeBytes,
         rewriteDataOpBatchLimit,
         gcsCountersOptions,
-        gcsOptions.getGoogleCloudStorageReadOptions());
+        gcsOptions.getGoogleCloudStorageReadOptions(),
+        gcsOptions.getGcsEndpoint());
   }
 
   @VisibleForTesting
@@ -254,6 +255,31 @@ public boolean shouldRetry(IOException e) {
       @Nullable Integer rewriteDataOpBatchLimit,
       GcsCountersOptions gcsCountersOptions,
       GoogleCloudStorageReadOptions gcsReadOptions) {
+    this(
+        storageClient,
+        httpRequestInitializer,
+        executorService,
+        shouldUseGrpc,
+        credentials,
+        uploadBufferSizeBytes,
+        rewriteDataOpBatchLimit,
+        gcsCountersOptions,
+        gcsReadOptions,
+        null);
+  }
+
+  @VisibleForTesting
+  GcsUtilV1(
+      Storage storageClient,
+      HttpRequestInitializer httpRequestInitializer,
+      ExecutorService executorService,
+      Boolean shouldUseGrpc,
+      Credentials credentials,
+      @Nullable Integer uploadBufferSizeBytes,
+      @Nullable Integer rewriteDataOpBatchLimit,
+      GcsCountersOptions gcsCountersOptions,
+      GoogleCloudStorageReadOptions gcsReadOptions,
+      @Nullable String gcsEndpoint) {
     this.storageClient = storageClient;
     this.httpRequestInitializer = httpRequestInitializer;
     this.uploadBufferSizeBytes = uploadBufferSizeBytes;
@@ -516,6 +542,11 @@ void setCloudStorageImpl(GoogleCloudStorageOptions g) {
     googleCloudStorageOptions = g;
   }
 
+  @VisibleForTesting
+  GoogleCloudStorageOptions getGoogleCloudStorageOptions() {
+    return googleCloudStorageOptions;
+  }
+
   /**
    * Create an integer consumer that updates the counter identified by a prefix and a bucket name.
    */

sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java

@@ -1704,7 +1704,7 @@ public static GcsUtilV1Mock createMock(PipelineOptions options) {
               gcsOptions.getEnableBucketWriteMetricCounter()
                   ? gcsOptions.getGcsWriteCounterPrefix()
                   : null),
-          gcsOptions.getGoogleCloudStorageReadOptions());
+          gcsOptions);
     }
 
     private GcsUtilV1Mock(
@@ -1716,7 +1716,7 @@ private GcsUtilV1Mock(
         @Nullable Integer uploadBufferSizeBytes,
         @Nullable Integer rewriteDataOpBatchLimit,
         GcsUtilV1.GcsCountersOptions gcsCountersOptions,
-        GoogleCloudStorageReadOptions gcsReadOptions) {
+        GcsOptions gcsOptions) {
       super(
           storageClient,
           httpRequestInitializer,
@@ -1726,7 +1726,7 @@ private GcsUtilV1Mock(
           uploadBufferSizeBytes,
           rewriteDataOpBatchLimit,
           gcsCountersOptions,
-          gcsReadOptions);
+          gcsOptions);
     }
 
     @Override

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java

@@ -497,6 +497,11 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re
     String instructionId = request.getInstructionId();
     String dataStreamId = request.getProcessBundle().getDataStreamId();
     @Nullable BundleProcessor bundleProcessor = null;
+    if (!dataStreamId.isEmpty()) {
+      // Keep the named data stream open for the duration of the bundle. Once a named data stream
+      // is no longer being used by any bundle it may be closed, freeing its underlying resources.
+      beamFnDataClient.retainDataStream(dataStreamId);
+    }
     try {
       bundleProcessor =
           Preconditions.checkNotNull(
@@ -615,6 +620,10 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re
       // Ensure that if more data arrives for the instruction it is discarded.
       beamFnDataClient.poisonInstructionId(instructionId);
       throw e;
+    } finally {
+      if (!dataStreamId.isEmpty()) {
+        beamFnDataClient.releaseDataStream(dataStreamId);
+      }
     }
   }
 

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

@@ -75,6 +75,25 @@ void unregisterReceiver(
   StreamObserver<Elements> getOutboundObserver(
       Endpoints.ApiServiceDescriptor apiServiceDescriptor, String dataStreamId);
 
+  /**
+   * Indicates that an instruction is going to use the specified named data stream.
+   *
+   * <p>Callers must {@link #releaseDataStream release} the data stream once the instruction has
+   * finished using it. Named data streams are kept open for as long as they are retained by at
+   * least one instruction. The default (empty) data stream is kept open for the lifetime of the
+   * client and calls for it are a no-op.
+   */
+  default void retainDataStream(String dataStreamId) {}
+
+  /**
+   * Indicates that an instruction is done using the specified named data stream.
+   *
+   * <p>Once a named data stream is no longer retained by any instruction its underlying resources
+   * may be released. Subsequent usage of the same data stream id will establish a new physical
+   * stream. The default (empty) data stream is never closed and calls for it are a no-op.
+   */
+  default void releaseDataStream(String dataStreamId) {}
+
   @Override
   default void close() throws IOException {
     // Default to no-op

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

@@ -17,7 +17,11 @@
  */
 package org.apache.beam.fn.harness.data;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -79,12 +83,27 @@ public int hashCode() {
   private final Function<Endpoints.ApiServiceDescriptor, ManagedChannel> channelFactory;
   private final OutboundObserverFactory outboundObserverFactory;
 
+  /**
+   * Guards creation and removal of entries in {@link #multiplexerCache} as well as all accesses to
+   * {@link #dataStreamRefCounts} so that a named data stream is not concurrently created and
+   * closed.
+   */
+  private final Object dataStreamLifecycleLock = new Object();
+
+  /**
+   * The number of instructions currently retaining each named data stream. Guarded by {@link
+   * #dataStreamLifecycleLock}. The default (empty) data stream is not tracked as it is kept open
+   * for the lifetime of the client.
+   */
+  private final Map<String, Integer> dataStreamRefCounts;
+
   public BeamFnDataGrpcClient(
       Function<Endpoints.ApiServiceDescriptor, ManagedChannel> channelFactory,
       OutboundObserverFactory outboundObserverFactory) {
     this.channelFactory = channelFactory;
     this.outboundObserverFactory = outboundObserverFactory;
     this.multiplexerCache = new ConcurrentHashMap<>();
+    this.dataStreamRefCounts = new HashMap<>();
   }
 
   @Override
@@ -136,27 +155,88 @@ public StreamObserver<Elements> getOutboundObserver(
     return getMultiplexer(apiServiceDescriptor, dataStreamId).getOutboundObserver();
   }
 
+  @Override
+  public void retainDataStream(String dataStreamId) {
+    if (dataStreamId == null || dataStreamId.isEmpty()) {
+      // The default data stream is kept open for the lifetime of the client.
+      return;
+    }
+    synchronized (dataStreamLifecycleLock) {
+      dataStreamRefCounts.merge(dataStreamId, 1, Integer::sum);
+    }
+  }
+
+  @Override
+  public void releaseDataStream(String dataStreamId) {
+    if (dataStreamId == null || dataStreamId.isEmpty()) {
+      // The default data stream is kept open for the lifetime of the client.
+      return;
+    }
+    List<BeamFnDataGrpcMultiplexer> multiplexersToClose = new ArrayList<>();
+    synchronized (dataStreamLifecycleLock) {
+      Integer refCount = dataStreamRefCounts.get(dataStreamId);
+      if (refCount == null) {
+        LOG.warn("Released data stream {} which was not retained.", dataStreamId);
+        return;
+      }
+      if (refCount > 1) {
+        dataStreamRefCounts.put(dataStreamId, refCount - 1);
+        return;
+      }
+      dataStreamRefCounts.remove(dataStreamId);
+      Iterator<Map.Entry<MultiplexerKey, BeamFnDataGrpcMultiplexer>> iterator =
+          multiplexerCache.entrySet().iterator();
+      while (iterator.hasNext()) {
+        Map.Entry<MultiplexerKey, BeamFnDataGrpcMultiplexer> entry = iterator.next();
+        if (dataStreamId.equals(entry.getKey().dataStreamId)) {
+          multiplexersToClose.add(entry.getValue());
+          iterator.remove();
+        }
+      }
+    }
+    // Close outside of the lock as closing terminates the underlying gRPC stream and may block.
+    for (BeamFnDataGrpcMultiplexer multiplexer : multiplexersToClose) {
+      LOG.debug("Closing multiplexer for released data stream {}", dataStreamId);
+      try {
+        multiplexer.close();
+      } catch (Exception e) {
+        LOG.warn("Failed to close multiplexer for data stream {}", dataStreamId, e);
+      }
+    }
+  }
+
   private BeamFnDataGrpcMultiplexer getMultiplexer(
       Endpoints.ApiServiceDescriptor apiServiceDescriptor, String dataStreamId) {
     MultiplexerKey key = new MultiplexerKey(apiServiceDescriptor, dataStreamId);
-    return multiplexerCache.computeIfAbsent(
-        key,
-        k -> {
-          OutboundObserverFactory.BasicFactory<Elements, Elements> baseOutboundObserverFactory =
-              inboundObserver -> {
-                BeamFnDataGrpc.BeamFnDataStub stub =
-                    BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor));
-                if (dataStreamId != null && !dataStreamId.isEmpty()) {
-                  Metadata headers = new Metadata();
-                  headers.put(
-                      Metadata.Key.of("data_stream_id", Metadata.ASCII_STRING_MARSHALLER),
-                      dataStreamId);
-                  stub = stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers));
-                }
-                return stub.data(inboundObserver);
-              };
-          return new BeamFnDataGrpcMultiplexer(
-              apiServiceDescriptor, outboundObserverFactory, baseOutboundObserverFactory);
-        });
+    BeamFnDataGrpcMultiplexer existingMultiplexer = multiplexerCache.get(key);
+    if (existingMultiplexer != null) {
+      return existingMultiplexer;
+    }
+    // Create under the lifecycle lock so that a named data stream being concurrently closed by
+    // releaseDataStream is not observed in a partially removed state. Callers are expected to
+    // retain named data streams for the duration of their usage which prevents the returned
+    // multiplexer from being closed while in use.
+    synchronized (dataStreamLifecycleLock) {
+      return multiplexerCache.computeIfAbsent(
+          key,
+          k -> {
+            OutboundObserverFactory.BasicFactory<Elements, Elements> baseOutboundObserverFactory =
+                inboundObserver -> {
+                  BeamFnDataGrpc.BeamFnDataStub stub =
+                      BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor));
+                  if (dataStreamId != null && !dataStreamId.isEmpty()) {
+                    Metadata headers = new Metadata();
+                    headers.put(
+                        Metadata.Key.of("data_stream_id", Metadata.ASCII_STRING_MARSHALLER),
+                        dataStreamId);
+                    stub =
+                        stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers));
+                  }
+                  return stub.data(inboundObserver);
+                };
+            return new BeamFnDataGrpcMultiplexer(
+                apiServiceDescriptor, outboundObserverFactory, baseOutboundObserverFactory);
+          });
+    }
   }
 }