[pull] master from apache:master

.github/actions/setup-environment-action/action.yml

@@ -56,6 +56,24 @@ runs:
           sdks/python/setup.py
           sdks/python/tox.ini
 
+    # Multiline python-version (e.g. 3.10\n3.11\n...) must not be embedded raw
+    # in cache keys: actions/cache treats each newline-separated line as a
+    # separate key and fails when there are more than 10 ("Keys are limited to
+    # a maximum of 10"). Collapse whitespace into a single key fragment.
+    - name: Sanitize Python version for tox cache key
+      if: ${{ inputs.python-version != '' && inputs.tox-cache == 'true' }}
+      id: py-cache-key
+      shell: bash
+      env:
+        PYTHON_VERSION_INPUT: ${{ inputs.python-version }}
+      run: |
+        if [[ "$PYTHON_VERSION_INPUT" == "default" ]]; then
+          sanitized="310"
+        else
+          sanitized=$(printf '%s' "$PYTHON_VERSION_INPUT" | tr -s '[:space:]' '-' | sed 's/^-//;s/-$//')
+        fi
+        echo "sanitized=$sanitized" >> "$GITHUB_OUTPUT"
+
     - name: Cache tox environments
       if: ${{ inputs.python-version != '' && inputs.tox-cache == 'true' }}
       uses: actions/cache@v6
@@ -64,10 +82,10 @@ runs:
           sdks/python/target/.tox
           !sdks/python/target/.tox/**/log
           !sdks/python/target/.tox/.package_cache
-        key: tox-${{ runner.os }}-py${{ inputs.python-version == 'default' && '310' || inputs.python-version }}-${{ hashFiles('sdks/python/tox.ini') }}-${{ hashFiles('sdks/python/setup.py') }}
+        key: tox-${{ runner.os }}-py${{ steps.py-cache-key.outputs.sanitized }}-${{ hashFiles('sdks/python/tox.ini') }}-${{ hashFiles('sdks/python/setup.py') }}
         restore-keys: |
-          tox-${{ runner.os }}-py${{ inputs.python-version == 'default' && '310' || inputs.python-version }}-${{ hashFiles('sdks/python/tox.ini') }}-
-          tox-${{ runner.os }}-py${{ inputs.python-version == 'default' && '310' || inputs.python-version }}-
+          tox-${{ runner.os }}-py${{ steps.py-cache-key.outputs.sanitized }}-${{ hashFiles('sdks/python/tox.ini') }}-
+          tox-${{ runner.os }}-py${{ steps.py-cache-key.outputs.sanitized }}-
 
     - name: Install Java
       if: ${{ inputs.java-version != '' }}

.github/workflows/build_release_candidate.yml

@@ -316,7 +316,7 @@ jobs:
         # settings.xml file
         run: rm ~/.m2/settings.xml || true
       - name: Login to Docker Hub
-        uses: docker/login-action@abd2ef45e78c5afb21d64d4ca52ee8550d9572c7
+        uses: docker/login-action@371161bbe7024a29a25c5e19bfcbc0804fe9ad2c
         with:
           username: ${{ secrets.DOCKERHUB_USER }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}

.github/workflows/finalize_release.yml

@@ -44,7 +44,7 @@ jobs:
     runs-on: [self-hosted, ubuntu-24.04, main]
     steps:
     - name: Login to Docker Hub
-      uses: docker/login-action@abd2ef45e78c5afb21d64d4ca52ee8550d9572c7
+      uses: docker/login-action@371161bbe7024a29a25c5e19bfcbc0804fe9ad2c
       with:
         username: ${{ secrets.DOCKERHUB_USER }}
         password: ${{ secrets.DOCKERHUB_TOKEN }}

.github/workflows/python_dependency_tests.yml

@@ -1,6 +1,7 @@
 name: Python Dependency Tests
 
 on:
+  workflow_dispatch:
   schedule:
     - cron: '10 2 * * *'
   pull_request:

.github/workflows/tour_of_beam_backend_integration.yml

@@ -20,6 +20,7 @@
 name: Tour of Beam Go integration tests
 
 on:
+  workflow_dispatch:
   push:
     branches: ['master', 'release-*']
     tags: 'v*'

.github/workflows/update_python_dependencies.yml

@@ -69,6 +69,8 @@ jobs:
           java-version: default
           go-version: default
           disable-cache: true
+          # generatePythonRequirementsAll does not use tox.
+          tox-cache: false
       - name: Update Python Dependencies
         uses: ./.github/actions/gradle-command-self-hosted-action
         with:

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