Merge pull request #38168 from acrites/cherrypick_pr_37954

Cherry-pick PR #37954 into 2.73 release

Author: Amar3tto

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json

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

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming_Engine.json

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

runners/google-cloud-dataflow-java/build.gradle

@@ -491,8 +491,6 @@ task validatesRunnerStreaming {
   description "Validates Dataflow runner forcing streaming mode"
   dependsOn(createLegacyWorkerValidatesRunnerTest(validatesRunnerStreamingConfig + [
     name: 'validatesRunnerLegacyWorkerTestStreaming',
-    // Streaming appliance currently fails bundle finalizer tests.
-    excludedCategories: validatesRunnerStreamingConfig.excludedCategories + [ 'org.apache.beam.sdk.testing.UsesBundleFinalizer', ],
   ]))
 }
 

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

@@ -401,12 +401,14 @@ private StreamingWorkerHarnessFactoryOutput createFanOutStreamingEngineWorkerHar
                 watermarks,
                 processingContext,
                 drainMode,
+                appliedFinalizeIds,
                 getWorkStreamLatencies) ->
                 checkNotNull(computationStateCache)
                     .get(processingContext.computationId())
                     .ifPresent(
                         computationState -> {
                           memoryMonitor.waitForResources("GetWork");
+                          streamingWorkScheduler.queueAppliedFinalizeIds(appliedFinalizeIds);
                           streamingWorkScheduler.scheduleWork(
                               computationState,
                               workItem,

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java

@@ -158,12 +158,16 @@ private void streamingEngineDispatchLoop(
                   drainMode,
                   workItem,
                   serializedWorkItemSize,
+                  appliedFinalizeIds,
                   getWorkStreamLatencies) ->
                   computationStateFetcher
                       .apply(computationId)
                       .ifPresent(
                           computationState -> {
                             waitForResources.run();
+                            if (!appliedFinalizeIds.isEmpty()) {
+                              streamingWorkScheduler.queueAppliedFinalizeIds(appliedFinalizeIds);
+                            }
                             streamingWorkScheduler.scheduleWork(
                                 computationState,
                                 workItem,
@@ -214,6 +218,12 @@ private void applianceDispatchLoop(Supplier<Windmill.GetWorkResponse> getWorkFn)
         sleepUninterruptibly(backoff, TimeUnit.MILLISECONDS);
         backoff = Math.min(1000, backoff * 2);
       } while (isRunning.get());
+      ImmutableList<Long> appliedFinalizeIds =
+          ImmutableList.copyOf(
+              Preconditions.checkNotNull(workResponse).getAppliedFinalizeIdsList());
+      if (!appliedFinalizeIds.isEmpty()) {
+        streamingWorkScheduler.queueAppliedFinalizeIds(appliedFinalizeIds);
+      }
       for (Windmill.ComputationWorkItems computationWork :
           Preconditions.checkNotNull(workResponse).getWorkList()) {
         String computationId = computationWork.getComputationId();

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

@@ -100,12 +100,12 @@ List<AssembledWorkItem> append(Windmill.StreamingGetWorkResponseChunk chunk) {
   private Optional<AssembledWorkItem> flushToWorkItem() {
     try {
       workItemBuilder.mergeFrom(data);
-      workItemBuilder.addAllAppliedFinalizeIds(appliedFinalizeIds);
       return Optional.of(
           AssembledWorkItem.create(
               workItemBuilder.build(),
               Preconditions.checkNotNull(metadata),
               workTimingInfosTracker.getLatencyAttributions(),
+              ImmutableList.copyOf(appliedFinalizeIds),
               bufferedSize));
     } catch (IOException e) {
       LOG.error("Failed to parse work item from stream: ", e);
@@ -149,9 +149,10 @@ private static AssembledWorkItem create(
         WorkItem workItem,
         ComputationMetadata computationMetadata,
         ImmutableList<LatencyAttribution> latencyAttributions,
+        ImmutableList<Long> appliedFinalizeIds,
         long size) {
       return new AutoValue_GetWorkResponseChunkAssembler_AssembledWorkItem(
-          workItem, computationMetadata, latencyAttributions, size);
+          workItem, computationMetadata, latencyAttributions, appliedFinalizeIds, size);
     }
 
     abstract WorkItem workItem();
@@ -160,6 +161,8 @@ private static AssembledWorkItem create(
 
     abstract ImmutableList<LatencyAttribution> latencyAttributions();
 
+    abstract ImmutableList<Long> appliedFinalizeIds();
+
     abstract long bufferedSize();
   }
 }

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

@@ -282,6 +282,7 @@ private void consumeAssembledWorkItem(AssembledWorkItem assembledWorkItem) {
         createWatermarks(workItem, metadata),
         createProcessingContext(metadata.computationId()),
         metadata.drainMode(),
+        assembledWorkItem.appliedFinalizeIds(),
         assembledWorkItem.latencyAttributions());
     budgetTracker.recordBudgetReceived(assembledWorkItem.bufferedSize());
     GetWorkBudget extension = budgetTracker.computeBudgetExtension();

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

@@ -206,6 +206,7 @@ private void consumeAssembledWorkItem(AssembledWorkItem assembledWorkItem) {
         assembledWorkItem.computationMetadata().drainMode(),
         assembledWorkItem.workItem(),
         assembledWorkItem.bufferedSize(),
+        assembledWorkItem.appliedFinalizeIds(),
         assembledWorkItem.latencyAttributions());
 
     // Record the fact that there are now fewer outstanding messages and bytes on the stream.

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemReceiver.java

@@ -33,5 +33,6 @@ void receiveWork(
       boolean drainMode,
       Windmill.WorkItem workItem,
       long serializedWorkItemSize,
+      ImmutableList<Long> appliedFinalizeIds,
       ImmutableList<LatencyAttribution> getWorkStreamLatencies);
 }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemScheduler.java

@@ -36,6 +36,7 @@ public interface WorkItemScheduler {
    * @param watermarks processing watermarks for the workItem.
    * @param processingContext for processing the workItem.
    * @param drainMode is job is draining.
+   * @param appliedFinalizeIds Any applied finalize ids that should have their callbacks run.
    * @param getWorkStreamLatencies Latencies per processing stage for the WorkItem for reporting
    *     back to Streaming Engine backend.
    */
@@ -45,5 +46,6 @@ void scheduleWork(
       Watermarks watermarks,
       Work.ProcessingContext processingContext,
       boolean drainMode,
+      ImmutableList<Long> appliedFinalizeIds,
       ImmutableList<LatencyAttribution> getWorkStreamLatencies);
 }