[Dataflow Streaming] Support to receive multiple work items in a single StreamingGetWorkResponseChunk

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java

@@ -136,6 +136,11 @@ public interface DataflowStreamingPipelineOptions extends PipelineOptions {
 
   void setUseSeparateWindmillHeartbeatStreams(Boolean value);
 
+  @Description("If true, GetWorkStreams will request multiple work items in a response chunk.")
+  boolean getWindmillRequestBatchedGetWorkResponse();
+
+  void setWindmillRequestBatchedGetWorkResponse(boolean value);
+
   @Description("The number of streams to use for GetData requests.")
   @Default.Integer(1)
   int getWindmillGetDataStreamCount();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

@@ -788,7 +788,8 @@ private static GrpcWindmillStreamFactory.Builder createGrpcwindmillStreamFactory
         .setSendKeyedGetDataRequests(
             !options.isEnableStreamingEngine()
                 || DataflowRunner.hasExperiment(
-                    options, "streaming_engine_disable_new_heartbeat_requests"));
+                    options, "streaming_engine_disable_new_heartbeat_requests"))
+        .setRequestBatchedGetWorkResponse(options.getWindmillRequestBatchedGetWorkResponse());
   }
 
   private static JobHeader createJobHeader(DataflowWorkerHarnessOptions options, long clientId) {

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java

@@ -19,6 +19,7 @@
 
 import com.google.auto.value.AutoValue;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import javax.annotation.Nullable;
@@ -43,6 +44,7 @@
  */
 @NotThreadSafe
 final class GetWorkResponseChunkAssembler {
+
   private static final Logger LOG = LoggerFactory.getLogger(GetWorkResponseChunkAssembler.class);
 
   private final GetWorkTimingInfosTracker workTimingInfosTracker;
@@ -61,17 +63,26 @@ final class GetWorkResponseChunkAssembler {
    * Appends the response chunk bytes to the {@link #data }byte buffer. Return the assembled
    * WorkItem if all response chunks for a WorkItem have been received.
    */
-  Optional<AssembledWorkItem> append(Windmill.StreamingGetWorkResponseChunk chunk) {
+  List<AssembledWorkItem> append(Windmill.StreamingGetWorkResponseChunk chunk) {
     if (chunk.hasComputationMetadata()) {
       metadata = ComputationMetadata.fromProto(chunk.getComputationMetadata());
     }
-
-    data = data.concat(chunk.getSerializedWorkItem());
-    bufferedSize += chunk.getSerializedWorkItem().size();
     workTimingInfosTracker.addTimingInfo(chunk.getPerWorkItemTimingInfosList());
 
-    // If the entire WorkItem has been received, assemble the WorkItem.
-    return chunk.getRemainingBytesForWorkItem() == 0 ? flushToWorkItem() : Optional.empty();
+    List<AssembledWorkItem> response = new ArrayList<>();
+    for (int i = 0; i < chunk.getSerializedWorkItemList().size(); i++) {
+      ByteString serializedWorkItem = chunk.getSerializedWorkItemList().get(i);
+      data = data.concat(serializedWorkItem);
+      bufferedSize += serializedWorkItem.size();
+      long remainingSize = 0;
+      if (i == chunk.getSerializedWorkItemList().size() - 1) {
+        remainingSize = chunk.getRemainingBytesForWorkItem();
+      }
+      if (remainingSize == 0) {
+        flushToWorkItem().ifPresent(response::add);
+      }
+    }
+    return response;
   }
 
   /**
@@ -100,6 +111,7 @@ private Optional<AssembledWorkItem> flushToWorkItem() {
 
   @AutoValue
   abstract static class ComputationMetadata {
+
     private static ComputationMetadata fromProto(
         Windmill.ComputationWorkItemMetadata metadataProto) {
       return new AutoValue_GetWorkResponseChunkAssembler_ComputationMetadata(

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java

@@ -61,6 +61,7 @@ final class GrpcDirectGetWorkStream
     extends AbstractWindmillStream<StreamingGetWorkRequest, StreamingGetWorkResponseChunk>
     implements GetWorkStream {
   private static final Logger LOG = LoggerFactory.getLogger(GrpcDirectGetWorkStream.class);
+
   private static final StreamingGetWorkRequest HEALTH_CHECK_REQUEST =
       StreamingGetWorkRequest.newBuilder()
           .setRequestExtension(
@@ -88,6 +89,8 @@ final class GrpcDirectGetWorkStream
    */
   private final ConcurrentMap<Long, GetWorkResponseChunkAssembler> workItemAssemblers;
 
+  private final boolean requestBatchedGetWorkResponse;
+
   private GrpcDirectGetWorkStream(
       String backendWorkerToken,
       Function<
@@ -99,6 +102,7 @@ private GrpcDirectGetWorkStream(
       StreamObserverFactory streamObserverFactory,
       Set<AbstractWindmillStream<?, ?>> streamRegistry,
       int logEveryNStreamFailures,
+      boolean requestBatchedGetWorkResponse,
       ThrottleTimer getWorkThrottleTimer,
       HeartbeatSender heartbeatSender,
       GetDataClient getDataClient,
@@ -127,6 +131,7 @@ private GrpcDirectGetWorkStream(
                 .setItems(requestHeader.getMaxItems())
                 .setBytes(requestHeader.getMaxBytes())
                 .build());
+    this.requestBatchedGetWorkResponse = requestBatchedGetWorkResponse;
   }
 
   static GrpcDirectGetWorkStream create(
@@ -140,6 +145,7 @@ static GrpcDirectGetWorkStream create(
       StreamObserverFactory streamObserverFactory,
       Set<AbstractWindmillStream<?, ?>> streamRegistry,
       int logEveryNStreamFailures,
+      boolean requestBatchedGetWorkResponse,
       ThrottleTimer getWorkThrottleTimer,
       HeartbeatSender heartbeatSender,
       GetDataClient getDataClient,
@@ -153,6 +159,7 @@ static GrpcDirectGetWorkStream create(
         streamObserverFactory,
         streamRegistry,
         logEveryNStreamFailures,
+        requestBatchedGetWorkResponse,
         getWorkThrottleTimer,
         heartbeatSender,
         getDataClient,
@@ -209,6 +216,7 @@ protected synchronized void onNewStream() throws WindmillStreamShutdownException
                     .setMaxItems(initialGetWorkBudget.items())
                     .setMaxBytes(initialGetWorkBudget.bytes())
                     .build())
+            .setSupportsMultipleWorkItemsInChunk(requestBatchedGetWorkResponse)
             .build();
     lastRequest.set(request);
     budgetTracker.recordBudgetRequested(initialGetWorkBudget);
@@ -243,7 +251,7 @@ protected void onResponse(StreamingGetWorkResponseChunk chunk) {
     workItemAssemblers
         .computeIfAbsent(chunk.getStreamId(), unused -> new GetWorkResponseChunkAssembler())
         .append(chunk)
-        .ifPresent(this::consumeAssembledWorkItem);
+        .forEach(this::consumeAssembledWorkItem);
   }
 
   private void consumeAssembledWorkItem(AssembledWorkItem assembledWorkItem) {

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java

@@ -45,13 +45,19 @@ final class GrpcGetWorkStream
     implements GetWorkStream {
 
   private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkStream.class);
+  private static final StreamingGetWorkRequest HEALTH_CHECK =
+      StreamingGetWorkRequest.newBuilder()
+          .setRequestExtension(
+              StreamingGetWorkRequestExtension.newBuilder().setMaxItems(0).setMaxBytes(0).build())
+          .build();
 
   private final GetWorkRequest request;
   private final WorkItemReceiver receiver;
   private final ThrottleTimer getWorkThrottleTimer;
   private final Map<Long, GetWorkResponseChunkAssembler> workItemAssemblers;
   private final AtomicLong inflightMessages;
   private final AtomicLong inflightBytes;
+  private final boolean requestBatchedGetWorkResponse;
 
   private GrpcGetWorkStream(
       String backendWorkerToken,
@@ -64,6 +70,7 @@ private GrpcGetWorkStream(
       StreamObserverFactory streamObserverFactory,
       Set<AbstractWindmillStream<?, ?>> streamRegistry,
       int logEveryNStreamFailures,
+      boolean requestBatchedGetWorkResponse,
       ThrottleTimer getWorkThrottleTimer,
       WorkItemReceiver receiver) {
     super(
@@ -81,6 +88,7 @@ private GrpcGetWorkStream(
     this.workItemAssemblers = new ConcurrentHashMap<>();
     this.inflightMessages = new AtomicLong();
     this.inflightBytes = new AtomicLong();
+    this.requestBatchedGetWorkResponse = requestBatchedGetWorkResponse;
   }
 
   public static GrpcGetWorkStream create(
@@ -94,6 +102,7 @@ public static GrpcGetWorkStream create(
       StreamObserverFactory streamObserverFactory,
       Set<AbstractWindmillStream<?, ?>> streamRegistry,
       int logEveryNStreamFailures,
+      boolean requestBatchedGetWorkResponse,
       ThrottleTimer getWorkThrottleTimer,
       WorkItemReceiver receiver) {
     return new GrpcGetWorkStream(
@@ -104,6 +113,7 @@ public static GrpcGetWorkStream create(
         streamObserverFactory,
         streamRegistry,
         logEveryNStreamFailures,
+        requestBatchedGetWorkResponse,
         getWorkThrottleTimer,
         receiver);
   }
@@ -132,7 +142,11 @@ protected synchronized void onNewStream() throws WindmillStreamShutdownException
     workItemAssemblers.clear();
     inflightMessages.set(request.getMaxItems());
     inflightBytes.set(request.getMaxBytes());
-    trySend(StreamingGetWorkRequest.newBuilder().setRequest(request).build());
+    trySend(
+        StreamingGetWorkRequest.newBuilder()
+            .setSupportsMultipleWorkItemsInChunk(requestBatchedGetWorkResponse)
+            .setRequest(request)
+            .build());
   }
 
   @Override
@@ -153,11 +167,7 @@ public void appendSpecificHtml(PrintWriter writer) {
 
   @Override
   public void sendHealthCheck() throws WindmillStreamShutdownException {
-    trySend(
-        StreamingGetWorkRequest.newBuilder()
-            .setRequestExtension(
-                StreamingGetWorkRequestExtension.newBuilder().setMaxItems(0).setMaxBytes(0).build())
-            .build());
+    trySend(HEALTH_CHECK);
   }
 
   @Override
@@ -166,7 +176,7 @@ protected void onResponse(StreamingGetWorkResponseChunk chunk) {
     workItemAssemblers
         .computeIfAbsent(chunk.getStreamId(), unused -> new GetWorkResponseChunkAssembler())
         .append(chunk)
-        .ifPresent(this::consumeAssembledWorkItem);
+        .forEach(this::consumeAssembledWorkItem);
   }
 
   private void consumeAssembledWorkItem(AssembledWorkItem assembledWorkItem) {

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServer.java

@@ -185,6 +185,8 @@ static GrpcWindmillServer newTestInstance(
             .setSendKeyedGetDataRequests(sendKeyedGetDataRequests)
             .setHealthCheckIntervalMillis(
                 testOptions.getWindmillServiceStreamingRpcHealthCheckPeriodMs())
+            .setRequestBatchedGetWorkResponse(
+                testOptions.getWindmillRequestBatchedGetWorkResponse())
             .build();
 
     return new GrpcWindmillServer(testOptions, windmillStreamFactory, dispatcherClient);

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillStreamFactory.java

@@ -91,6 +91,7 @@ public class GrpcWindmillStreamFactory implements StatusDataProvider {
   // If true, then active work refreshes will be sent as KeyedGetDataRequests. Otherwise, use the
   // newer ComputationHeartbeatRequests.
   private final boolean sendKeyedGetDataRequests;
+  private final boolean requestBatchedGetWorkResponse;
   private final Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses;
 
   private GrpcWindmillStreamFactory(
@@ -99,6 +100,7 @@ private GrpcWindmillStreamFactory(
       int streamingRpcBatchLimit,
       int windmillMessagesBetweenIsReadyChecks,
       boolean sendKeyedGetDataRequests,
+      boolean requestBatchedGetWorkResponse,
       Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses,
       Supplier<Duration> maxBackOffSupplier) {
     this.jobHeader = jobHeader;
@@ -115,6 +117,7 @@ private GrpcWindmillStreamFactory(
                     .backoff());
     this.streamRegistry = ConcurrentHashMap.newKeySet();
     this.sendKeyedGetDataRequests = sendKeyedGetDataRequests;
+    this.requestBatchedGetWorkResponse = requestBatchedGetWorkResponse;
     this.processHeartbeatResponses = processHeartbeatResponses;
     this.streamIdGenerator = new AtomicLong();
   }
@@ -126,6 +129,7 @@ static GrpcWindmillStreamFactory create(
       int streamingRpcBatchLimit,
       int windmillMessagesBetweenIsReadyChecks,
       boolean sendKeyedGetDataRequests,
+      boolean requestBatchedGetWorkResponse,
       Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses,
       Supplier<Duration> maxBackOffSupplier,
       int healthCheckIntervalMillis) {
@@ -136,6 +140,7 @@ static GrpcWindmillStreamFactory create(
             streamingRpcBatchLimit,
             windmillMessagesBetweenIsReadyChecks,
             sendKeyedGetDataRequests,
+            requestBatchedGetWorkResponse,
             processHeartbeatResponses,
             maxBackOffSupplier);
 
@@ -174,6 +179,7 @@ public static GrpcWindmillStreamFactory.Builder of(JobHeader jobHeader) {
         .setStreamingRpcBatchLimit(DEFAULT_STREAMING_RPC_BATCH_LIMIT)
         .setHealthCheckIntervalMillis(NO_HEALTH_CHECKS)
         .setSendKeyedGetDataRequests(true)
+        .setRequestBatchedGetWorkResponse(false)
         .setProcessHeartbeatResponses(ignored -> {});
   }
 
@@ -209,6 +215,7 @@ public GetWorkStream createGetWorkStream(
         newStreamObserverFactory(),
         streamRegistry,
         logEveryNStreamFailures,
+        requestBatchedGetWorkResponse,
         getWorkThrottleTimer,
         processWorkItem);
   }
@@ -229,6 +236,7 @@ public GetWorkStream createDirectGetWorkStream(
         newStreamObserverFactory(),
         streamRegistry,
         logEveryNStreamFailures,
+        requestBatchedGetWorkResponse,
         getWorkThrottleTimer,
         heartbeatSender,
         getDataClient,
@@ -356,6 +364,8 @@ Builder setProcessHeartbeatResponses(
 
     Builder setHealthCheckIntervalMillis(int healthCheckIntervalMillis);
 
+    Builder setRequestBatchedGetWorkResponse(boolean enabled);
+
     GrpcWindmillStreamFactory build();
   }
 }

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java

@@ -242,7 +242,7 @@ public void testStreamsStartCorrectly() throws InterruptedException {
   @Test
   public void testOnNewWorkerMetadata_correctlyRemovesStaleWindmillServers()
       throws InterruptedException {
-    TestGetWorkBudgetDistributor getWorkBudgetDistributor = spy(new TestGetWorkBudgetDistributor());
+    GetWorkBudgetDistributor getWorkBudgetDistributor = mock(GetWorkBudgetDistributor.class);
     fanOutStreamingEngineWorkProvider =
         newFanOutStreamingEngineWorkerHarness(
             GetWorkBudget.builder().setItems(1).setBytes(1).build(),