[Drain] OnTimer - propagate caused by drain bit up to DoFnRunner by stankiewicz · Pull Request #37012 · apache/beam

stankiewicz · 2025-12-05T17:17:20Z

This pull request enhances the Apache Beam Java SDK's timer mechanism by propagating a causedByDrain flag. This flag allows DoFns to distinguish between timers that fire as part of normal pipeline execution and those that are explicitly triggered during a pipeline draining process. By propagating this information through the core timer interfaces and their implementations, it provides more granular control and context for DoFns reacting to timer events e.g. in future within SDF or ReduceFnRunner.

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

gemini-code-assist · 2025-12-05T17:18:33Z

Summary of Changes

Hello @stankiewicz, 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 enhances the Apache Beam Java SDK's timer mechanism by introducing a causedByDrain flag. This flag allows DoFns to distinguish between timers that fire as part of normal pipeline execution and those that are explicitly triggered during a pipeline draining process. By propagating this information through the core timer interfaces and their implementations, it provides more granular control and context for DoFns reacting to timer events, which is crucial for robust pipeline shutdown and data consistency during draining.

Highlights

Propagate 'causedByDrain' Flag: Introduced a boolean causedByDrain parameter to the DoFnRunner.onTimer method and TimerInternals.TimerData to indicate if a timer firing was triggered by a draining operation.
DoFn Context Updates: Added causedByDrain() abstract methods to DoFn.ProcessContext and DoFn.OnTimerContext to expose the drain status to user DoFn logic.
Extensive Integration: Updated numerous DoFnRunner implementations, TimerData factory methods, and Timer utility classes across various runners (Core Java, Flink, Dataflow, Samza, Spark) and SDK components to correctly pass and handle the new causedByDrain flag.
Test Coverage: Modified existing test cases to accommodate the new causedByDrain parameter in TimerData.of and Timer.of calls, ensuring consistency and correctness.

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 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 counter productive. 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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

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. ↩

github-actions · 2025-12-06T00:28:59Z

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

github-actions · 2026-02-04T12:48:18Z

This pull request has been marked as stale due to 60 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that’s incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@beam.apache.org list. Thank you for your contributions.

codecov · 2026-02-05T12:58:41Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 40.06%. Comparing base (4e76198) to head (2612296).

Additional details and impacted files

@@            Coverage Diff            @@
##             master   #37012   +/-   ##
=========================================
  Coverage     40.06%   40.06%           
  Complexity     3404     3404           
=========================================
  Files          1177     1177           
  Lines        187083   187083           
  Branches       3581     3581           
=========================================
+ Hits          74947    74953    +6     
+ Misses       108744   108739    -5     
+ Partials       3392     3391    -1

Flag	Coverage Δ
java	`71.93% <ø> (+0.02%)`	⬆️

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

☔ View full report in Codecov by Sentry.
📢 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.

stankiewicz · 2026-02-17T13:13:05Z

                    .setPaneInfo(currentTimer.getPaneInfo())
+                    .setCausedByDrain(causedByDrain)
                    .setReceiver(
                        windowedValue ->


@kennknowles , ptal.

causedBy drain will be part of windowedValue but because receiver is lambda from WV to context.outputWindowedValue, I'm pretty sure I'm loosing all WV metadata.

stankiewicz · 2026-02-17T13:15:39Z

            .setReceiver(
                windowedValue -> {
                  checkTimerTimestamp(windowedValue.getTimestamp());
                  outputTo(mainOutputConsumer, windowedValue);


@kennknowles here, windowedValue is sent differently, why is that?

I am not sure. This file was a huge mess, with actually many types of transforms executed in the same file via switch statements. It may just be an accidental piece of history and/or a part that I failed to refactor when I introduced OutputBuilder. My quick read is that this one is the idea way to do it, because it passes the whole WindowedValue to the FnDataReceiver so it doesn't lose dasta. In the other places where it goes through the OnTimerContext it loses metadata.

stankiewicz · 2026-02-17T13:20:21Z

        Instant timestamp,
        Collection<? extends BoundedWindow> windows,
        PaneInfo paneInfo) {
      builder(value).setTimestamp(timestamp).setWindows(windows).setPaneInfo(paneInfo).output();


@kennknowles , I assume that by introducing builder we freeze outputWindowedValue and have some limited functionality, but as part of receiver we should use outputWindowedValue method as we risk losing any new fields we've added.

I am not sure if I understand the comment. Here is what I think are answers:

Yes, we freeze the API for outputWindowedValue.

No, we shouldn't lose fields, because builder(value) is an abstract method and it is the job of OutputReceiver.builder(...) to make sure to set all the builder values to defaults, and propagate the values from the current element or timer context.

…ions. Mostly passthrough.

github-actions · 2026-02-17T16:25:28Z

Assigning reviewers:

R: @chamikaramj for label java.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

kennknowles · 2026-02-17T18:00:23Z


+    @Override
+    public CausedByDrain causedByDrain() {
+      return CausedByDrain.NORMAL;


Should this come from the element metadata?

kennknowles · 2026-02-17T18:01:28Z

+        timestamp,
+        outputTimestamp,
+        timeDomain,
+        CausedByDrain.NORMAL);


I think it should propagate here? will that be a follow-up that adds it to the pushback side input DoFnRunner?

ack, adding it. I've missed that interface.

kennknowles · 2026-02-17T18:02:06Z

        new Instant(0),
-        TimeDomain.EVENT_TIME);
+        TimeDomain.EVENT_TIME,
+        CausedByDrain.NORMAL);


Maybe this unit test file should test caused by drain propagation somehow?

ack, adding test

kennknowles · 2026-02-17T18:04:06Z

          new Instant(((Number) kvMap.get("holdTimestamp")).longValue()),
-          paneInfo);
+          paneInfo,
+          CausedByDrain.NORMAL);


I think this should actually come from the data? It is a weird method and I have no context for why it is the way it is...

I think here this requires follow up PR for runner v2 in the future - https://github.com/stankiewicz/beam/blob/e4a5eea37a52429ccec55e3a3a1e0701ba5ff1ee/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto#L901

kennknowles · 2026-02-17T18:07:19Z

        Instant timestamp,
        Collection<? extends BoundedWindow> windows,
        PaneInfo paneInfo) {
      builder(value).setTimestamp(timestamp).setWindows(windows).setPaneInfo(paneInfo).output();


I am not sure if I understand the comment. Here is what I think are answers:

Yes, we freeze the API for outputWindowedValue.

No, we shouldn't lose fields, because builder(value) is an abstract method and it is the job of OutputReceiver.builder(...) to make sure to set all the builder values to defaults, and propagate the values from the current element or timer context.

kennknowles · 2026-02-17T18:07:43Z

            HOLD_TIME,
-            PaneInfo.NO_FIRING));
+            PaneInfo.NO_FIRING,
+            CausedByDrain.NORMAL));


We need some caused by drain tests?

github-actions · 2026-02-25T12:27:47Z

Reminder, please take a look at this pr: @chamikaramj

kennknowles · 2026-02-25T17:01:44Z

waiting on author

fix flink rewrites

stankiewicz · 2026-02-27T15:01:07Z

Run Flink Container PreCommit

kennknowles · 2026-02-27T15:47:48Z

R: @kennknowles

github-actions · 2026-02-27T15:49:01Z

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

kennknowles

I am happy with merging this right now. Regarding timers: there is a gap where we do not propagate or encode metadata, right? So then "caused by drain" will not be set on elements output from a timer during drain unless we always check the is_draining bit. I think this is fine, since during drain every timer could be caused by drain so it won't be mistakenly processed as "not caused by drain" even if it is not truly caused by drain (like if it is set to a past watermark that was before the drain started) but was set because of processing an element that was caused by drain... interesting situation, perhaps.

kennknowles · 2026-02-27T15:53:53Z

        InstantCoder.of().encode(timer.getFireTimestamp(), outStream);
        InstantCoder.of().encode(timer.getHoldTimestamp(), outStream);
        PaneInfoCoder.INSTANCE.encode(timer.getPaneInfo(), outStream);
+        // todo maybe similarly to windowedValue, should we propagate metadata with paneinfo bit


Yes, good point! We need the same metadata capability on encoded timers.

correct, so today, metadata is lost when we set timer (we don't encode metadata anywhere in persistence) and timer expires we don't recover metadata from anywhere. I think this will require follow up exploration as I don't know if windmill would accept additional metadata.

stankiewicz · 2026-02-27T16:14:17Z

I am happy with merging this right now. Regarding timers: there is a gap where we do not propagate or encode metadata, right? So then "caused by drain" will not be set on elements output from a timer during drain unless we always check the is_draining bit. I think this is fine, since during drain every timer could be caused by drain so it won't be mistakenly processed as "not caused by drain" even if it is not truly caused by drain (like if it is set to a past watermark that was before the drain started) but was set because of processing an element that was caused by drain... interesting situation, perhaps.

Correct on gap - if someone sets the timer during draining, metadata is not propagating (e.g tracing) for subsequent onTimer call. but for caused by drain - it will be taken from is_draining bit from windmill.

github-actions Bot added java runners spark dataflow direct core flink fn-execution jet samza labels Dec 5, 2025

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch 2 times, most recently from 0dbde3d to 69b077f Compare December 5, 2025 21:33

github-actions Bot added the stale label Feb 4, 2026

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from 69b077f to 9ff470a Compare February 5, 2026 09:14

github-actions Bot removed the direct label Feb 5, 2026

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch 2 times, most recently from 2d77410 to 2612296 Compare February 5, 2026 11:04

github-actions Bot removed the stale label Feb 5, 2026

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from 2612296 to bf072d2 Compare February 5, 2026 13:30

stankiewicz commented Feb 5, 2026

View reviewed changes

Comment thread runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java Outdated

stankiewicz commented Feb 5, 2026

View reviewed changes

Comment thread runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java Outdated

stankiewicz commented Feb 5, 2026

View reviewed changes

Comment thread sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from bf072d2 to 602243b Compare February 17, 2026 12:54

stankiewicz commented Feb 17, 2026

View reviewed changes

add causedByDrain to DoFnRunner.onTimer interface and all implementat…

03aa66f

…ions. Mostly passthrough.

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from 43050a8 to 03aa66f Compare February 17, 2026 14:23

github-actions Bot added the Next Action: Reviewers label Feb 17, 2026

stankiewicz requested a review from kennknowles February 17, 2026 16:37

kennknowles requested changes Feb 17, 2026

View reviewed changes

github-actions Bot added the slow-review label Feb 25, 2026

github-actions Bot added Next Action: Author direct Next Action: Reviewers and removed Next Action: Reviewers slow-review direct Next Action: Author labels Feb 25, 2026

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from 840eca7 to 38ace4c Compare February 27, 2026 11:30

github-actions Bot added direct and removed direct labels Feb 27, 2026

add missing propagations of causedByDrain

047b7e2

fix flink rewrites

stankiewicz force-pushed the ontimer_propagate_causedByDrain branch from 38ace4c to 047b7e2 Compare February 27, 2026 12:00

github-actions Bot added direct and removed direct labels Feb 27, 2026

stankiewicz requested a review from kennknowles February 27, 2026 15:01

kennknowles approved these changes Feb 27, 2026

View reviewed changes

kennknowles merged commit 5a6f763 into apache:master Feb 27, 2026
30 of 31 checks passed

Conversation

stankiewicz commented Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist Bot commented Dec 5, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions Bot commented Dec 6, 2025

Uh oh!

github-actions Bot commented Feb 4, 2026

Uh oh!

codecov Bot commented Feb 5, 2026

Codecov Report

Uh oh!

Uh oh!

Uh oh!

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!

github-actions Bot commented Feb 17, 2026

Uh oh!

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!

github-actions Bot commented Feb 25, 2026

Uh oh!

kennknowles commented Feb 25, 2026

Uh oh!

stankiewicz commented Feb 27, 2026

Uh oh!

kennknowles commented Feb 27, 2026

Uh oh!

github-actions Bot commented Feb 27, 2026

Uh oh!

kennknowles 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!

Uh oh!

stankiewicz commented Feb 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

stankiewicz commented Dec 5, 2025 •

edited

Loading