[wip] Multikey commit validation

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/MultiKeyBundleOptions.java

@@ -0,0 +1,138 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker;
+
+import com.google.auto.value.AutoValue;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.options.ExperimentalOptions;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@AutoValue
+@Internal
+public abstract class MultiKeyBundleOptions {
+  // TODO: consider moving this to be part of PipelineOptions after the feature is stable
+
+  private static final Logger LOG = LoggerFactory.getLogger(MultiKeyBundleOptions.class);
+
+  // Don't use. Experiment guarding multi key bundles. The feature is work in progress and
+  // incomplete.
+  public static final String UNSTABLE_ENABLE_MULTI_KEY_BUNDLE = "unstable_enable_multi_key_bundle";
+
+  private static final String WINDMILL_MAX_KEY_GROUP_BATCH_SIZE =
+      "windmill_max_key_group_batch_size";
+  private static final String WINDMILL_MAX_KEY_GROUP_BATCH_TIME_MS =
+      "windmill_max_key_group_batch_time_ms";
+  private static final String WINDMILL_MAX_KEY_GROUP_BATCH_SINK_BYTES =
+      "windmill_max_key_group_batch_sink_bytes";
+
+  public abstract int maxKeyGroupBatchSize();
+
+  public abstract long maxKeyGroupBatchTimeNanos();
+
+  public abstract boolean multiKeyBundleEnabled();
+
+  public abstract long maxKeyGroupBatchSinkBytes();
+
+  public static Builder builder() {
+    return new AutoValue_MultiKeyBundleOptions.Builder();
+  }
+
+  public static MultiKeyBundleOptions fromOptions(PipelineOptions options) {
+    int maxKeyGroupBatchSize =
+        tryParseInt(
+            ExperimentalOptions.getExperimentValue(options, WINDMILL_MAX_KEY_GROUP_BATCH_SIZE),
+            100,
+            WINDMILL_MAX_KEY_GROUP_BATCH_SIZE);
+
+    long batchTimeMs =
+        tryParseLong(
+            ExperimentalOptions.getExperimentValue(options, WINDMILL_MAX_KEY_GROUP_BATCH_TIME_MS),
+            100,
+            WINDMILL_MAX_KEY_GROUP_BATCH_TIME_MS);
+
+    boolean multiKeyBundleEnabled =
+        ExperimentalOptions.hasExperiment(options, UNSTABLE_ENABLE_MULTI_KEY_BUNDLE);
+
+    long maxKeyGroupBatchSinkBytes =
+        tryParseLong(
+            ExperimentalOptions.getExperimentValue(
+                options, WINDMILL_MAX_KEY_GROUP_BATCH_SINK_BYTES),
+            StreamingDataflowWorker.MAX_SINK_BYTES,
+            WINDMILL_MAX_KEY_GROUP_BATCH_SINK_BYTES);
+
+    return builder()
+        .setMaxKeyGroupBatchSize(maxKeyGroupBatchSize)
+        .setMaxKeyGroupBatchTimeNanos(TimeUnit.MILLISECONDS.toNanos(batchTimeMs))
+        .setMultiKeyBundleEnabled(multiKeyBundleEnabled)
+        .setMaxKeyGroupBatchSinkBytes(maxKeyGroupBatchSinkBytes)
+        .build();
+  }
+
+  private static int tryParseInt(@Nullable String value, int defaultValue, String experimentName) {
+    if (value == null) {
+      return defaultValue;
+    }
+    try {
+      return Integer.parseInt(value);
+    } catch (NumberFormatException e) {
+      LOG.warn(
+          "Failed to parse experiment {} value '{}' as integer, falling back to default: {}",
+          experimentName,
+          value,
+          defaultValue,
+          e);
+      return defaultValue;
+    }
+  }
+
+  private static long tryParseLong(
+      @Nullable String value, long defaultValue, String experimentName) {
+    if (value == null) {
+      return defaultValue;
+    }
+    try {
+      return Long.parseLong(value);
+    } catch (NumberFormatException e) {
+      LOG.warn(
+          "Failed to parse experiment {} value '{}' as long, falling back to default: {}",
+          experimentName,
+          value,
+          defaultValue,
+          e);
+      return defaultValue;
+    }
+  }
+
+  @AutoValue.Builder
+  public abstract static class Builder {
+
+    public abstract Builder setMaxKeyGroupBatchSize(int size);
+
+    public abstract Builder setMaxKeyGroupBatchTimeNanos(long nanos);
+
+    public abstract Builder setMultiKeyBundleEnabled(boolean enabled);
+
+    public abstract Builder setMaxKeyGroupBatchSinkBytes(long bytes);
+
+    public abstract MultiKeyBundleOptions build();
+  }
+}

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

@@ -179,9 +179,6 @@ public final class StreamingDataflowWorker {
   // Experiment make the monitor within BoundedQueueExecutor fair
   public static final String BOUNDED_QUEUE_EXECUTOR_USE_FAIR_MONITOR_EXPERIMENT =
       "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";
 
   private final WindmillStateCache stateCache;
   private AtomicReference<StreamingWorkerStatusPages> statusPages = new AtomicReference<>();
@@ -211,7 +208,7 @@ public final class StreamingDataflowWorker {
   private StreamingDataflowWorker(
       WindmillServerStub windmillServer,
       long clientId,
-      ComputationConfig.Fetcher configFetcher,
+      Fetcher configFetcher,
       ComputationStateCache computationStateCache,
       WindmillStateCache windmillStateCache,
       BoundedQueueExecutor workUnitExecutor,
@@ -228,7 +225,8 @@ private StreamingDataflowWorker(
       GrpcWindmillStreamFactory windmillStreamFactory,
       ScheduledExecutorService activeWorkRefreshExecutorFn,
       ConcurrentMap<String, StageInfo> stageInfoMap,
-      @Nullable GrpcDispatcherClient dispatcherClient) {
+      @Nullable GrpcDispatcherClient dispatcherClient,
+      MultiKeyBundleOptions multiKeyBundleOptions) {
     // Register standard file systems.
     FileSystems.setDefaultPipelineOptions(options);
     this.configFetcher = configFetcher;
@@ -257,6 +255,7 @@ private StreamingDataflowWorker(
     this.streamingWorkScheduler =
         StreamingWorkScheduler.create(
             options,
+            multiKeyBundleOptions,
             clock,
             readerCache,
             mapTaskExecutorFactory,
@@ -627,7 +626,8 @@ public static StreamingDataflowWorker fromOptions(DataflowWorkerHarnessOptions o
     ConcurrentMap<String, StageInfo> stageInfo = new ConcurrentHashMap<>();
     StreamingCounters streamingCounters = StreamingCounters.create();
     WorkUnitClient dataflowServiceClient = new DataflowWorkUnitClient(options, LOG);
-    BoundedQueueExecutor workExecutor = createWorkUnitExecutor(options);
+    MultiKeyBundleOptions multiKeyBundleOptions = MultiKeyBundleOptions.fromOptions(options);
+    BoundedQueueExecutor workExecutor = createWorkUnitExecutor(options, multiKeyBundleOptions);
     ScheduledExecutorService commitFinalizerCleanupExecutor =
         Executors.newScheduledThreadPool(
             1,
@@ -726,7 +726,8 @@ public static StreamingDataflowWorker fromOptions(DataflowWorkerHarnessOptions o
         Executors.newSingleThreadScheduledExecutor(
             new ThreadFactoryBuilder().setNameFormat("RefreshWork").build()),
         stageInfo,
-        configFetcherComputationStateCacheAndWindmillClient.windmillDispatcherClient());
+        configFetcherComputationStateCacheAndWindmillClient.windmillDispatcherClient(),
+        multiKeyBundleOptions);
   }
 
   /**
@@ -876,7 +877,8 @@ static StreamingDataflowWorker forTesting(
       StreamingCounters streamingCounters,
       WindmillStubFactoryFactory stubFactory) {
     ConcurrentMap<String, StageInfo> stageInfo = new ConcurrentHashMap<>();
-    BoundedQueueExecutor workExecutor = createWorkUnitExecutor(options);
+    MultiKeyBundleOptions multiKeyBundleOptions = MultiKeyBundleOptions.fromOptions(options);
+    BoundedQueueExecutor workExecutor = createWorkUnitExecutor(options, multiKeyBundleOptions);
     ScheduledExecutorService commitFinalizerCleanupExecutor =
         Executors.newScheduledThreadPool(
             1,
@@ -990,7 +992,8 @@ static StreamingDataflowWorker forTesting(
             : windmillStreamFactory.build(),
         executorSupplier.apply("RefreshWork"),
         stageInfo,
-        grpcDispatcherClient);
+        grpcDispatcherClient,
+        multiKeyBundleOptions);
   }
 
   private static GrpcWindmillStreamFactory.Builder createGrpcwindmillStreamFactoryBuilder(
@@ -1020,11 +1023,11 @@ private static JobHeader createJobHeader(DataflowWorkerHarnessOptions options, l
         .build();
   }
 
-  private static BoundedQueueExecutor createWorkUnitExecutor(DataflowWorkerHarnessOptions options) {
+  private static BoundedQueueExecutor createWorkUnitExecutor(
+      DataflowWorkerHarnessOptions options, MultiKeyBundleOptions multiKeyBundleOptions) {
     boolean useFairMonitor =
         DataflowRunner.hasExperiment(options, BOUNDED_QUEUE_EXECUTOR_USE_FAIR_MONITOR_EXPERIMENT);
-    boolean useKeyGroupWorkQueue =
-        DataflowRunner.hasExperiment(options, UNSTABLE_ENABLE_MULTI_KEY_BUNDLE);
+    boolean useKeyGroupWorkQueue = multiKeyBundleOptions.multiKeyBundleEnabled();
     return new BoundedQueueExecutor(
         chooseMaxThreads(options),
         THREAD_EXPIRATION_TIME_SEC,
@@ -1206,9 +1209,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