[Dataflow Streaming] [Multi Key] MultiKey failure handling + Integration by arunpandianp · Pull Request #38919 · apache/beam

arunpandianp · 2026-06-11T12:11:03Z

The change connects adds failure handling for multi key commits.
Integrates StreamingWorkScheduler and multikey commit methods.
Updates StreamingModeExecutionContext::advance to pull in more items from BoundedWorkQueue

All changes are behind the experiment unstable_enable_multi_key_bundle and does not affect default logic.

…lients - Add MultiKeyWorkItemCommitRequest to windmill.proto. - Support MultiKey commits in Commit model and StreamingEngineWorkCommitter. - Update GrpcCommitWorkStream to batch and stream MultiKey commit requests.

…oring for multi-key execution.

…nContext

# Conflicts: # runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java

Resolved conflicts in StreamingModeExecutionContext.java and StreamingModeExecutionContextTest.java. Fixed compilation error in Work.java by removing duplicate getComputationId() method. TAG=agy CONV=143daaa5-e902-4d26-820d-cf1af2babb84

arunpandianp · 2026-06-11T12:11:51Z

R: @scwhittle This change is on top of #38814 and #38768 PTAL

github-actions · 2026-06-11T12:13:09Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

gemini-code-assist · 2026-06-11T12:17:00Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces multi-key commit support for Dataflow streaming, allowing multiple work items to be committed in a single transaction. It enhances failure handling by enabling re-execution of specific work items within a batch if a retryable failure occurs. Additionally, it optimizes the execution context to pull more work items from the BoundedWorkQueue, thereby increasing processing efficiency for multi-key bundles.

Highlights

Multi-Key Commit Failure Handling: Introduced robust failure handling for multi-key commits, allowing for partial re-execution of work items upon retryable failures.
Streaming Scheduler Integration: Integrated StreamingWorkScheduler with multi-key commit methods to enable efficient batching of work items.
StreamingModeExecutionContext Updates: Updated StreamingModeExecutionContext::advance to pull more items from the BoundedWorkQueue, improving throughput for multi-key bundles.
Experimental Feature Flag: All changes are guarded by the 'unstable_enable_multi_key_bundle' experiment flag, ensuring no impact on default logic.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces multi-key bundle support for streaming in the Google Cloud Dataflow Java worker, enabling batching and advancing through multiple keys within a key group, as well as transactionally committing multi-key work batches. Feedback suggests adding robust error handling when parsing experimental options to prevent worker crashes, and implementing defensive null and bounds checks when accessing read operation receivers to avoid potential runtime exceptions.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-11T12:20:25Z

+    String batchSizeStr =
+        ExperimentalOptions.getExperimentValue(options, WINDMILL_MAX_KEY_GROUP_BATCH_SIZE);
+    this.maxKeyGroupBatchSize = batchSizeStr != null ? Integer.parseInt(batchSizeStr) : 100;
+
+    String batchTimeStr =
+        ExperimentalOptions.getExperimentValue(options, WINDMILL_MAX_KEY_GROUP_BATCH_TIME_MS);
+    this.maxKeyGroupBatchTimeNanos =
+        TimeUnit.MILLISECONDS.toNanos(batchTimeStr != null ? Long.parseLong(batchTimeStr) : 100);


Parsing user-provided experimental options directly using Integer.parseInt and Long.parseLong without error handling can cause the worker to crash if the values are malformed. It is safer to wrap these in a try-catch block and fall back to the default values with a warning log.

String batchSizeStr = ExperimentalOptions.getExperimentValue(options, WINDMILL_MAX_KEY_GROUP_BATCH_SIZE); int batchSize = 100; if (batchSizeStr != null) { try { batchSize = Integer.parseInt(batchSizeStr); } catch (NumberFormatException e) { LOG.warn("Failed to parse {} as integer, using default of 100", WINDMILL_MAX_KEY_GROUP_BATCH_SIZE, e); } } this.maxKeyGroupBatchSize = batchSize; String batchTimeStr = ExperimentalOptions.getExperimentValue(options, WINDMILL_MAX_KEY_GROUP_BATCH_TIME_MS); long batchTimeMs = 100; if (batchTimeStr != null) { try { batchTimeMs = Long.parseLong(batchTimeStr); } catch (NumberFormatException e) { LOG.warn("Failed to parse {} as long, using default of 100", WINDMILL_MAX_KEY_GROUP_BATCH_TIME_MS, e); } } this.maxKeyGroupBatchTimeNanos = TimeUnit.MILLISECONDS.toNanos(batchTimeMs);

gemini-code-assist · 2026-06-11T12:20:25Z

+    HashMap<String, ElementCounter> counters =
+        ((DataflowMapTaskExecutor) workExecutor)
+            .getReadOperation()
+            .receivers[0]
+            .getOutputCounters();


Defensive programming: Accessing receivers[0] directly without checking if getReadOperation() is null, or if receivers is null or empty, can lead to NullPointerException or ArrayIndexOutOfBoundsException. Adding appropriate guards ensures robust execution.

DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor; if (mapTaskExecutor.getReadOperation() == null || mapTaskExecutor.getReadOperation().receivers == null || mapTaskExecutor.getReadOperation().receivers.length == 0) { return 0L; } HashMap<String, ElementCounter> counters = mapTaskExecutor.getReadOperation().receivers[0].getOutputCounters(); if (counters == null) { return 0L; }

codecov · 2026-06-11T15:03:49Z

Codecov Report

❌ Patch coverage is 88.44884% with 35 lines in your changes missing coverage. Please review.
✅ Project coverage is 55.39%. Comparing base (1cf3545) to head (0df9f54).
⚠️ Report is 8 commits behind head on master.

Files with missing lines	Patch %	Lines
...e/beam/runners/dataflow/worker/streaming/Work.java	69.44%	10 Missing and 1 partial ⚠️
...dataflow/worker/StreamingModeExecutionContext.java	94.59%	8 Missing and 2 partials ⚠️
...nners/dataflow/worker/WindowingWindmillReader.java	81.08%	3 Missing and 4 partials ⚠️
...ers/dataflow/worker/util/BoundedQueueExecutor.java	75.00%	0 Missing and 3 partials ⚠️
...nners/dataflow/worker/StreamingDataflowWorker.java	60.00%	1 Missing and 1 partial ⚠️
...unners/dataflow/worker/WorkCancelingException.java	71.42%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@              Coverage Diff              @@
##             master   #38919       +/-   ##
=============================================
- Coverage     59.33%   55.39%    -3.95%     
+ Complexity    16593     2243    -14350     
=============================================
  Files          2845     1104     -1741     
  Lines        291334   171355   -119979     
  Branches      14421     1437    -12984     
=============================================
- Hits         172859    94915    -77944     
+ Misses       111065    74000    -37065     
+ Partials       7410     2440     -4970

Flag	Coverage Δ
java	`75.76% <88.44%> (+9.65%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

derrickaw · 2026-07-08T15:04:10Z

Hi @scwhittle, can you review this at your convenience? Thanks

scwhittle

waiting on tests until previous change has been merged

scwhittle · 2026-07-22T17:39:02Z

@@ -88,6 +88,11 @@ static ActiveWorkState create(WindmillStateCache.ForComputation computationState
    return new ActiveWorkState(new HashMap<>(), computationStateCache);
  }

+  synchronized Optional<ExecutableWork> getActiveWork(ShardedKey shardedKey, WorkId workId) {


should we just return nullable ExecutableWork instead? With annotations that seems preferrable to me to avoid the optional allocation

scwhittle · 2026-07-22T17:39:58Z

@@ -131,6 +131,10 @@ public void completeWorkAndScheduleNextWorkForKey(ShardedKey shardedKey, WorkId
        .ifPresent(this::forceExecute);
  }

+  public void reExecuteActiveWork(ShardedKey shardedKey, WorkId workId) {


nit: I would uncapitalized E since reexecute is one word

scwhittle · 2026-07-22T17:41:32Z

    this.latencyTrackingId =
        Long.toHexString(workItem.getShardingKey())
            + '-'
            + Long.toHexString(workItem.getWorkToken());
    this.currentState = TimedState.initialState(startTime);
    this.isFailed = false;
+    this.getWorkStreamLatencies = getWorkStreamLatencies;
  }

  public static Work create(


if we don't use this often, perhaps we should just add the ImmutableList.of() to the call-site

scwhittle · 2026-07-22T17:45:12Z

-    for (LatencyAttribution latency : getWorkStreamLatencies) {
-      totalDurationPerState.put(
-          latency.getState(), Duration.millis(latency.getTotalDurationMillis()));
+  public ImmutableList<LatencyAttribution> getWorkStreamLatencies() {


I don't see usage of this, can we just get rid of this.getWorkStreamLatencies and instead do the recordGetWorkStreamLatencies logic in the constructor?

removed getWorkStreamLatencies().

recordGetWorkStreamLatencies was initially inside the constructor, moved it out of the constructor to reduce load on the Getworkstream thread. recordGetWorkStreamLatencies runs in the processing thread.

Not sure if you sent comments, but I was also suggesting removing the member variable and public recordGetWorkStreamLatencies method as well. Seems it could just be done in constructor?

Sorry forgot to publish comments. Published now.

scwhittle · 2026-07-22T17:46:31Z

-      this.elements = elements;
+    private BoundedQueueExecutorWorkHandleImpl(Work work, long bytes) {
+      checkArgument(bytes >= 0);
+      this.workBatch = new ArrayList<>();


could workBatch be an ImmutableList.Builder?

changing it to ImmutableList.Builder is adding more copies in merge(). ArrayList seems to be better.

scwhittle · 2026-07-22T17:52:30Z

+    Windmill.MultiKeyWorkItemCommitRequest.Builder multiKeyBuilder =
+        Windmill.MultiKeyWorkItemCommitRequest.newBuilder();
+
+    Work primaryWork = workBatch.get(0);


checkState that workBatch is non-empty

scwhittle · 2026-07-22T17:53:09Z

+
+    for (int i = 0; i < workBatch.size(); i++) {
+      // TODO: Add commit size validation
+      Windmill.WorkItemCommitRequest commit = workItemCommits.get(i);


precondition that workitemcommits and workBatch are same size

scwhittle · 2026-07-22T17:57:43Z

+      boolean hotKeyLoggingEnabled,
+      String stepName,
+      String sourceBytesProcessCounterName,
+      PipelineOptions options,


instead of passing in full options, could pass some internal multikeyoptions struct

benefits is that the parsing doesn't have to be here and can be shared across all the contexts for now. But in the future we may want to configure differently for different fused stages based upon other information and that would allow us to do so as it would be separate from the single experiment value.

Changed to use MultiKeyBundleOptions.

scwhittle · 2026-07-22T17:58:41Z

+      return false;
+    }
+    if (workIsFailed()) {
+      throw new WorkItemCancelledException(checkStateNotNull(work).getWorkItem().getShardingKey());


move activeWork definition above and use it here instead of separate check

scwhittle · 2026-07-22T18:04:17Z

+ * Indicates that the work is no longer valid and should be canceled. It is thrown as a signal for
+ * upper layers to mark the work as failed.
+ */
+public class WorkCancelingException extends RuntimeException {


nit: can we use Cancelling with two ls to match other usage in beam

should we just use WorkItemCancelledException.java? If not can you explain the differences here?

Renamed and added differences. once is thrown before marking work as failed and the other is thrown after noticing that a work is failed.

scwhittle

Code looks mostly good, still going through tests

scwhittle · 2026-07-24T22:51:27Z

            + Long.toHexString(workItem.getWorkToken());
    this.currentState = TimedState.initialState(startTime);
    this.isFailed = false;
+    this.getWorkStreamLatencies = getWorkStreamLatencies;


instead of saving just populate totalDurationPerState here

this was intentionally moved out of the submission loop in #33736

can you add a comment here?
// We defer recordGetWorkStreamLatencies() to be called during bundle processing
// as these are constructed on the hot GetWork thread

Thanks, done.

scwhittle · 2026-07-24T22:54:27Z

+      switch (evaluateRetry(computationId, executableWork.work(), t)) {
+        case DO_NOT_RETRY:
+          // Consider the item invalid. It will eventually be retried by Windmill if it still needs
+          // to


nit, wrap comment

scwhittle · 2026-07-24T22:57:59Z

+    if (workBatch.isEmpty()) {
+      return;
+    }
+    if (workBatch.size() > 1 || multiKeyBundleOptions.multiKeyBundleEnabled()) {


is there a benefit to sending single keys in multi-key format? Otherwise seems likely more overhead in protos etc

The benefit is eventually we can remove the single key code path and have only the multikey path. Having one code path on a job also will make debugging easier.

scwhittle

finished looking at all tests, so think this is likely final round

scwhittle · 2026-07-27T19:01:14Z

            + Long.toHexString(workItem.getWorkToken());
    this.currentState = TimedState.initialState(startTime);
    this.isFailed = false;
+    this.getWorkStreamLatencies = getWorkStreamLatencies;


can you add a comment here?
// We defer recordGetWorkStreamLatencies() to be called during bundle processing
// as these are constructed on the hot GetWork thread

scwhittle · 2026-07-27T19:05:04Z

+        Windmill.MultiKeyWorkItemCommitRequest.newBuilder();
+
+    Work primaryWork = workBatch.get(0);
+    Work.KeyGroup keyGroup = primaryWork.getKeyGroup();


since we are using this for single keys also that likely don't have a group, should we avoid setting the key group nested field if 0?

scwhittle · 2026-07-27T19:06:21Z

+        Windmill.Uint128Proto.newBuilder().setHigh(keyGroup.high()).setLow(keyGroup.low()).build());
+
+    for (int i = 0; i < workBatch.size(); i++) {
+      // TODO: Retry on commit truncations


should we throw an exception for now?

scwhittle · 2026-07-27T19:09:02Z

  // Key switch listener to delegate MDC logging context and thread name updates
  public interface KeyTransitionListener {
-    void onKeyTransition(Work oldWork, Work newWork);
+    void onKeyTransition(@Nullable Work oldWork, Work newWork);


// oldWork is null when newWork is the first work for the bundle.

scwhittle · 2026-07-27T19:10:15Z

    currentBuilder.clear();
    currentBuilder.mergeFrom(truncationBuilder.build());
+
+    // TODO: throw and retry when truncation is not on a single key bundle.


can we throw an exception if multikey bundles are enabled to make sure we don't lose data if we forget to address this?

scwhittle · 2026-07-27T19:32:37Z

+        makeWorker(
+            defaultWorkerParams(
+                    "--experiments=unstable_enable_multi_key_bundle,windmill_max_key_group_batch_time_ms=50000",
+                    "--numberOfWorkerHarnessThreads=1")


do we need this harness threads because otherwise we start new threads to process the same group in parallel?

If so I wonder if that something we might want to try to prevent in the future? If we have more threads than # of CPU, it is likely better to not parallelize a key group.

Yes, there is room to improve the batching. Planning to tackle it separately.

scwhittle · 2026-07-27T19:33:37Z

+                          .setKey(keyRequest.getKey())
+                          .setShardingKey(keyRequest.getShardingKey());
+                  if (keyRequest.getWorkToken() == 2) {
+                    keyBuilder.setFailed(true);


if we do add support for the generic response to the server, we could also add the ability to set work tokens to fail get data requests.

Added StreamingDataflowWorkerTest::emptyDataResponderWithFailedWorkTokens that take the failed work tokens.

scwhittle · 2026-07-27T19:35:08Z

    options = PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class);
+    options
+        .as(ExperimentalOptions.class)
+        .setExperiments(Arrays.asList("unstable_enable_multi_key_bundle"));


scwhittle · 2026-07-27T19:36:19Z

+        work1, workExecutor, mockExecutor, mockHandle, null, (oldWork, newWork) -> {});
+
+    assertTrue(executionContext.advance());
+    assertEquals("key2", executionContext.getSerializedKey().toStringUtf8());


advance again?

scwhittle · 2026-07-27T19:37:11Z

+
+    work1.setFailed();
+
+    assertThrows(WorkItemCancelledException.class, () -> executionContext.advance());


assert no interactions on executor

arunpandianp added 24 commits June 2, 2026 11:33

[Dataflow Streaming] [Multi Key] StreamingModeExecutionContext refact…

73faa68

…oring for multi-key execution.

trigger postsubmit tests

53bc9a6

fix tests

7720cf8

fix tests

72740d2

improve work synchronization

0f96e72

cleanup logic

51b9257

cleanup logic

9a4e7be

address comments

3f36afd

improve WindowingWindmillReader

f3cc628

spotless fix

58e0ef9

[Dataflow Streaming] Fix nullness supression in StreamingModeExecutio…

3dceab0

…nContext

make windmillTagEncoding final

e199438

address comments

700dfbc

Move SideInputStateFetcherFactory from start to constructor

bc5bee2

Merge branch 'contextnullness' into multikey_context_review

6d7f28e

# Conflicts: # runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java

Merge remote-tracking branch 'beam/master' into multikey_context_review

c4a52c1

Merge beam/master into multikey_context_review

1107a60

Resolved conflicts in StreamingModeExecutionContext.java and StreamingModeExecutionContextTest.java. Fixed compilation error in Work.java by removing duplicate getComputationId() method. TAG=agy CONV=143daaa5-e902-4d26-820d-cf1af2babb84

Address comment

47eb7d6

Address comment

24505dd

Fix UnderInitialization

4e0d174

Merge branch 'multikey_commit' into tmp

5f9fef5

Merge branch 'multikey_context_review' into tmp

0e8157f

Multikey commit failure handling and integration

93de23a

github-actions Bot added runners dataflow labels Jun 11, 2026

gemini-code-assist Bot reviewed Jun 11, 2026

View reviewed changes

arunpandianp added 2 commits June 11, 2026 13:04

Fix tests and add sink byte limit for batching

9de3d84

spotless

0df9f54

arunpandianp added 4 commits June 29, 2026 08:55

Merge remote-tracking branch 'beam/master' into multikey_commit

18ae844

Merge remote-tracking branch 'beam/master' into multikey_failure

ad177c0

Merge branch 'multikey_commit' into multikey_failure

557f95e

fix merge

00e5fbf

derrickaw added the Next Action: Reviewers label Jul 8, 2026

derrickaw requested a review from scwhittle July 8, 2026 15:04

scwhittle requested changes Jul 22, 2026

View reviewed changes

arunpandianp added 6 commits July 23, 2026 08:55

Address comments

9ff2cad

Merge branch 'multikey_commit' into multikey_failure

6c5481a

Merge remote-tracking branch 'beam/master' into multikey_failure

5ddc9d4

fix merge

ab33927

fix merge

ef77b1f

Address comments

43c7f1a

arunpandianp requested a review from scwhittle July 24, 2026 07:42

scwhittle requested changes Jul 24, 2026

View reviewed changes

fix comment

68fc0bc

scwhittle requested changes Jul 27, 2026

View reviewed changes

Address comments

60fec05

arunpandianp requested a review from scwhittle July 29, 2026 13:49

scwhittle approved these changes Jul 29, 2026

View reviewed changes

scwhittle merged commit 3c9f9ce into apache:master Jul 29, 2026
16 checks passed


		work1.setFailed();

		assertThrows(WorkItemCancelledException.class, () -> executionContext.advance());

Uh oh!

Conversation

arunpandianp commented Jun 11, 2026

Uh oh!

arunpandianp commented Jun 11, 2026

Uh oh!

github-actions Bot commented Jun 11, 2026

Uh oh!

gemini-code-assist Bot commented Jun 11, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Jun 11, 2026

Codecov Report

Uh oh!

derrickaw commented Jul 8, 2026

Uh oh!

scwhittle left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

scwhittle left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!