[Dataflow Streaming] [Multi Key] MultiKey failure handling + Integration

.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":  1,
+  "modification":  2,
 }

.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/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowExecutionContext.java

@@ -150,6 +150,10 @@ boolean isSinkFullHintSet() {
     //               the state size might grow unbounded.
   }
 
+  protected final long getBytesSinked() {
+    return bytesSinked;
+  }
+
   /**
    * Sets a flag to indicate that a sink has enough data written to it. This hint is read by
    * upstream producers to stop producing if they can. Mainly used in streaming.

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

@@ -181,7 +181,7 @@ public final class StreamingDataflowWorker {
       "windmill_bounded_queue_executor_use_fair_monitor";
   // Don't use. Experiment guarding multi key bundles. The feature is work in progress and
   // incomplete.
-  private static final String UNSTABLE_ENABLE_MULTI_KEY_BUNDLE = "unstable_enable_multi_key_bundle";
+  public static final String UNSTABLE_ENABLE_MULTI_KEY_BUNDLE = "unstable_enable_multi_key_bundle";
 
   private final WindmillStateCache stateCache;
   private AtomicReference<StreamingWorkerStatusPages> statusPages = new AtomicReference<>();
@@ -257,6 +257,7 @@ private StreamingDataflowWorker(
     this.streamingWorkScheduler =
         StreamingWorkScheduler.create(
             options,
+            DataflowRunner.hasExperiment(options, UNSTABLE_ENABLE_MULTI_KEY_BUNDLE),
             clock,
             readerCache,
             mapTaskExecutorFactory,
@@ -1198,9 +1199,14 @@ private void onCompleteCommit(CompleteCommit completeCommit) {
     computationStateCache
         .getIfPresent(completeCommit.computationId())
         .ifPresent(
-            state ->
+            state -> {
+              if (completeCommit.retryableFailure()) {
+                state.reExecuteActiveWork(completeCommit.shardedKey(), completeCommit.workId());
+              } else {
                 state.completeWorkAndScheduleNextWorkForKey(
-                    completeCommit.shardedKey(), completeCommit.workId()));
+                    completeCommit.shardedKey(), completeCommit.workId());
+              }
+            });
   }
 
   @AutoValue