[FnApi Java] Add support for separate named data streams to provide bundle isolation (#38863)

* [FnApi Java] Add support for separate named data streams to provide bundle isolation.

This is advertised to the runner via a new NAMED_DATA_STREAMS protocol capability. The runner is then free to assign bundles to named data streams as it chooses
to isolate bundle processing from each other. Instead of single data stream
from the sdk, the sdk will create a data stream for each name.  The benefit of
doing so is that the multiplexing currently performed on data stream messages being received allows a slow bundle to fill up buffers and block the shared stream.  With separate named streams, bundles on other data streams have separate grpc flow control from the blocked stream and are not affected.

Author: scwhittle

.github/trigger_files/beam_PostCommit_Java_DataflowV2.json

@@ -1,4 +1,4 @@
 {
-  "modification": 6,
+  "modification": 7,
   "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface"
 }

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_V2.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run!",
-  "modification": 4
+  "modification": 5
 }

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark.json

@@ -1,18 +1,4 @@
 {
-  "https://github.com/apache/beam/pull/34902": "Introducing OutputBuilder",
-  "https://github.com/apache/beam/pull/35213": "Eliminating getPane() in favor of getPaneInfo()",
-  "https://github.com/apache/beam/pull/35177": "Introducing WindowedValueReceiver to runners",
-  "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "https://github.com/apache/beam/pull/31156": "noting that PR #31156 should run this test",
-  "https://github.com/apache/beam/pull/31798": "noting that PR #31798 should run this test",
-  "https://github.com/apache/beam/pull/32546": "noting that PR #32546 should run this test",
-  "https://github.com/apache/beam/pull/33267": "noting that PR #33267 should run this test",
-  "https://github.com/apache/beam/pull/33322": "noting that PR #33322 should run this test",
-  "https://github.com/apache/beam/pull/34123": "noting that PR #34123 should run this test",
-  "https://github.com/apache/beam/pull/34080": "noting that PR #34080 should run this test",
-  "https://github.com/apache/beam/pull/34155": "noting that PR #34155 should run this test",
-  "https://github.com/apache/beam/pull/34560": "noting that PR #34560 should run this test",
-  "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface",
-  "https://github.com/apache/beam/pull/35316": "noting that PR #35316 should run this test",
-  "https://github.com/apache/beam/pull/38011": "noting that PR #38011 should run this test"
+    "comment": "Modify this file in a trivial way to cause this test suite to run",
+    "modification": 1
 }

CHANGES.md

@@ -73,6 +73,7 @@
 
 * (Java) Enabled state tag encoding v2 by default for new Dataflow Streaming Engine jobs. It can be disabled by passing `--experiments=disable_streaming_engine_state_tag_encoding_v2` or `--updateCompatibilityVersion=2.74.0` pipeline option. Note that the tag encoding version cannot change during a job update. Jobs using tag encoding v2 (enabled by default for new jobs on 2.75.0+) cannot be downgraded to Beam versions prior to 2.73.0, as only versions 2.73.0 and later support tag encoding v2. ([#38705](https://github.com/apache/beam/issues/38705)).
 * (Python) Added instrumentation to support off-the-shelf profiling agents when launching Python SDK Harness ([#38853](https://github.com/apache/beam/issues/38853)).
+* (Java) Added support to the FnApi Data stream protocol allowing runners to isolate bundles slowly processing input from other bundles. ([#39001](https://github.com/apache/beam/issues/39001)).
 
 ## Breaking Changes
 

model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto

@@ -120,6 +120,9 @@ 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.
@@ -130,6 +133,9 @@ 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) {}
 }
@@ -416,14 +422,22 @@ message ProcessBundleRequest {
   // at https://s.apache.org/beam-fn-api-control-data-embedding.
   Elements elements = 3;
 
-  // indicates that the runner has no stare for the keys in this bundle
+  // Indicates that the runner has no state 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 {
@@ -834,7 +848,15 @@ message Elements {
 
 // Stable
 service BeamFnData {
-  // Used to send data between harnesses.
+  // 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.
   rpc Data(
       // A stream of data representing input.
       stream Elements)
@@ -900,6 +922,9 @@ 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)
@@ -1295,6 +1320,9 @@ 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)
@@ -1356,6 +1384,8 @@ 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) {}
 }

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto

@@ -1689,6 +1689,10 @@ 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"];
   }
 }
 

runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java

@@ -298,7 +298,7 @@ public ActiveBundle newBundle(
       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(

runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/FnDataService.java

@@ -17,7 +17,6 @@
  */
 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;
@@ -69,5 +68,5 @@ public interface FnDataService {
    * <p>The returned aggregator is not thread safe.
    */
   BeamFnDataOutboundAggregator createOutboundAggregator(
-      Supplier<String> processBundleRequestIdSupplier, boolean collectElementsIfNoFlushes);
+      String processBundleId, boolean collectElementsIfNoFlushes);
 }

runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/GrpcDataService.java

@@ -23,7 +23,6 @@
 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;
@@ -175,13 +174,13 @@ public void unregisterReceiver(String instructionId) {
 
   @Override
   public BeamFnDataOutboundAggregator createOutboundAggregator(
-      Supplier<String> processBundleRequestIdSupplier, boolean collectElementsIfNoFlushes) {
+      String instructionId, boolean collectElementsIfNoFlushes) {
     try {
-      return new BeamFnDataOutboundAggregator(
-          options,
-          processBundleRequestIdSupplier,
-          connectedClient.get(3, TimeUnit.MINUTES).getOutboundObserver(),
-          collectElementsIfNoFlushes);
+      BeamFnDataOutboundAggregator aggregator =
+          new BeamFnDataOutboundAggregator(options, collectElementsIfNoFlushes);
+      aggregator.prepareForInstruction(
+          instructionId, connectedClient.get(3, TimeUnit.MINUTES).getOutboundObserver());
+      return aggregator;
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
       throw new RuntimeException(e);

runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/data/GrpcDataServiceTest.java

@@ -102,7 +102,7 @@ public void testMessageReceivedBySingleClientWhenThereAreMultipleClients() throw
       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);