This is an automated email from the ASF dual-hosted git repository. Amar3tto pushed a commit to branch fix-postreleasenightly in repository https://gitbox.apache.org/repos/asf/beam.git
commit 8b5f5ce73d45e3499b8ead3d6ebc612bf5273a9d Author: Vitaly Terentyev <[email protected]> AuthorDate: Wed Jun 24 15:36:52 2026 +0400 Revert #38863 --- .../beam/model/fn_execution/v1/beam_fn_api.proto | 38 +----- .../beam/model/pipeline/v1/beam_runner_api.proto | 4 - .../fnexecution/control/SdkHarnessClient.java | 2 +- .../runners/fnexecution/data/FnDataService.java | 3 +- .../runners/fnexecution/data/GrpcDataService.java | 13 +- .../fnexecution/data/GrpcDataServiceTest.java | 2 +- .../sdk/fn/data/BeamFnDataGrpcMultiplexer.java | 2 +- .../sdk/fn/data/BeamFnDataOutboundAggregator.java | 131 ++++++-------------- .../beam/sdk/util/construction/Environments.java | 1 - .../fn/data/BeamFnDataOutboundAggregatorTest.java | 135 +++++++++------------ .../java/org/apache/beam/fn/harness/FnHarness.java | 49 +++----- .../fn/harness/control/ProcessBundleHandler.java | 50 +++----- .../beam/fn/harness/data/BeamFnDataClient.java | 28 +++-- .../beam/fn/harness/data/BeamFnDataGrpcClient.java | 91 +++++--------- .../beam/fn/harness/BeamFnDataWriteRunnerTest.java | 44 +++---- .../beam/fn/harness/FnApiDoFnRunnerTest.java | 2 +- .../PTransformRunnerFactoryTestContext.java | 12 +- .../harness/control/ProcessBundleHandlerTest.java | 47 ++++--- .../fn/harness/data/BeamFnDataGrpcClientTest.java | 41 ++++--- 19 files changed, 283 insertions(+), 412 deletions(-) diff --git a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto index 80a0fa6f7f2..ecef3f2e7a9 100644 --- a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto +++ b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto @@ -120,9 +120,6 @@ message RemoteGrpcPort { service BeamFnControl { // Instructions sent by the runner to the SDK requesting different types // of work. - // - // Header metadata has the specified keys pairs: - // - "worker_id": the id of the sdk rpc Control( // A stream of responses to instructions the SDK was asked to be // performed. @@ -133,9 +130,6 @@ service BeamFnControl { // Used to get the full process bundle descriptors for bundles one // is asked to process. - // - // Header metadata has the specified keys pairs: - // - "worker_id": the id of the sdk rpc GetProcessBundleDescriptor(GetProcessBundleDescriptorRequest) returns ( ProcessBundleDescriptor) {} } @@ -422,22 +416,14 @@ message ProcessBundleRequest { // at https://s.apache.org/beam-fn-api-control-data-embedding. Elements elements = 3; - // Indicates that the runner has no state for the keys in this bundle + // indicates that the runner has no stare for the keys in this bundle // so SDk can safely begin stateful processing with a locally-generated - // initial empty state. + // initial empty state bool has_no_state = 4; - // Indicates that the runner will never process another bundle for the keys + // indicates that the runner will never process another bundle for the keys // in this bundle so state need not be included in the bundle commit. bool only_bundle_for_keys = 5; - - // (Optional) If non-empty, the ID of the data stream to use for all data - // requests related to this bundle. See comments at BeamFnData.Data for - // more details. - // - // The runner should only populate this field if the sdk advertises the - // beam:protocol:named_data_streams:v1 capability. - string data_stream_id = 6; } message ProcessBundleResponse { @@ -848,15 +834,7 @@ message Elements { // Stable service BeamFnData { - // Used to send data between harnesses. Sdks default to using an unnamed data stream - // (without "data_stream_id" header value) for bundles unless the runner requests another named stream to be - // used for a bundle. SDKs can advertise that they support named data streams with the capability - // `beam:protocol:named_data_streams:v1`. - // - // Header metadata has the specified keys pairs: - // - "worker_id": value is the id of the sdk - // - "data_stream_id": value is the id of the data stream, distinguishing it from other data streams from the same - // sdk. This field should only be populated if requested in a received ProcessBundleRequest from the runner. + // Used to send data between harnesses. rpc Data( // A stream of data representing input. stream Elements) @@ -922,9 +900,6 @@ message StateResponse { service BeamFnState { // Used to get/append/clear state stored by the runner on behalf of the SDK. - // - // Header metadata has the specified keys pairs: - // - "worker_id": the id of the sdk rpc State( // A stream of state instructions requested of the runner. stream StateRequest) @@ -1320,9 +1295,6 @@ message LogControl {} service BeamFnLogging { // Allows for the SDK to emit log entries which the runner can // associate with the active job. - // - // Header metadata has the specified keys pairs: - // - "worker_id": the id of the sdk rpc Logging( // A stream of log entries batched into lists emitted by the SDK harness. stream LogEntry.List) @@ -1384,8 +1356,6 @@ message WorkerStatusResponse { // API for SDKs to report debug-related statuses to runner during pipeline execution. service BeamFnWorkerStatus { - // Header metadata has the specified keys pairs: - // - "worker_id": the id of the sdk rpc WorkerStatus (stream WorkerStatusResponse) returns (stream WorkerStatusRequest) {} } diff --git a/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto b/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto index 5824c9bf4b7..67df8b9e800 100644 --- a/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto +++ b/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto @@ -1689,10 +1689,6 @@ message StandardProtocols { // Indicates whether the SDK supports multimap state. MULTIMAP_STATE = 12 [(beam_urn) = "beam:protocol:multimap_state:v1"]; - - // Indicates whether the SDK supports data stream ids being requested by the runner in - // ProcessBundleRequests. - NAMED_DATA_STREAMS = 13 [(beam_urn) = "beam:protocol:named_data_streams:v1"]; } } diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java index 704d298a195..682c45e3079 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java @@ -298,7 +298,7 @@ public class SdkHarnessClient implements AutoCloseable { ImmutableMap.Builder<LogicalEndpoint, FnDataReceiver<?>> receiverBuilder = ImmutableMap.builder(); BeamFnDataOutboundAggregator beamFnDataOutboundAggregator = - fnApiDataService.createOutboundAggregator(bundleId, false); + fnApiDataService.createOutboundAggregator(() -> bundleId, false); for (RemoteInputDestination remoteInput : remoteInputs) { LogicalEndpoint endpoint = LogicalEndpoint.data(bundleId, remoteInput.getPTransformId()); receiverBuilder.put( diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/FnDataService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/FnDataService.java index 657ec74553b..7c5f110eab2 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/FnDataService.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/FnDataService.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.fnexecution.data; +import java.util.function.Supplier; import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements; import org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator; import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver; @@ -68,5 +69,5 @@ public interface FnDataService { * <p>The returned aggregator is not thread safe. */ BeamFnDataOutboundAggregator createOutboundAggregator( - String processBundleId, boolean collectElementsIfNoFlushes); + Supplier<String> processBundleRequestIdSupplier, boolean collectElementsIfNoFlushes); } diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/GrpcDataService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/GrpcDataService.java index a3a8c324404..d4e45c8ccf8 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/GrpcDataService.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/GrpcDataService.java @@ -23,6 +23,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.function.Supplier; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements; import org.apache.beam.model.fnexecution.v1.BeamFnDataGrpc; @@ -174,13 +175,13 @@ public class GrpcDataService extends BeamFnDataGrpc.BeamFnDataImplBase @Override public BeamFnDataOutboundAggregator createOutboundAggregator( - String instructionId, boolean collectElementsIfNoFlushes) { + Supplier<String> processBundleRequestIdSupplier, boolean collectElementsIfNoFlushes) { try { - BeamFnDataOutboundAggregator aggregator = - new BeamFnDataOutboundAggregator(options, collectElementsIfNoFlushes); - aggregator.prepareForInstruction( - instructionId, connectedClient.get(3, TimeUnit.MINUTES).getOutboundObserver()); - return aggregator; + return new BeamFnDataOutboundAggregator( + options, + processBundleRequestIdSupplier, + connectedClient.get(3, TimeUnit.MINUTES).getOutboundObserver(), + collectElementsIfNoFlushes); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new RuntimeException(e); diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/data/GrpcDataServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/data/GrpcDataServiceTest.java index 363367f1087..f8446707750 100644 --- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/data/GrpcDataServiceTest.java +++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/data/GrpcDataServiceTest.java @@ -102,7 +102,7 @@ public class GrpcDataServiceTest { for (int i = 0; i < 3; ++i) { final String instructionId = Integer.toString(i); BeamFnDataOutboundAggregator aggregator = - service.createOutboundAggregator(instructionId, false); + service.createOutboundAggregator(() -> instructionId, false); aggregator.start(); FnDataReceiver<WindowedValue<String>> consumer = aggregator.registerOutputDataLocation(TRANSFORM_ID, CODER); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java index 0b9d6adab4f..8fec8b455cc 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java @@ -63,7 +63,7 @@ public class BeamFnDataGrpcMultiplexer implements AutoCloseable { private final Cache</*instructionId=*/ String, /*unused=*/ Boolean> poisonedInstructionIds; private static class PoisonedException extends RuntimeException { - private PoisonedException() { + public PoisonedException() { super("Instruction poisoned"); } }; diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregator.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregator.java index 04c7e6ef96b..9b9603706b4 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregator.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregator.java @@ -17,9 +17,6 @@ */ package org.apache.beam.sdk.fn.data; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; - import java.io.IOException; import java.util.Collections; import java.util.HashMap; @@ -31,6 +28,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; import javax.annotation.Nullable; import javax.annotation.concurrent.NotThreadSafe; import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements; @@ -55,7 +53,10 @@ import org.slf4j.LoggerFactory; * <p>The default time-based buffer threshold can be overridden by specifying the experiment {@code * data_buffer_time_limit_ms=<milliseconds>} */ -// The calling thread that invokes sendOrCollectBufferedDataAndFinishOutboundStreams synchronizes on +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +// The calling thread that invokes sendBufferedDataAndFinishOutboundStreams synchronizes on // flushLock effectively making the periodic flushing no longer read or mutate hasFlushedForBundle // and allowing the calling thread to read and mutate hasFlushedForBundle safely without needing to // create another memory barrier. Also note that flush is always invoked when synchronizing on @@ -71,54 +72,31 @@ public class BeamFnDataOutboundAggregator { private static final Logger LOG = LoggerFactory.getLogger(BeamFnDataOutboundAggregator.class); private final int sizeLimit; private final long timeLimit; - // The instructionId is set between prepareForInstruction and finishInstruction/discard. - private @Nullable String instructionId = null; - @VisibleForTesting final Map<String, Receiver<?>> outputDataReceivers = new HashMap<>(); - @VisibleForTesting final Map<TimerEndpoint, Receiver<?>> outputTimersReceivers = new HashMap<>(); - @Nullable private StreamObserver<Elements> outboundObserver; + private final Supplier<String> processBundleRequestIdSupplier; + @VisibleForTesting final Map<String, Receiver<?>> outputDataReceivers; + @VisibleForTesting final Map<TimerEndpoint, Receiver<?>> outputTimersReceivers; + private final StreamObserver<Elements> outboundObserver; @Nullable @VisibleForTesting ScheduledFuture<?> flushFuture; - private long bytesWrittenSinceFlush = 0; - private final Object flushLock = new Object(); + private long bytesWrittenSinceFlush; + private final Object flushLock; private final boolean collectElementsIfNoFlushes; - private boolean hasFlushedForBundle = false; + private boolean hasFlushedForBundle; - public BeamFnDataOutboundAggregator(PipelineOptions options, boolean collectElementsIfNoFlushes) { + public BeamFnDataOutboundAggregator( + PipelineOptions options, + Supplier<String> processBundleRequestIdSupplier, + StreamObserver<Elements> outboundObserver, + boolean collectElementsIfNoFlushes) { this.sizeLimit = getSizeLimit(options); this.timeLimit = getTimeLimit(options); this.collectElementsIfNoFlushes = collectElementsIfNoFlushes; - } - - public void prepareForInstruction( - String instructionId, StreamObserver<Elements> outboundObserver) { - if (timeLimit > 0) { - synchronized (flushLock) { - checkState(this.instructionId == null && this.outboundObserver == null); - this.instructionId = instructionId; - this.outboundObserver = outboundObserver; - } - } else { - checkState(this.instructionId == null && this.outboundObserver == null); - this.instructionId = instructionId; - this.outboundObserver = outboundObserver; - } - } - - public void finishInstruction() { - if (flushFuture != null) { - synchronized (flushLock) { - checkState( - this.instructionId != null && this.outboundObserver != null, - "instruction was not started or previously completed"); - checkState(bytesWrittenSinceFlush == 0, "bytes were not flushed for instruction"); - this.instructionId = null; - this.outboundObserver = null; - } - } else { - checkState(this.instructionId != null && this.outboundObserver != null); - checkState(bytesWrittenSinceFlush == 0, "bytes were not flushed for instruction"); - this.instructionId = null; - this.outboundObserver = null; - } + this.outputDataReceivers = new HashMap<>(); + this.outputTimersReceivers = new HashMap<>(); + this.outboundObserver = outboundObserver; + this.processBundleRequestIdSupplier = processBundleRequestIdSupplier; + this.bytesWrittenSinceFlush = 0L; + this.flushLock = new Object(); + this.hasFlushedForBundle = false; } /** Starts the flushing daemon thread if data_buffer_time_limit_ms is set. */ @@ -188,7 +166,7 @@ public class BeamFnDataOutboundAggregator { } Elements.Builder elements = convertBufferForTransmission(); if (elements.getDataCount() > 0 || elements.getTimersCount() > 0) { - checkNotNull(outboundObserver).onNext(elements.build()); + outboundObserver.onNext(elements.build()); } hasFlushedForBundle = true; } @@ -199,15 +177,10 @@ public class BeamFnDataOutboundAggregator { * collectElementsIfNoFlushes=true, and there was no previous flush in this bundle, otherwise * returns null. */ - @Nullable public Elements sendOrCollectBufferedDataAndFinishOutboundStreams() { if (outputTimersReceivers.isEmpty() && outputDataReceivers.isEmpty()) { return null; } - String instructionId = - checkNotNull( - this.instructionId, - "This method should only be called between prepareForInstruction and finishInstruction"); Elements.Builder bufferedElements; if (timeLimit > 0) { synchronized (flushLock) { @@ -218,14 +191,14 @@ public class BeamFnDataOutboundAggregator { } LOG.debug( "Closing streams for instruction {} and outbound data {} and timers {}.", - instructionId, + processBundleRequestIdSupplier.get(), outputDataReceivers, outputTimersReceivers); for (Map.Entry<String, Receiver<?>> entry : outputDataReceivers.entrySet()) { String pTransformId = entry.getKey(); bufferedElements .addDataBuilder() - .setInstructionId(instructionId) + .setInstructionId(processBundleRequestIdSupplier.get()) .setTransformId(pTransformId) .setIsLast(true); entry.getValue().resetStats(); @@ -234,60 +207,35 @@ public class BeamFnDataOutboundAggregator { TimerEndpoint timerKey = entry.getKey(); bufferedElements .addTimersBuilder() - .setInstructionId(instructionId) + .setInstructionId(processBundleRequestIdSupplier.get()) .setTransformId(timerKey.pTransformId) .setTimerFamilyId(timerKey.timerFamilyId) .setIsLast(true); entry.getValue().resetStats(); } - // This is the end of the bundle so we reset state to prepare for future bundles. if (collectElementsIfNoFlushes && !hasFlushedForBundle) { return bufferedElements.build(); } - checkNotNull(outboundObserver).onNext(bufferedElements.build()); + outboundObserver.onNext(bufferedElements.build()); + // This is now at the end of a bundle, so we reset hasFlushedForBundle to prepare for new + // bundles. hasFlushedForBundle = false; return null; } // Send the elements to the StreamObserver associated with this aggregator. public void sendElements(Elements elements) { - if (timeLimit > 0) { - synchronized (flushLock) { - checkNotNull(outboundObserver).onNext(elements); - } - } else { - checkNotNull(outboundObserver).onNext(elements); - } + outboundObserver.onNext(elements); } - // Prepares for discarding the aggregator without preserving its output or - // preparing it for reuse. public void discard() { - if (timeLimit > 0) { - // Short-circuit the possibly concurrently running flush. - synchronized (flushLock) { - bytesWrittenSinceFlush = 0L; - finishInstruction(); - } - if (flushFuture != null) { - flushFuture.cancel(false); - } - } else { - bytesWrittenSinceFlush = 0L; - finishInstruction(); + if (flushFuture != null) { + flushFuture.cancel(true); } } private Elements.Builder convertBufferForTransmission() { Elements.Builder bufferedElements = Elements.newBuilder(); - if (bytesWrittenSinceFlush == 0) { - return bufferedElements; - } - bytesWrittenSinceFlush = 0L; - String instructionId = - checkNotNull( - this.instructionId, - "This method should only be called between prepareForInstruction and finishInstruction"); for (Map.Entry<String, Receiver<?>> entry : outputDataReceivers.entrySet()) { if (!entry.getValue().hasBufferedOutput()) { continue; @@ -295,7 +243,7 @@ public class BeamFnDataOutboundAggregator { ByteString bytes = entry.getValue().toByteStringAndResetBuffer(); bufferedElements .addDataBuilder() - .setInstructionId(instructionId) + .setInstructionId(processBundleRequestIdSupplier.get()) .setTransformId(entry.getKey()) .setData(bytes); } @@ -306,11 +254,12 @@ public class BeamFnDataOutboundAggregator { ByteString bytes = entry.getValue().toByteStringAndResetBuffer(); bufferedElements .addTimersBuilder() - .setInstructionId(instructionId) + .setInstructionId(processBundleRequestIdSupplier.get()) .setTransformId(entry.getKey().pTransformId) .setTimerFamilyId(entry.getKey().timerFamilyId) .setTimers(bytes); } + bytesWrittenSinceFlush = 0L; return bufferedElements; } @@ -328,7 +277,7 @@ public class BeamFnDataOutboundAggregator { /** Check if the flush thread failed with an exception. */ private void checkFlushThreadException() throws IOException { - if (flushFuture != null && flushFuture.isDone()) { + if (timeLimit > 0 && flushFuture.isDone()) { try { flushFuture.get(); throw new IOException("Periodic flushing thread finished unexpectedly."); @@ -404,12 +353,10 @@ public class BeamFnDataOutboundAggregator { } } - @VisibleForTesting public long getByteCount() { return perBundleByteCount; } - @VisibleForTesting public long getElementCount() { return perBundleElementCount; } @@ -445,7 +392,7 @@ public class BeamFnDataOutboundAggregator { } @Override - public boolean equals(@Nullable Object o) { + public boolean equals(Object o) { if (this == o) { return true; } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java index c3b1a7a5235..969bda88d07 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java @@ -522,7 +522,6 @@ public class Environments { capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.SDK_CONSUMING_RECEIVED_DATA)); capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.ORDERED_LIST_STATE)); capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.MULTIMAP_STATE)); - capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.NAMED_DATA_STREAMS)); return capabilities.build(); } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregatorTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregatorTest.java index 9bcf615d638..092ba200c94 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregatorTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregatorTest.java @@ -20,7 +20,6 @@ package org.apache.beam.sdk.fn.data; import static org.hamcrest.Matchers.empty; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertThrows; import static org.junit.Assert.fail; import java.io.IOException; @@ -76,12 +75,13 @@ public class BeamFnDataOutboundAggregatorTest { final List<Elements> values = new ArrayList<>(); final AtomicBoolean onCompletedWasCalled = new AtomicBoolean(); BeamFnDataOutboundAggregator aggregator = - new BeamFnDataOutboundAggregator(PipelineOptionsFactory.create(), false); - aggregator.prepareForInstruction( - endpoint.getInstructionId(), - TestStreams.<Elements>withOnNext(values::add) - .withOnCompleted(() -> onCompletedWasCalled.set(true)) - .build()); + new BeamFnDataOutboundAggregator( + PipelineOptionsFactory.create(), + endpoint::getInstructionId, + TestStreams.<Elements>withOnNext(values::add) + .withOnCompleted(() -> onCompletedWasCalled.set(true)) + .build(), + false); // Test that nothing is emitted till the default buffer size is surpassed. FnDataReceiver<byte[]> dataReceiver = registerOutputLocation(aggregator, endpoint, CODER); @@ -124,12 +124,14 @@ public class BeamFnDataOutboundAggregatorTest { options .as(ExperimentalOptions.class) .setExperiments(Arrays.asList("data_buffer_size_limit=100")); - BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, false); - aggregator.prepareForInstruction( - endpoint.getInstructionId(), - TestStreams.<Elements>withOnNext(values::add) - .withOnCompleted(() -> onCompletedWasCalled.set(true)) - .build()); + BeamFnDataOutboundAggregator aggregator = + new BeamFnDataOutboundAggregator( + options, + endpoint::getInstructionId, + TestStreams.<Elements>withOnNext(values::add) + .withOnCompleted(() -> onCompletedWasCalled.set(true)) + .build(), + false); // Test that nothing is emitted till the default buffer size is surpassed. FnDataReceiver<byte[]> dataReceiver = registerOutputLocation(aggregator, endpoint, CODER); aggregator.start(); @@ -185,16 +187,18 @@ public class BeamFnDataOutboundAggregatorTest { .as(ExperimentalOptions.class) .setExperiments(Arrays.asList("data_buffer_time_limit_ms=1")); final CountDownLatch waitForFlush = new CountDownLatch(1); - BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, false); - aggregator.prepareForInstruction( - endpoint.getInstructionId(), - TestStreams.withOnNext( - (Consumer<Elements>) - e -> { - values.add(e); - waitForFlush.countDown(); - }) - .build()); + BeamFnDataOutboundAggregator aggregator = + new BeamFnDataOutboundAggregator( + options, + endpoint::getInstructionId, + TestStreams.withOnNext( + (Consumer<Elements>) + e -> { + values.add(e); + waitForFlush.countDown(); + }) + .build(), + false); // Test that it emits when time passed the time limit FnDataReceiver<byte[]> dataReceiver = registerOutputLocation(aggregator, endpoint, CODER); @@ -210,15 +214,17 @@ public class BeamFnDataOutboundAggregatorTest { options .as(ExperimentalOptions.class) .setExperiments(Arrays.asList("data_buffer_time_limit_ms=1")); - BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, false); - aggregator.prepareForInstruction( - endpoint.getInstructionId(), - TestStreams.withOnNext( - (Consumer<Elements>) - e -> { - throw new RuntimeException(""); - }) - .build()); + BeamFnDataOutboundAggregator aggregator = + new BeamFnDataOutboundAggregator( + options, + endpoint::getInstructionId, + TestStreams.withOnNext( + (Consumer<Elements>) + e -> { + throw new RuntimeException(""); + }) + .build(), + false); // Test that it emits when time passed the time limit FnDataReceiver<byte[]> dataReceiver = registerOutputLocation(aggregator, endpoint, CODER); @@ -237,15 +243,17 @@ public class BeamFnDataOutboundAggregatorTest { // expected } - aggregator = new BeamFnDataOutboundAggregator(options, false); - aggregator.prepareForInstruction( - endpoint.getInstructionId(), - TestStreams.withOnNext( - (Consumer<Elements>) - e -> { - throw new RuntimeException(""); - }) - .build()); + aggregator = + new BeamFnDataOutboundAggregator( + options, + endpoint::getInstructionId, + TestStreams.withOnNext( + (Consumer<Elements>) + e -> { + throw new RuntimeException(""); + }) + .build(), + false); dataReceiver = registerOutputLocation(aggregator, endpoint, CODER); aggregator.start(); dataReceiver.accept(new byte[1]); @@ -271,12 +279,14 @@ public class BeamFnDataOutboundAggregatorTest { options .as(ExperimentalOptions.class) .setExperiments(Arrays.asList("data_buffer_size_limit=100")); - BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, false); - aggregator.prepareForInstruction( - endpoint.getInstructionId(), - TestStreams.<Elements>withOnNext(values::add) - .withOnCompleted(() -> onCompletedWasCalled.set(true)) - .build()); + BeamFnDataOutboundAggregator aggregator = + new BeamFnDataOutboundAggregator( + options, + endpoint::getInstructionId, + TestStreams.<Elements>withOnNext(values::add) + .withOnCompleted(() -> onCompletedWasCalled.set(true)) + .build(), + false); // Test that nothing is emitted till the default buffer size is surpassed. LogicalEndpoint additionalEndpoint = LogicalEndpoint.data( @@ -324,37 +334,6 @@ public class BeamFnDataOutboundAggregatorTest { checkEqualInAnyOrder(builder.build(), values.get(1)); } - @Test - public void testInstructionLifecycle() { - BeamFnDataOutboundAggregator aggregator = - new BeamFnDataOutboundAggregator(PipelineOptionsFactory.create(), false); - assertThrows( - NullPointerException.class, () -> aggregator.sendElements(Elements.getDefaultInstance())); - aggregator.prepareForInstruction( - "testInstruction", - TestStreams.withOnNext( - (Consumer<Elements>) - e -> { - throw new RuntimeException(""); - }) - .build()); - assertThrows( - IllegalStateException.class, - () -> - aggregator.prepareForInstruction( - "testInstruction", - TestStreams.withOnNext( - (Consumer<Elements>) - e -> { - throw new RuntimeException(""); - }) - .build())); - aggregator.finishInstruction(); - assertThrows( - NullPointerException.class, () -> aggregator.sendElements(Elements.getDefaultInstance())); - assertThrows(IllegalStateException.class, aggregator::finishInstruction); - } - private void checkEqualInAnyOrder(Elements first, Elements second) { MatcherAssert.assertThat( first.getDataList(), Matchers.containsInAnyOrder(second.getDataList().toArray())); diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java index 60e83251f14..703e726739a 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java @@ -17,8 +17,6 @@ */ package org.apache.beam.fn.harness; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; - import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import java.nio.charset.StandardCharsets; @@ -32,6 +30,7 @@ import java.util.Map; import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.function.Function; +import javax.annotation.Nullable; import org.apache.beam.fn.harness.control.BeamFnControlClient; import org.apache.beam.fn.harness.control.ExecutionStateSampler; import org.apache.beam.fn.harness.control.FinalizeBundleHandler; @@ -73,7 +72,6 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.Vi 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.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.MoreExecutors; -import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -94,6 +92,9 @@ import org.slf4j.LoggerFactory; * for further details. * </ul> */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) public class FnHarness { private static final String HARNESS_ID = "HARNESS_ID"; private static final String CONTROL_API_SERVICE_DESCRIPTOR = "CONTROL_API_SERVICE_DESCRIPTOR"; @@ -137,31 +138,22 @@ public class FnHarness { } public static void main(String[] args) throws Exception { - Function<String, @Nullable String> environmentVarGetter = System::getenv; - main(environmentVarGetter); + main(System::getenv); } @VisibleForTesting - public static void main(Function<String, @Nullable String> environmentVarGetter) - throws Exception { + public static void main(Function<String, String> environmentVarGetter) throws Exception { JvmInitializers.runOnStartup(); Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor = - getApiServiceDescriptor( - checkNotNull( - environmentVarGetter.apply(LOGGING_API_SERVICE_DESCRIPTOR), - "LOGGING_API_SERVICE_DESCRIPTOR env var must be set.")); + getApiServiceDescriptor(environmentVarGetter.apply(LOGGING_API_SERVICE_DESCRIPTOR)); Endpoints.ApiServiceDescriptor controlApiServiceDescriptor = - getApiServiceDescriptor( - checkNotNull( - environmentVarGetter.apply(CONTROL_API_SERVICE_DESCRIPTOR), - "CONTROL_API_SERVICE_DESCRIPTOR env var must be set.")); - - @Nullable String envVar = environmentVarGetter.apply(STATUS_API_SERVICE_DESCRIPTOR); - Endpoints.@Nullable ApiServiceDescriptor statusApiServiceDescriptor = - (envVar == null) ? null : getApiServiceDescriptor(envVar); - String id = - checkNotNull(environmentVarGetter.apply(HARNESS_ID), "HARNESS_ID env var must be set."); + getApiServiceDescriptor(environmentVarGetter.apply(CONTROL_API_SERVICE_DESCRIPTOR)); + Endpoints.ApiServiceDescriptor statusApiServiceDescriptor = + environmentVarGetter.apply(STATUS_API_SERVICE_DESCRIPTOR) == null + ? null + : getApiServiceDescriptor(environmentVarGetter.apply(STATUS_API_SERVICE_DESCRIPTOR)); + String id = environmentVarGetter.apply(HARNESS_ID); System.out.format("SDK Fn Harness started%n"); System.out.format("Harness ID %s%n", id); @@ -169,11 +161,11 @@ public class FnHarness { System.out.format("Control location %s%n", controlApiServiceDescriptor); System.out.format("Status location %s%n", statusApiServiceDescriptor); + String pipelineOptionsJson = environmentVarGetter.apply(PIPELINE_OPTIONS); // Try looking for a file first. If that exists it should override PIPELINE_OPTIONS to avoid // maxing out the kernel's environment space - @Nullable String pipelineOptionsJson = null; try { - @Nullable String pipelineOptionsPath = environmentVarGetter.apply(PIPELINE_OPTIONS_FILE); + String pipelineOptionsPath = environmentVarGetter.apply(PIPELINE_OPTIONS_FILE); System.out.format("Pipeline Options File %s%n", pipelineOptionsPath); if (pipelineOptionsPath != null) { Path filePath = Paths.get(pipelineOptionsPath); @@ -187,12 +179,11 @@ public class FnHarness { } catch (Exception e) { System.out.format("Problem loading pipeline options from file: %s%n", e.getMessage()); } - if (pipelineOptionsJson == null) { - pipelineOptionsJson = checkNotNull(environmentVarGetter.apply(PIPELINE_OPTIONS)); - } + System.out.format("Pipeline options %s%n", pipelineOptionsJson); // TODO: https://github.com/apache/beam/issues/30301 pipelineOptionsJson = removeNestedKey(pipelineOptionsJson, "impersonateServiceAccount"); + PipelineOptions options = PipelineOptionsTranslation.fromJson(pipelineOptionsJson); String runnerCapabilitesOrNull = environmentVarGetter.apply(RUNNER_CAPABILITIES); @@ -228,7 +219,7 @@ public class FnHarness { Set<String> runnerCapabilities, Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor, Endpoints.ApiServiceDescriptor controlApiServiceDescriptor, - Endpoints.@Nullable ApiServiceDescriptor statusApiServiceDescriptor) + @Nullable Endpoints.ApiServiceDescriptor statusApiServiceDescriptor) throws Exception { ManagedChannelFactory channelFactory; if (ExperimentalOptions.hasExperiment(options, "beam_fn_api_epoll")) { @@ -272,7 +263,7 @@ public class FnHarness { Set<String> runnerCapabilites, Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor, Endpoints.ApiServiceDescriptor controlApiServiceDescriptor, - Endpoints.@Nullable ApiServiceDescriptor statusApiServiceDescriptor, + Endpoints.ApiServiceDescriptor statusApiServiceDescriptor, ManagedChannelFactory channelFactory, OutboundObserverFactory outboundObserverFactory, Cache<Object, Object> processWideCache) @@ -327,7 +318,7 @@ public class FnHarness { BeamFnControlGrpc.newBlockingStub(channel); BeamFnDataGrpcClient beamFnDataMultiplexer = - new BeamFnDataGrpcClient(channelFactory::forDescriptor, outboundObserverFactory); + new BeamFnDataGrpcClient(options, channelFactory::forDescriptor, outboundObserverFactory); BeamFnStateGrpcClientCache beamFnStateGrpcClientCache = new BeamFnStateGrpcClientCache(idGenerator, channelFactory, outboundObserverFactory); 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 449afd6a024..5a57b137bf6 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 @@ -224,6 +224,7 @@ public class ProcessBundleHandler { private void addRunnerAndConsumersForPTransformRecursively( BeamFnStateClient beamFnStateClient, + BeamFnDataClient queueingClient, String pTransformId, PTransform pTransform, Supplier<String> processBundleInstructionId, @@ -256,6 +257,7 @@ public class ProcessBundleHandler { for (String consumingPTransformId : pCollectionIdsToConsumingPTransforms.get(pCollectionId)) { addRunnerAndConsumersForPTransformRecursively( beamFnStateClient, + queueingClient, consumingPTransformId, processBundleDescriptor.getTransformsMap().get(consumingPTransformId), processBundleInstructionId, @@ -313,7 +315,7 @@ public class ProcessBundleHandler { @Override public BeamFnDataClient getBeamFnDataClient() { - return beamFnDataClient; + return queueingClient; } @Override @@ -376,8 +378,9 @@ public class ProcessBundleHandler { outboundAggregatorMap.computeIfAbsent( apiServiceDescriptor, asd -> - new BeamFnDataOutboundAggregator( - options, + queueingClient.createOutboundAggregator( + asd, + processBundleInstructionId, runnerCapabilities.contains( BeamUrns.getUrn( StandardRunnerProtocols.Enum @@ -388,19 +391,21 @@ public class ProcessBundleHandler { @Override public <T> FnDataReceiver<Timer<T>> addOutgoingTimersEndpoint( String timerFamilyId, org.apache.beam.sdk.coders.Coder<Timer<T>> coder) { + BeamFnDataOutboundAggregator aggregator; if (!processBundleDescriptor.hasTimerApiServiceDescriptor()) { throw new IllegalStateException( String.format( - "Timers are unsupported because the ProcessBundleRequest %s does not" - + " provide a timer ApiServiceDescriptor.", + "Timers are unsupported because the " + + "ProcessBundleRequest %s does not provide a timer ApiServiceDescriptor.", processBundleInstructionId.get())); } - BeamFnDataOutboundAggregator aggregator = + aggregator = outboundAggregatorMap.computeIfAbsent( processBundleDescriptor.getTimerApiServiceDescriptor(), asd -> - new BeamFnDataOutboundAggregator( - options, + queueingClient.createOutboundAggregator( + asd, + processBundleInstructionId, runnerCapabilities.contains( BeamUrns.getUrn( StandardRunnerProtocols.Enum @@ -494,8 +499,6 @@ public class ProcessBundleHandler { */ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest request) throws Exception { - String instructionId = request.getInstructionId(); - String dataStreamId = request.getProcessBundle().getDataStreamId(); @Nullable BundleProcessor bundleProcessor = null; try { bundleProcessor = @@ -512,20 +515,13 @@ public class ProcessBundleHandler { } })); - for (Map.Entry<ApiServiceDescriptor, BeamFnDataOutboundAggregator> entry : - bundleProcessor.getOutboundAggregators().entrySet()) { - BeamFnDataOutboundAggregator aggregator = entry.getValue(); - aggregator.prepareForInstruction( - instructionId, beamFnDataClient.getOutboundObserver(entry.getKey(), dataStreamId)); - } - PTransformFunctionRegistry startFunctionRegistry = bundleProcessor.getStartFunctionRegistry(); PTransformFunctionRegistry finishFunctionRegistry = bundleProcessor.getFinishFunctionRegistry(); ExecutionStateTracker stateTracker = bundleProcessor.getStateTracker(); ProcessBundleResponse.Builder response = ProcessBundleResponse.newBuilder(); try (HandleStateCallsForBundle beamFnStateClient = bundleProcessor.getBeamFnStateClient()) { - stateTracker.start(instructionId); + stateTracker.start(request.getInstructionId()); try { // Already in reverse topological order so we don't need to do anything. for (ThrowingRunnable startFunction : startFunctionRegistry.getFunctions()) { @@ -549,14 +545,12 @@ public class ProcessBundleHandler { } else if (!bundleProcessor.getInboundEndpointApiServiceDescriptors().isEmpty()) { BeamFnDataInboundObserver observer = bundleProcessor.getInboundObserver(); beamFnDataClient.registerReceiver( - instructionId, - dataStreamId, + request.getInstructionId(), bundleProcessor.getInboundEndpointApiServiceDescriptors(), observer); observer.awaitCompletion(); beamFnDataClient.unregisterReceiver( - instructionId, - dataStreamId, + request.getInstructionId(), bundleProcessor.getInboundEndpointApiServiceDescriptors()); } @@ -587,7 +581,7 @@ public class ProcessBundleHandler { if (!bundleProcessor.getBundleFinalizationCallbackRegistrations().isEmpty()) { finalizeBundleHandler.registerCallbacks( - instructionId, + bundleProcessor.getInstructionId(), ImmutableList.copyOf(bundleProcessor.getBundleFinalizationCallbackRegistrations())); response.setRequiresFinalization(true); } @@ -605,7 +599,7 @@ public class ProcessBundleHandler { } catch (Exception e) { LOG.debug( "Error processing bundle {} with bundleProcessor for {} after exception", - instructionId, + request.getInstructionId(), request.getProcessBundle().getProcessBundleDescriptorId(), e); if (bundleProcessor != null) { @@ -613,7 +607,7 @@ public class ProcessBundleHandler { bundleProcessorCache.discard(bundleProcessor); } // Ensure that if more data arrives for the instruction it is discarded. - beamFnDataClient.poisonInstructionId(instructionId); + beamFnDataClient.poisonInstructionId(request.getInstructionId()); throw e; } } @@ -635,10 +629,6 @@ public class ProcessBundleHandler { collectedElements.add(elements); } if (!hasFlushedAggregator) { - for (BeamFnDataOutboundAggregator aggregator : - bundleProcessor.getOutboundAggregators().values()) { - aggregator.finishInstruction(); - } Elements.Builder elementsToEmbed = Elements.newBuilder(); for (Elements collectedElement : collectedElements) { elementsToEmbed.mergeFrom(collectedElement); @@ -655,7 +645,6 @@ public class ProcessBundleHandler { if (elements != null) { aggregator.sendElements(elements); } - aggregator.finishInstruction(); } } } @@ -886,6 +875,7 @@ public class ProcessBundleHandler { addRunnerAndConsumersForPTransformRecursively( beamFnStateClient, + beamFnDataClient, entry.getKey(), entry.getValue(), bundleProcessor::getInstructionId, 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 1a50f5b448c..94d59d0fcb6 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 @@ -18,12 +18,13 @@ package org.apache.beam.fn.harness.data; import java.util.List; +import java.util.function.Supplier; import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements; import org.apache.beam.model.pipeline.v1.Endpoints; import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor; +import org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator; import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver; import org.apache.beam.sdk.fn.data.FnDataReceiver; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; /** * The {@link BeamFnDataClient} is able to forward inbound elements to a {@link FnDataReceiver} and @@ -46,7 +47,6 @@ public interface BeamFnDataClient { */ void registerReceiver( String instructionId, - String dataStreamId, List<ApiServiceDescriptor> apiServiceDescriptors, CloseableFnDataReceiver<Elements> receiver); @@ -58,8 +58,7 @@ public interface BeamFnDataClient { * to the {@link BeamFnDataClient} during a future {@link FnDataReceiver#accept} invocation or via * a call to {@link #poisonInstructionId}. */ - void unregisterReceiver( - String instructionId, String dataStreamId, List<ApiServiceDescriptor> apiServiceDescriptors); + void unregisterReceiver(String instructionId, List<ApiServiceDescriptor> apiServiceDescriptors); /** * Poisons the instruction id, indicating that future data arriving for it should be discarded. @@ -69,7 +68,22 @@ public interface BeamFnDataClient { */ void poisonInstructionId(String instructionId); - /** Get the outbound observer for the specified apiServiceDescriptor and dataStreamId. */ - StreamObserver<Elements> getOutboundObserver( - Endpoints.ApiServiceDescriptor apiServiceDescriptor, String dataStreamId); + /** + * Creates a {@link BeamFnDataOutboundAggregator} for buffering and sending outbound data and + * timers over the data plane. It is important that {@link + * BeamFnDataOutboundAggregator#sendOrCollectBufferedDataAndFinishOutboundStreams()} is called on + * the returned BeamFnDataOutboundAggregator at the end of each bundle. If + * collectElementsIfNoFlushes is set to true, {@link + * BeamFnDataOutboundAggregator#sendOrCollectBufferedDataAndFinishOutboundStreams()} returns the + * buffered elements instead of sending it through the outbound StreamObserver if there's no + * previous flush. + * + * <p>Closing the returned aggregator signals the end of the streams. + * + * <p>The returned aggregator is not thread safe. + */ + BeamFnDataOutboundAggregator createOutboundAggregator( + Endpoints.ApiServiceDescriptor apiServiceDescriptor, + Supplier<String> processBundleRequestIdSupplier, + boolean collectElementsIfNoFlushes); } 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 2f2a6b0fc66..499d816f8cc 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 @@ -18,22 +18,20 @@ package org.apache.beam.fn.harness.data; import java.util.List; -import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.function.Function; -import javax.annotation.Nullable; +import java.util.function.Supplier; import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements; import org.apache.beam.model.fnexecution.v1.BeamFnDataGrpc; import org.apache.beam.model.pipeline.v1.Endpoints; import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor; import org.apache.beam.sdk.fn.data.BeamFnDataGrpcMultiplexer; +import org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator; import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver; import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; +import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Metadata; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.MetadataUtils; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -46,42 +44,17 @@ public class BeamFnDataGrpcClient implements BeamFnDataClient { private static final Logger LOG = LoggerFactory.getLogger(BeamFnDataGrpcClient.class); - private static class MultiplexerKey { - private final Endpoints.ApiServiceDescriptor apiServiceDescriptor; - private final String dataStreamId; - - private MultiplexerKey( - Endpoints.ApiServiceDescriptor apiServiceDescriptor, String dataStreamId) { - this.apiServiceDescriptor = apiServiceDescriptor; - this.dataStreamId = dataStreamId; - } - - @Override - public boolean equals(@Nullable Object o) { - if (this == o) { - return true; - } - if (!(o instanceof MultiplexerKey)) { - return false; - } - MultiplexerKey that = (MultiplexerKey) o; - return Objects.equals(dataStreamId, that.dataStreamId) - && Objects.equals(apiServiceDescriptor, that.apiServiceDescriptor); - } - - @Override - public int hashCode() { - return Objects.hash(apiServiceDescriptor, dataStreamId); - } - } - - private final ConcurrentMap<MultiplexerKey, BeamFnDataGrpcMultiplexer> multiplexerCache; + private final ConcurrentMap<Endpoints.ApiServiceDescriptor, BeamFnDataGrpcMultiplexer> + multiplexerCache; private final Function<Endpoints.ApiServiceDescriptor, ManagedChannel> channelFactory; private final OutboundObserverFactory outboundObserverFactory; + private final PipelineOptions options; public BeamFnDataGrpcClient( + PipelineOptions options, Function<Endpoints.ApiServiceDescriptor, ManagedChannel> channelFactory, OutboundObserverFactory outboundObserverFactory) { + this.options = options; this.channelFactory = channelFactory; this.outboundObserverFactory = outboundObserverFactory; this.multiplexerCache = new ConcurrentHashMap<>(); @@ -90,22 +63,21 @@ public class BeamFnDataGrpcClient implements BeamFnDataClient { @Override public void registerReceiver( String instructionId, - String dataStreamId, List<ApiServiceDescriptor> apiServiceDescriptors, CloseableFnDataReceiver<Elements> receiver) { LOG.debug("Registering consumer for {}", instructionId); for (int i = 0, size = apiServiceDescriptors.size(); i < size; i++) { - BeamFnDataGrpcMultiplexer client = getMultiplexer(apiServiceDescriptors.get(i), dataStreamId); + BeamFnDataGrpcMultiplexer client = getClientFor(apiServiceDescriptors.get(i)); client.registerConsumer(instructionId, receiver); } } @Override public void unregisterReceiver( - String instructionId, String dataStreamId, List<ApiServiceDescriptor> apiServiceDescriptors) { + String instructionId, List<ApiServiceDescriptor> apiServiceDescriptors) { LOG.debug("Unregistering consumer for {}", instructionId); for (int i = 0, size = apiServiceDescriptors.size(); i < size; i++) { - BeamFnDataGrpcMultiplexer client = getMultiplexer(apiServiceDescriptors.get(i), dataStreamId); + BeamFnDataGrpcMultiplexer client = getClientFor(apiServiceDescriptors.get(i)); client.unregisterConsumer(instructionId); } } @@ -119,32 +91,25 @@ public class BeamFnDataGrpcClient implements BeamFnDataClient { } @Override - public StreamObserver<Elements> getOutboundObserver( - ApiServiceDescriptor apiServiceDescriptor, String dataStreamId) { - return getMultiplexer(apiServiceDescriptor, dataStreamId).getOutboundObserver(); + public BeamFnDataOutboundAggregator createOutboundAggregator( + ApiServiceDescriptor apiServiceDescriptor, + Supplier<String> processBundleRequestIdSupplier, + boolean collectElementsIfNoFlushes) { + return new BeamFnDataOutboundAggregator( + options, + processBundleRequestIdSupplier, + getClientFor(apiServiceDescriptor).getOutboundObserver(), + collectElementsIfNoFlushes); } - private BeamFnDataGrpcMultiplexer getMultiplexer( - Endpoints.ApiServiceDescriptor apiServiceDescriptor, String dataStreamId) { - MultiplexerKey key = new MultiplexerKey(apiServiceDescriptor, dataStreamId); + private BeamFnDataGrpcMultiplexer getClientFor( + Endpoints.ApiServiceDescriptor apiServiceDescriptor) { 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); - }); + apiServiceDescriptor, + (Endpoints.ApiServiceDescriptor descriptor) -> + new BeamFnDataGrpcMultiplexer( + descriptor, + outboundObserverFactory, + BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor))::data)); } } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java index 2882b75a259..70a894e7b37 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java @@ -32,6 +32,7 @@ import java.util.List; import java.util.Map; import java.util.ServiceLoader; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; import org.apache.beam.fn.harness.PTransformRunnerFactory.Registrar; import org.apache.beam.fn.harness.data.BeamFnDataClient; import org.apache.beam.model.fnexecution.v1.BeamFnApi; @@ -106,29 +107,19 @@ public class BeamFnDataWriteRunnerTest { MockitoAnnotations.initMocks(this); } - @Test - public void testReuseForMultipleBundles() throws Exception { - AtomicReference<String> bundleId = new AtomicReference<>("0"); - String localInputId = "inputPC"; - RunnerApi.PTransform pTransform = - RemoteGrpcPortWrite.writeToPort(localInputId, PORT_SPEC).toPTransform(); - - List<WindowedValue<String>> output0 = new ArrayList<>(); - List<WindowedValue<String>> output1 = new ArrayList<>(); - Map<ApiServiceDescriptor, BeamFnDataOutboundAggregator> aggregators = new HashMap<>(); - + private BeamFnDataOutboundAggregator createRecordingAggregator( + Map<String, List<WindowedValue<String>>> output, Supplier<String> bundleId) { PipelineOptions options = PipelineOptionsFactory.create(); options.as(ExperimentalOptions.class).setExperiments(Arrays.asList("data_buffer_size_limit=0")); - BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, false); - - Map<String, List<WindowedValue<String>>> outputs = ImmutableMap.of("0", output0, "1", output1); - StreamObserver<Elements> observer = + return new BeamFnDataOutboundAggregator( + options, + bundleId, new StreamObserver<Elements>() { @Override public void onNext(Elements elements) { for (Data data : elements.getDataList()) { try { - outputs.get(bundleId.get()).add(WIRE_CODER.decode(data.getData().newInput())); + output.get(bundleId.get()).add(WIRE_CODER.decode(data.getData().newInput())); } catch (IOException e) { throw new RuntimeException("Failed to decode output."); } @@ -140,9 +131,22 @@ public class BeamFnDataWriteRunnerTest { @Override public void onCompleted() {} - }; + }, + false); + } + + @Test + public void testReuseForMultipleBundles() throws Exception { + AtomicReference<String> bundleId = new AtomicReference<>("0"); + String localInputId = "inputPC"; + RunnerApi.PTransform pTransform = + RemoteGrpcPortWrite.writeToPort(localInputId, PORT_SPEC).toPTransform(); - aggregator.prepareForInstruction(bundleId.get(), observer); + List<WindowedValue<String>> output0 = new ArrayList<>(); + List<WindowedValue<String>> output1 = new ArrayList<>(); + Map<ApiServiceDescriptor, BeamFnDataOutboundAggregator> aggregators = new HashMap<>(); + BeamFnDataOutboundAggregator aggregator = + createRecordingAggregator(ImmutableMap.of("0", output0, "1", output1), bundleId::get); aggregators.put(PORT_SPEC.getApiServiceDescriptor(), aggregator); PTransformRunnerFactoryTestContext context = @@ -168,20 +172,18 @@ public class BeamFnDataWriteRunnerTest { FnDataReceiver<Object> pCollectionConsumer = context.getPCollectionConsumer(localInputId); pCollectionConsumer.accept(valueInGlobalWindow("ABC")); pCollectionConsumer.accept(valueInGlobalWindow("DEF")); + assertThat(output0, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF"))); - aggregator.finishInstruction(); output0.clear(); // Process for bundle id 1 bundleId.set("1"); - aggregator.prepareForInstruction(bundleId.get(), observer); pCollectionConsumer.accept(valueInGlobalWindow("GHI")); pCollectionConsumer.accept(valueInGlobalWindow("JKL")); assertThat(output1, contains(valueInGlobalWindow("GHI"), valueInGlobalWindow("JKL"))); - aggregator.finishInstruction(); verifyNoMoreInteractions(mockBeamFnDataClient); } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java index 2aa555e83cd..50a2fec0b5a 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java @@ -839,7 +839,7 @@ public class FnApiDoFnRunnerTest implements Serializable { private Supplier<String> processBundleRequestIdSupplier; public TestBeamFnDataOutboundAggregator(Supplier<String> bundleIdSupplier) { - super(PipelineOptionsFactory.create(), false); + super(PipelineOptionsFactory.create(), bundleIdSupplier, null, false); this.timers = new HashMap<>(); this.dataOutput = new HashMap<>(); this.processBundleRequestIdSupplier = bundleIdSupplier; diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java index 51e49953b40..7b4387738a4 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java @@ -54,7 +54,6 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.DoFn.BundleFinalizer; import org.apache.beam.sdk.util.construction.Timer; import org.apache.beam.sdk.values.WindowedValue; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; import org.joda.time.Instant; /** @@ -75,7 +74,6 @@ public abstract class PTransformRunnerFactoryTestContext @Override public void registerReceiver( String instructionId, - String dataStreamId, List<ApiServiceDescriptor> apiServiceDescriptors, CloseableFnDataReceiver<Elements> receiver) { throw new UnsupportedOperationException("Unexpected call during test."); @@ -83,15 +81,15 @@ public abstract class PTransformRunnerFactoryTestContext @Override public void unregisterReceiver( - String instructionId, - String dataStreamId, - List<ApiServiceDescriptor> apiServiceDescriptors) { + String instructionId, List<ApiServiceDescriptor> apiServiceDescriptors) { throw new UnsupportedOperationException("Unexpected call during test."); } @Override - public StreamObserver<Elements> getOutboundObserver( - ApiServiceDescriptor apiServiceDescriptor, String dataStreamId) { + public BeamFnDataOutboundAggregator createOutboundAggregator( + ApiServiceDescriptor apiServiceDescriptor, + Supplier<String> processBundleRequestIdSupplier, + boolean collectElementsIfNoFlushes) { throw new UnsupportedOperationException("Unexpected call during test."); } 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 c03f8272674..47f85178b0a 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 @@ -37,6 +37,7 @@ import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.Mockito.argThat; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.eq; @@ -1070,20 +1071,28 @@ public class ProcessBundleHandlerTest { dataOutput.add(input.getValue()); })); - Mockito.when(beamFnDataClient.getOutboundObserver(any(), any())) - .thenReturn( - new StreamObserver<Elements>() { - @Override - public void onNext(Elements elements) { - timerOutput.addAll(elements.getTimersList()); - } + Mockito.doAnswer( + (invocation) -> + new BeamFnDataOutboundAggregator( + PipelineOptionsFactory.create(), + invocation.getArgument(1), + new StreamObserver<Elements>() { + @Override + public void onNext(Elements elements) { + for (Timers timer : elements.getTimersList()) { + timerOutput.addAll(elements.getTimersList()); + } + } - @Override - public void onError(Throwable throwable) {} + @Override + public void onError(Throwable throwable) {} - @Override - public void onCompleted() {} - }); + @Override + public void onCompleted() {} + }, + invocation.getArgument(2))) + .when(beamFnDataClient) + .createOutboundAggregator(any(), any(), anyBoolean()); return new ProcessBundleHandler( PipelineOptionsFactory.create(), @@ -1400,7 +1409,7 @@ public class ProcessBundleHandlerTest { (invocation) -> { String instructionId = invocation.getArgument(0, String.class); CloseableFnDataReceiver<BeamFnApi.Elements> data = - invocation.getArgument(3, CloseableFnDataReceiver.class); + invocation.getArgument(2, CloseableFnDataReceiver.class); data.accept( BeamFnApi.Elements.newBuilder() .addData( @@ -1412,7 +1421,7 @@ public class ProcessBundleHandlerTest { return null; }) .when(beamFnDataClient) - .registerReceiver(any(), any(), any(), any()); + .registerReceiver(any(), any(), any()); ProcessBundleHandler handler = new ProcessBundleHandler( @@ -1442,8 +1451,8 @@ public class ProcessBundleHandlerTest { .build()); // Ensure that we unregister during successful processing - verify(beamFnDataClient).registerReceiver(eq("instructionId"), any(), any(), any()); - verify(beamFnDataClient).unregisterReceiver(eq("instructionId"), any(), any()); + verify(beamFnDataClient).registerReceiver(eq("instructionId"), any(), any()); + verify(beamFnDataClient).unregisterReceiver(eq("instructionId"), any()); verifyNoMoreInteractions(beamFnDataClient); } @@ -1466,7 +1475,7 @@ public class ProcessBundleHandlerTest { StringUtf8Coder.of().encode("A", encodedData); String instructionId = invocation.getArgument(0, String.class); CloseableFnDataReceiver<BeamFnApi.Elements> data = - invocation.getArgument(3, CloseableFnDataReceiver.class); + invocation.getArgument(2, CloseableFnDataReceiver.class); data.accept( BeamFnApi.Elements.newBuilder() .addData( @@ -1480,7 +1489,7 @@ public class ProcessBundleHandlerTest { return null; }) .when(beamFnDataClient) - .registerReceiver(any(), any(), any(), any()); + .registerReceiver(any(), any(), any()); ProcessBundleHandler handler = new ProcessBundleHandler( @@ -1517,7 +1526,7 @@ public class ProcessBundleHandlerTest { .build())); // Ensure that we unregister during successful processing - verify(beamFnDataClient).registerReceiver(eq("instructionId"), any(), any(), any()); + verify(beamFnDataClient).registerReceiver(eq("instructionId"), any(), any()); verify(beamFnDataClient).poisonInstructionId(eq("instructionId")); verifyNoMoreInteractions(beamFnDataClient); } 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 9d9efa0b9c4..15f83f2582c 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,8 +23,8 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.empty; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; import java.util.Arrays; import java.util.Collection; @@ -49,7 +49,6 @@ import org.apache.beam.sdk.fn.data.FnDataReceiver; import org.apache.beam.sdk.fn.data.LogicalEndpoint; import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; import org.apache.beam.sdk.fn.test.TestStreams; -import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.values.WindowedValue; @@ -170,6 +169,7 @@ public class BeamFnDataGrpcClientTest { BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient( + PipelineOptionsFactory.create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial()); @@ -183,7 +183,7 @@ public class BeamFnDataGrpcClientTest { Collections.emptyList()); clientFactory.registerReceiver( - INSTRUCTION_ID_A, "", Arrays.asList(apiServiceDescriptor), observerA); + INSTRUCTION_ID_A, Arrays.asList(apiServiceDescriptor), observerA); waitForClientToConnect.await(); outboundServerObserver.get().onNext(ELEMENTS_A_1); @@ -193,7 +193,7 @@ public class BeamFnDataGrpcClientTest { Thread.sleep(100); clientFactory.registerReceiver( - INSTRUCTION_ID_B, "", Arrays.asList(apiServiceDescriptor), observerB); + INSTRUCTION_ID_B, Arrays.asList(apiServiceDescriptor), observerB); // Show that out of order stream completion can occur. observerB.awaitCompletion(); @@ -245,6 +245,7 @@ public class BeamFnDataGrpcClientTest { BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient( + PipelineOptionsFactory.create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial()); @@ -261,7 +262,7 @@ public class BeamFnDataGrpcClientTest { Collections.emptyList()); clientFactory.registerReceiver( - INSTRUCTION_ID_A, "", Arrays.asList(apiServiceDescriptor), observer); + INSTRUCTION_ID_A, Arrays.asList(apiServiceDescriptor), observer); waitForClientToConnect.await(); @@ -269,8 +270,12 @@ public class BeamFnDataGrpcClientTest { outboundServerObserver.get().onNext(ELEMENTS_A_1); outboundServerObserver.get().onNext(ELEMENTS_A_2); - Exception e = assertThrows(Exception.class, observer::awaitCompletion); - assertEquals(exceptionToThrow, e); + try { + observer.awaitCompletion(); + fail("Expected channel to fail"); + } catch (Exception e) { + assertEquals(exceptionToThrow, e); + } // The server should not have received any values assertThat(inboundServerValues, empty()); // The consumer should have only been invoked once @@ -316,6 +321,7 @@ public class BeamFnDataGrpcClientTest { BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient( + PipelineOptionsFactory.create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial()); @@ -341,7 +347,7 @@ public class BeamFnDataGrpcClientTest { }); clientFactory.registerReceiver( - INSTRUCTION_ID_A, "", Arrays.asList(apiServiceDescriptor), observerA); + INSTRUCTION_ID_A, Arrays.asList(apiServiceDescriptor), observerA); waitForClientToConnect.await(); outboundServerObserver.get().onNext(ELEMENTS_B_1); @@ -352,9 +358,11 @@ public class BeamFnDataGrpcClientTest { assertTrue(receivedAElement.await(5, TimeUnit.SECONDS)); clientFactory.poisonInstructionId(INSTRUCTION_ID_A); - // We expect the awaitCompletion to fail due to closing. - // Expected. - assertThrows(Exception.class, future::get); + try { + future.get(); + fail(); // We expect the awaitCompletion to fail due to closing. + } catch (Exception ignored) { + } outboundServerObserver.get().onNext(ELEMENTS_A_2); @@ -396,15 +404,16 @@ public class BeamFnDataGrpcClientTest { ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); - PipelineOptions options = - PipelineOptionsFactory.fromArgs("--experiments=data_buffer_size_limit=20").create(); BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient( + PipelineOptionsFactory.fromArgs( + new String[] {"--experiments=data_buffer_size_limit=20"}) + .create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial()); - BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, false); - aggregator.prepareForInstruction( - INSTRUCTION_ID_A, clientFactory.getOutboundObserver(apiServiceDescriptor, "")); + BeamFnDataOutboundAggregator aggregator = + clientFactory.createOutboundAggregator( + apiServiceDescriptor, () -> INSTRUCTION_ID_A, false); FnDataReceiver<WindowedValue<String>> fnDataReceiver = aggregator.registerOutputDataLocation(TRANSFORM_ID_A, CODER); fnDataReceiver.accept(valueInGlobalWindow("ABC"));
