From 00eb482640b2f52f2ba690845e123c9b42af61da Mon Sep 17 00:00:00 2001 From: Krishna Manchikalapudi Date: Wed, 8 Jul 2026 17:44:09 -0700 Subject: [PATCH 1/3] [FnApi Java] Close named data stream multiplexers once no bundle is using them Named data streams introduced in #38863 provide bundle isolation, but the SDK harness cached one multiplexer (and its underlying gRPC stream) per (endpoint, dataStreamId) forever. A runner assigning fresh data stream ids over time would leak gRPC streams and memory in the harness. This reference counts named data stream usage per bundle in ProcessBundleHandler and closes the multiplexers for a named data stream once no bundle is retaining it. The default (unnamed) data stream keeps its lifetime-of-the-harness behavior. Also adds a fork-friendly GitHub-hosted CI workflow that builds and tests the Java SDK harness since upstream PreCommit workflows require self-hosted runners. Fixes #39001 (Java SDK harness portion) --- .github/workflows/fork_ci_java_harness.yml | 72 ++++++++ CHANGES.md | 1 + .../harness/control/ProcessBundleHandler.java | 9 + .../fn/harness/data/BeamFnDataClient.java | 19 +++ .../fn/harness/data/BeamFnDataGrpcClient.java | 118 ++++++++++--- .../control/ProcessBundleHandlerTest.java | 155 ++++++++++++++++++ .../data/BeamFnDataGrpcClientTest.java | 112 +++++++++++++ 7 files changed, 467 insertions(+), 19 deletions(-) create mode 100644 .github/workflows/fork_ci_java_harness.yml diff --git a/.github/workflows/fork_ci_java_harness.yml b/.github/workflows/fork_ci_java_harness.yml new file mode 100644 index 000000000000..c0afa3c95cd1 --- /dev/null +++ b/.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 diff --git a/CHANGES.md b/CHANGES.md index 1e09ade6c389..80fc41e57d49 100644 --- a/CHANGES.md +++ b/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 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 449afd6a0243..3f3875c672f3 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 @@ -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); + } } } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java index 1a50f5b448c5..e115a7f46a92 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java +++ b/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 getOutboundObserver( Endpoints.ApiServiceDescriptor apiServiceDescriptor, String dataStreamId); + + /** + * Indicates that an instruction is going to use the specified named data stream. + * + *

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. + * + *

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) {} } 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 2f2a6b0fc660..4f4d68c7dd5b 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 @@ -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 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 dataStreamRefCounts; + public BeamFnDataGrpcClient( Function channelFactory, OutboundObserverFactory outboundObserverFactory) { this.channelFactory = channelFactory; this.outboundObserverFactory = outboundObserverFactory; this.multiplexerCache = new ConcurrentHashMap<>(); + this.dataStreamRefCounts = new HashMap<>(); } @Override @@ -124,27 +143,88 @@ public StreamObserver 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 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> iterator = + multiplexerCache.entrySet().iterator(); + while (iterator.hasNext()) { + Map.Entry 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 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 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); + }); + } } } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index c03f82726740..067adf08f256 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -160,6 +160,7 @@ import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.mockito.InOrder; import org.mockito.Mock; import org.mockito.Mockito; import org.mockito.MockitoAnnotations; @@ -1447,6 +1448,160 @@ public void testInstructionIsUnregisteredFromBeamFnDataClientOnSuccess() throws verifyNoMoreInteractions(beamFnDataClient); } + @Test + public void testNamedDataStreamIsRetainedForTheDurationOfTheBundle() throws Exception { + BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = + BeamFnApi.ProcessBundleDescriptor.newBuilder() + .putTransforms( + "2L", + RunnerApi.PTransform.newBuilder() + .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()) + .build()) + .build(); + Map fnApiRegistry = + ImmutableMap.of("1L", processBundleDescriptor); + + Mockito.doAnswer( + (invocation) -> { + String instructionId = invocation.getArgument(0, String.class); + CloseableFnDataReceiver data = + invocation.getArgument(3, CloseableFnDataReceiver.class); + data.accept( + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("2L") + .setIsLast(true)) + .build()); + return null; + }) + .when(beamFnDataClient) + .registerReceiver(any(), any(), any(), any()); + + ProcessBundleHandler handler = + new ProcessBundleHandler( + PipelineOptionsFactory.create(), + Collections.emptySet(), + fnApiRegistry::get, + beamFnDataClient, + null /* beamFnStateGrpcClientCache */, + null /* finalizeBundleHandler */, + new ShortIdMap(), + executionStateSampler, + ImmutableMap.of( + DATA_INPUT_URN, + (context) -> + context.addIncomingDataEndpoint( + ApiServiceDescriptor.getDefaultInstance(), + StringUtf8Coder.of(), + (input) -> {})), + Caches.noop(), + new BundleProcessorCache(Duration.ZERO), + null /* dataSampler */); + handler.processBundle( + BeamFnApi.InstructionRequest.newBuilder() + .setInstructionId("instructionId") + .setProcessBundle( + BeamFnApi.ProcessBundleRequest.newBuilder() + .setProcessBundleDescriptorId("1L") + .setDataStreamId("dataStreamId")) + .build()); + + // Ensure that the named data stream is retained for the duration of bundle processing and + // released once the bundle completes. + InOrder inOrder = Mockito.inOrder(beamFnDataClient); + inOrder.verify(beamFnDataClient).retainDataStream(eq("dataStreamId")); + inOrder + .verify(beamFnDataClient) + .registerReceiver(eq("instructionId"), eq("dataStreamId"), any(), any()); + inOrder + .verify(beamFnDataClient) + .unregisterReceiver(eq("instructionId"), eq("dataStreamId"), any()); + inOrder.verify(beamFnDataClient).releaseDataStream(eq("dataStreamId")); + verifyNoMoreInteractions(beamFnDataClient); + } + + @Test + public void testNamedDataStreamIsReleasedWhenBundleFails() throws Exception { + BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = + BeamFnApi.ProcessBundleDescriptor.newBuilder() + .putTransforms( + "2L", + RunnerApi.PTransform.newBuilder() + .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()) + .build()) + .build(); + Map fnApiRegistry = + ImmutableMap.of("1L", processBundleDescriptor); + + Mockito.doAnswer( + (invocation) -> { + ByteStringOutputStream encodedData = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("A", encodedData); + String instructionId = invocation.getArgument(0, String.class); + CloseableFnDataReceiver data = + invocation.getArgument(3, CloseableFnDataReceiver.class); + data.accept( + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("2L") + .setData(encodedData.toByteString()) + .setIsLast(true)) + .build()); + return null; + }) + .when(beamFnDataClient) + .registerReceiver(any(), any(), any(), any()); + + ProcessBundleHandler handler = + new ProcessBundleHandler( + PipelineOptionsFactory.create(), + Collections.emptySet(), + fnApiRegistry::get, + beamFnDataClient, + null /* beamFnStateGrpcClientCache */, + null /* finalizeBundleHandler */, + new ShortIdMap(), + executionStateSampler, + ImmutableMap.of( + DATA_INPUT_URN, + (context) -> + context.addIncomingDataEndpoint( + ApiServiceDescriptor.getDefaultInstance(), + StringUtf8Coder.of(), + (input) -> { + throw new IllegalStateException("TestException"); + })), + Caches.noop(), + new BundleProcessorCache(Duration.ZERO), + null /* dataSampler */); + assertThrows( + "TestException", + IllegalStateException.class, + () -> + handler.processBundle( + BeamFnApi.InstructionRequest.newBuilder() + .setInstructionId("instructionId") + .setProcessBundle( + BeamFnApi.ProcessBundleRequest.newBuilder() + .setProcessBundleDescriptorId("1L") + .setDataStreamId("dataStreamId")) + .build())); + + // Ensure that the named data stream is released even when the bundle fails. + InOrder inOrder = Mockito.inOrder(beamFnDataClient); + inOrder.verify(beamFnDataClient).retainDataStream(eq("dataStreamId")); + inOrder + .verify(beamFnDataClient) + .registerReceiver(eq("instructionId"), eq("dataStreamId"), any(), any()); + inOrder.verify(beamFnDataClient).poisonInstructionId(eq("instructionId")); + inOrder.verify(beamFnDataClient).releaseDataStream(eq("dataStreamId")); + verifyNoMoreInteractions(beamFnDataClient); + } + @Test public void testDataProcessingExceptionsArePropagated() throws Exception { BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java index 9d9efa0b9c49..5d99d75978b0 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java @@ -23,6 +23,8 @@ import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.empty; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; @@ -364,6 +366,116 @@ public StreamObserver data( } } + @Test + public void testNamedDataStreamClosedWhenNoLongerRetained() throws Exception { + String dataStreamId = "streamA"; + AtomicInteger connectionCount = new AtomicInteger(); + CountDownLatch firstConnection = new CountDownLatch(1); + CountDownLatch secondConnection = new CountDownLatch(2); + CountDownLatch streamTerminated = new CountDownLatch(1); + + Endpoints.ApiServiceDescriptor apiServiceDescriptor = + Endpoints.ApiServiceDescriptor.newBuilder() + .setUrl(this.getClass().getName() + "-" + UUID.randomUUID()) + .build(); + Server server = + InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()) + .addService( + new BeamFnDataGrpc.BeamFnDataImplBase() { + @Override + public StreamObserver data( + StreamObserver outboundObserver) { + connectionCount.incrementAndGet(); + firstConnection.countDown(); + secondConnection.countDown(); + return TestStreams.withOnNext(elements -> {}) + .withOnError(streamTerminated::countDown) + .withOnCompleted(streamTerminated::countDown) + .build(); + } + }) + .build(); + server.start(); + try { + ManagedChannel channel = + InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); + + BeamFnDataGrpcClient clientFactory = + new BeamFnDataGrpcClient( + (Endpoints.ApiServiceDescriptor descriptor) -> channel, + OutboundObserverFactory.trivial()); + + // Two bundles concurrently use the same named data stream. + clientFactory.retainDataStream(dataStreamId); + clientFactory.retainDataStream(dataStreamId); + clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId); + assertTrue(firstConnection.await(5, TimeUnit.SECONDS)); + assertEquals(1, connectionCount.get()); + + // Releasing one of the two usages should not close the stream. + clientFactory.releaseDataStream(dataStreamId); + assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS)); + + // Releasing the last usage should close the stream. + clientFactory.releaseDataStream(dataStreamId); + assertTrue(streamTerminated.await(5, TimeUnit.SECONDS)); + + // A subsequent usage of the same named data stream establishes a new stream. + clientFactory.retainDataStream(dataStreamId); + clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId); + assertTrue(secondConnection.await(5, TimeUnit.SECONDS)); + assertEquals(2, connectionCount.get()); + clientFactory.releaseDataStream(dataStreamId); + } finally { + server.shutdownNow(); + } + } + + @Test + public void testDefaultDataStreamIsNotClosedOnRelease() throws Exception { + CountDownLatch streamTerminated = new CountDownLatch(1); + + Endpoints.ApiServiceDescriptor apiServiceDescriptor = + Endpoints.ApiServiceDescriptor.newBuilder() + .setUrl(this.getClass().getName() + "-" + UUID.randomUUID()) + .build(); + Server server = + InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()) + .addService( + new BeamFnDataGrpc.BeamFnDataImplBase() { + @Override + public StreamObserver data( + StreamObserver outboundObserver) { + return TestStreams.withOnNext(elements -> {}) + .withOnError(streamTerminated::countDown) + .withOnCompleted(streamTerminated::countDown) + .build(); + } + }) + .build(); + server.start(); + try { + ManagedChannel channel = + InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); + + BeamFnDataGrpcClient clientFactory = + new BeamFnDataGrpcClient( + (Endpoints.ApiServiceDescriptor descriptor) -> channel, + OutboundObserverFactory.trivial()); + + clientFactory.retainDataStream(""); + StreamObserver outboundObserver = + clientFactory.getOutboundObserver(apiServiceDescriptor, ""); + clientFactory.releaseDataStream(""); + + // The default data stream is kept open and reused even after release. + assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS)); + assertSame(outboundObserver, clientFactory.getOutboundObserver(apiServiceDescriptor, "")); + } finally { + server.shutdownNow(); + } + } + @Test public void testForOutboundConsumer() throws Exception { CountDownLatch waitForInboundServerValuesCompletion = new CountDownLatch(2); From dce006fcae4e01aee0e9677aae97a4d269012469 Mon Sep 17 00:00:00 2001 From: Krishna Date: Wed, 8 Jul 2026 20:06:47 -0700 Subject: [PATCH 2/3] Fix validate-runner build by using public Jenkins repository URL --- .test-infra/validate-runner/build.gradle | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.test-infra/validate-runner/build.gradle b/.test-infra/validate-runner/build.gradle index 3992abb24dd4..5cc6bbae6495 100644 --- a/.test-infra/validate-runner/build.gradle +++ b/.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/" From 56b38bf2313dfd0a14a339009c2d9b151bb94800 Mon Sep 17 00:00:00 2001 From: Krishna Date: Thu, 9 Jul 2026 11:21:18 -0700 Subject: [PATCH 3/3] Fix PreCommit Java CI failure and resolve GcsUtil endpoint propagation --- .github/workflows/beam_PreCommit_Java.yml | 1 + runners/flink/flink_runner.gradle | 3 + .../sdk/extensions/gcp/util/GcsUtilV1.java | 55 +++++++++++++++++-- .../sdk/extensions/gcp/util/GcsUtilTest.java | 19 ++++++- 4 files changed, 70 insertions(+), 8 deletions(-) diff --git a/.github/workflows/beam_PreCommit_Java.yml b/.github/workflows/beam_PreCommit_Java.yml index b3e89b46218e..66e61d44c76f 100644 --- a/.github/workflows/beam_PreCommit_Java.yml +++ b/.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 diff --git a/runners/flink/flink_runner.gradle b/runners/flink/flink_runner.gradle index 837561ec71b7..240ac23b1a77 100644 --- a/runners/flink/flink_runner.gradle +++ b/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") } diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java index 1ad08f0ba1a2..ca4830088182 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java +++ b/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; @@ -261,12 +287,26 @@ public boolean shouldRetry(IOException e) { 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. */ diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java index d32ca162e3fd..4c3254998bfd 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java +++ b/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));