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/.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/.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/" 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/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)); 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);