[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

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

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

@@ -789,7 +789,7 @@ private static GrpcWindmillStreamFactory.Builder createGrpcwindmillStreamFactory
             !options.isEnableStreamingEngine()
                 || DataflowRunner.hasExperiment(
                     options, "streaming_engine_disable_new_heartbeat_requests"))
-        .setMultipleItemsInGetWorkResponse(options.getWindmillMultipleItemsInGetWorkResponse());
+        .setBatchedGetWorkResponse(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/GrpcDirectGetWorkStream.java

@@ -62,6 +62,15 @@ final class GrpcDirectGetWorkStream
     implements GetWorkStream {
   private static final Logger LOG = LoggerFactory.getLogger(GrpcDirectGetWorkStream.class);
 
+  private static final StreamingGetWorkRequest HEALTH_CHECK_REQUEST =
+      StreamingGetWorkRequest.newBuilder()
+          .setRequestExtension(
+              Windmill.StreamingGetWorkRequestExtension.newBuilder()
+                  .setMaxItems(0)
+                  .setMaxBytes(0)
+                  .build())
+          .build();
+
   private final GetWorkBudgetTracker budgetTracker;
   private final GetWorkRequest requestHeader;
   private final WorkItemScheduler workItemScheduler;
@@ -80,7 +89,7 @@ final class GrpcDirectGetWorkStream
    */
   private final ConcurrentMap<Long, GetWorkResponseChunkAssembler> workItemAssemblers;
 
-  private final boolean multipleItemsInGetWorkResponse;
+  private final boolean batchedGetWorkResponse;
 
   private GrpcDirectGetWorkStream(
       String backendWorkerToken,
@@ -93,7 +102,7 @@ private GrpcDirectGetWorkStream(
       StreamObserverFactory streamObserverFactory,
       Set<AbstractWindmillStream<?, ?>> streamRegistry,
       int logEveryNStreamFailures,
-      boolean multipleItemsInGetWorkResponse,
+      boolean batchedGetWorkResponse,
       ThrottleTimer getWorkThrottleTimer,
       HeartbeatSender heartbeatSender,
       GetDataClient getDataClient,
@@ -122,7 +131,7 @@ private GrpcDirectGetWorkStream(
                 .setItems(requestHeader.getMaxItems())
                 .setBytes(requestHeader.getMaxBytes())
                 .build());
-    this.multipleItemsInGetWorkResponse = multipleItemsInGetWorkResponse;
+    this.batchedGetWorkResponse = batchedGetWorkResponse;
   }
 
   static GrpcDirectGetWorkStream create(
@@ -182,7 +191,6 @@ private void maybeSendRequestExtension(GetWorkBudget extension) {
                         Windmill.StreamingGetWorkRequestExtension.newBuilder()
                             .setMaxItems(extension.items())
                             .setMaxBytes(extension.bytes()))
-                    .setSupportsMultipleWorkItemsInChunk(multipleItemsInGetWorkResponse)
                     .build();
             lastRequest.set(request);
             budgetTracker.recordBudgetRequested(extension);
@@ -208,7 +216,7 @@ protected synchronized void onNewStream() throws WindmillStreamShutdownException
                     .setMaxItems(initialGetWorkBudget.items())
                     .setMaxBytes(initialGetWorkBudget.bytes())
                     .build())
-            .setSupportsMultipleWorkItemsInChunk(multipleItemsInGetWorkResponse)
+            .setSupportsMultipleWorkItemsInChunk(batchedGetWorkResponse)
             .build();
     lastRequest.set(request);
     budgetTracker.recordBudgetRequested(initialGetWorkBudget);
@@ -231,15 +239,7 @@ public void appendSpecificHtml(PrintWriter writer) {
 
   @Override
   public void sendHealthCheck() throws WindmillStreamShutdownException {
-    trySend(
-        StreamingGetWorkRequest.newBuilder()
-            .setRequestExtension(
-                Windmill.StreamingGetWorkRequestExtension.newBuilder()
-                    .setMaxItems(0)
-                    .setMaxBytes(0)
-                    .build())
-            .setSupportsMultipleWorkItemsInChunk(multipleItemsInGetWorkResponse)
-            .build());
+    trySend(HEALTH_CHECK_REQUEST);
   }
 
   @Override

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

@@ -45,14 +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 multipleItemsInGetWorkResponse;
+  private final boolean batchedGetWorkResponse;
 
   private GrpcGetWorkStream(
       String backendWorkerToken,
@@ -65,7 +70,7 @@ private GrpcGetWorkStream(
       StreamObserverFactory streamObserverFactory,
       Set<AbstractWindmillStream<?, ?>> streamRegistry,
       int logEveryNStreamFailures,
-      boolean multipleItemsInGetWorkResponse,
+      boolean batchedGetWorkResponse,
       ThrottleTimer getWorkThrottleTimer,
       WorkItemReceiver receiver) {
     super(
@@ -83,7 +88,7 @@ private GrpcGetWorkStream(
     this.workItemAssemblers = new ConcurrentHashMap<>();
     this.inflightMessages = new AtomicLong();
     this.inflightBytes = new AtomicLong();
-    this.multipleItemsInGetWorkResponse = multipleItemsInGetWorkResponse;
+    this.batchedGetWorkResponse = batchedGetWorkResponse;
   }
 
   public static GrpcGetWorkStream create(
@@ -120,7 +125,6 @@ private void sendRequestExtension(long moreItems, long moreBytes) {
                 StreamingGetWorkRequestExtension.newBuilder()
                     .setMaxItems(moreItems)
                     .setMaxBytes(moreBytes))
-            .setSupportsMultipleWorkItemsInChunk(multipleItemsInGetWorkResponse)
             .build();
 
     executeSafely(
@@ -140,7 +144,7 @@ protected synchronized void onNewStream() throws WindmillStreamShutdownException
     inflightBytes.set(request.getMaxBytes());
     trySend(
         StreamingGetWorkRequest.newBuilder()
-            .setSupportsMultipleWorkItemsInChunk(multipleItemsInGetWorkResponse)
+            .setSupportsMultipleWorkItemsInChunk(batchedGetWorkResponse)
             .setRequest(request)
             .build());
   }
@@ -163,12 +167,7 @@ public void appendSpecificHtml(PrintWriter writer) {
 
   @Override
   public void sendHealthCheck() throws WindmillStreamShutdownException {
-    trySend(
-        StreamingGetWorkRequest.newBuilder()
-            .setRequestExtension(
-                StreamingGetWorkRequestExtension.newBuilder().setMaxItems(0).setMaxBytes(0).build())
-            .setSupportsMultipleWorkItemsInChunk(multipleItemsInGetWorkResponse)
-            .build());
+    trySend(HEALTH_CHECK);
   }
 
   @Override

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

@@ -185,8 +185,7 @@ static GrpcWindmillServer newTestInstance(
             .setSendKeyedGetDataRequests(sendKeyedGetDataRequests)
             .setHealthCheckIntervalMillis(
                 testOptions.getWindmillServiceStreamingRpcHealthCheckPeriodMs())
-            .setMultipleItemsInGetWorkResponse(
-                testOptions.getWindmillMultipleItemsInGetWorkResponse())
+            .setBatchedGetWorkResponse(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,7 +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 multipleItemsInGetWorkResponse;
+  private final boolean batchedGetWorkResponse;
   private final Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses;
 
   private GrpcWindmillStreamFactory(
@@ -100,7 +100,7 @@ private GrpcWindmillStreamFactory(
       int streamingRpcBatchLimit,
       int windmillMessagesBetweenIsReadyChecks,
       boolean sendKeyedGetDataRequests,
-      boolean multipleItemsInGetWorkResponse,
+      boolean batchedGetWorkResponse,
       Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses,
       Supplier<Duration> maxBackOffSupplier) {
     this.jobHeader = jobHeader;
@@ -117,7 +117,7 @@ private GrpcWindmillStreamFactory(
                     .backoff());
     this.streamRegistry = ConcurrentHashMap.newKeySet();
     this.sendKeyedGetDataRequests = sendKeyedGetDataRequests;
-    this.multipleItemsInGetWorkResponse = multipleItemsInGetWorkResponse;
+    this.batchedGetWorkResponse = batchedGetWorkResponse;
     this.processHeartbeatResponses = processHeartbeatResponses;
     this.streamIdGenerator = new AtomicLong();
   }
@@ -129,7 +129,7 @@ static GrpcWindmillStreamFactory create(
       int streamingRpcBatchLimit,
       int windmillMessagesBetweenIsReadyChecks,
       boolean sendKeyedGetDataRequests,
-      boolean multipleItemsInGetWorkResponse,
+      boolean batchedGetWorkResponse,
       Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses,
       Supplier<Duration> maxBackOffSupplier,
       int healthCheckIntervalMillis) {
@@ -140,7 +140,7 @@ static GrpcWindmillStreamFactory create(
             streamingRpcBatchLimit,
             windmillMessagesBetweenIsReadyChecks,
             sendKeyedGetDataRequests,
-            multipleItemsInGetWorkResponse,
+            batchedGetWorkResponse,
             processHeartbeatResponses,
             maxBackOffSupplier);
 
@@ -179,7 +179,7 @@ public static GrpcWindmillStreamFactory.Builder of(JobHeader jobHeader) {
         .setStreamingRpcBatchLimit(DEFAULT_STREAMING_RPC_BATCH_LIMIT)
         .setHealthCheckIntervalMillis(NO_HEALTH_CHECKS)
         .setSendKeyedGetDataRequests(true)
-        .setMultipleItemsInGetWorkResponse(false)
+        .setBatchedGetWorkResponse(false)
         .setProcessHeartbeatResponses(ignored -> {});
   }
 
@@ -215,7 +215,7 @@ public GetWorkStream createGetWorkStream(
         newStreamObserverFactory(),
         streamRegistry,
         logEveryNStreamFailures,
-        multipleItemsInGetWorkResponse,
+        batchedGetWorkResponse,
         getWorkThrottleTimer,
         processWorkItem);
   }
@@ -236,7 +236,7 @@ public GetWorkStream createDirectGetWorkStream(
         newStreamObserverFactory(),
         streamRegistry,
         logEveryNStreamFailures,
-        multipleItemsInGetWorkResponse,
+        batchedGetWorkResponse,
         getWorkThrottleTimer,
         heartbeatSender,
         getDataClient,
@@ -364,7 +364,7 @@ Builder setProcessHeartbeatResponses(
 
     Builder setHealthCheckIntervalMillis(int healthCheckIntervalMillis);
 
-    Builder setMultipleItemsInGetWorkResponse(boolean enabled);
+    Builder setBatchedGetWorkResponse(boolean enabled);
 
     GrpcWindmillStreamFactory build();
   }