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

.github/workflows/beam_PreCommit_Java.yml

@@ -180,6 +180,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

@@ -82,6 +82,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)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Security Fixes

runners/flink/flink_runner.gradle

@@ -187,6 +187,9 @@ test {
   // Change log level to debug only for the package and nested packages:
   // systemProperty "org.slf4j.simpleLogger.log.org.apache.beam.runners.flink.translation.wrappers.streaming", "debug"
   jvmArgs "-XX:-UseGCOverheadLimit"
+  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,19 +255,58 @@ 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;
     this.executorService = executorService;
     this.credentials = credentials;
     this.maxBytesRewrittenPerCall = null;
     this.numRewriteTokensUsed = null;
-    googleCloudStorageOptions =
+    GoogleCloudStorageOptions.Builder storageOptionsBuilder =
         GoogleCloudStorageOptions.builder()
             .setAppName("Beam")
             .setReadChannelOptions(gcsReadOptions)
-            .setGrpcEnabled(shouldUseGrpc)
-            .build();
+            .setGrpcEnabled(shouldUseGrpc);
+    if (gcsEndpoint != null) {
+      try {
+        java.net.URL url = new java.net.URL(gcsEndpoint);
+        String rootUrl =
+            url.getProtocol()
+                + "://"
+                + url.getHost()
+                + (url.getPort() > 0 ? ":" + url.getPort() : "");
+        storageOptionsBuilder.setStorageRootUrl(rootUrl);
+        storageOptionsBuilder.setStorageServicePath(url.getPath());
+      } catch (java.net.MalformedURLException e) {
+        throw new RuntimeException("Invalid URL: " + gcsEndpoint, e);
+      }
+    }
+    googleCloudStorageOptions = storageOptionsBuilder.build();
     try {
       googleCloudStorage =
           createGoogleCloudStorage(googleCloudStorageOptions, storageClient, credentials);
@@ -505,6 +545,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

@@ -1703,7 +1703,7 @@ public static GcsUtilV1Mock createMock(PipelineOptions options) {
               gcsOptions.getEnableBucketWriteMetricCounter()
                   ? gcsOptions.getGcsWriteCounterPrefix()
                   : null),
-          gcsOptions.getGoogleCloudStorageReadOptions());
+          gcsOptions);
     }
 
     private GcsUtilV1Mock(
@@ -1715,7 +1715,7 @@ private GcsUtilV1Mock(
         @Nullable Integer uploadBufferSizeBytes,
         @Nullable Integer rewriteDataOpBatchLimit,
         GcsUtilV1.GcsCountersOptions gcsCountersOptions,
-        GoogleCloudStorageReadOptions gcsReadOptions) {
+        GcsOptions gcsOptions) {
       super(
           storageClient,
           httpRequestInitializer,
@@ -1725,7 +1725,7 @@ private GcsUtilV1Mock(
           uploadBufferSizeBytes,
           rewriteDataOpBatchLimit,
           gcsCountersOptions,
-          gcsReadOptions);
+          gcsOptions);
     }
 
     @Override
@@ -1866,6 +1866,19 @@ public void testReadMetricsAreNotCollectedWhenNotEnabledOpenWithOptions() throws
     testReadMetrics(false, GoogleCloudStorageReadOptions.DEFAULT);
   }
 
+  @Test
+  public void testGcsEndpointPropagation() {
+    GcsOptions options = PipelineOptionsFactory.as(GcsOptions.class);
+    options.setGcpCredential(new TestCredential());
+    options.setGcsEndpoint("http://localhost:8080/storage/v1/");
+
+    GcsUtilV1 gcsUtilV1 = new GcsUtilV1.GcsUtilFactory().create(options);
+    GoogleCloudStorageOptions googleCloudStorageOptions = gcsUtilV1.getGoogleCloudStorageOptions();
+
+    assertEquals("http://localhost:8080", googleCloudStorageOptions.getStorageRootUrl());
+    assertEquals("/storage/v1/", googleCloudStorageOptions.getStorageServicePath());
+  }
+
   /** A helper to wrap a {@link GenericJson} object in a content stream. */
   private static InputStream toStream(String content) throws IOException {
     return new ByteArrayInputStream(content.getBytes(StandardCharsets.UTF_8));

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

@@ -72,4 +72,23 @@ void unregisterReceiver(
   /** Get the outbound observer for the specified apiServiceDescriptor and dataStreamId. */
   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) {}
 }