[Flink] Bound SDF self-checkpoint residual state on the portable runner by Eliaaazzz · Pull Request #39191 · apache/beam

Eliaaazzz · 2026-07-01T15:52:20Z

Splittable DoFns that self-checkpoint on the portable Flink runner leaked keyed state without bound. StateAndTimerBundleCheckpointHandler.onCheckpoint wrote each residual under a fresh, unique ValueState tag (sdf_checkpoint:<id>:<index>), so a polling SDF (for example one using tracker.defer_remainder()) registered a new Flink keyed-state descriptor in registeredKVStates on every self-checkpoint, and the job's heap grew until it OOMed.

This stores all residuals for a key and window under a single stable MapState (sdf_checkpoint_residuals) keyed by the per-residual checkpoint id, and removes each entry once its timer fires and the residual is re-processed. Both read paths, streaming ExecutableStageDoFnOperator.onTimer and batch FlinkExecutableStageFunction, look the residual up by timer id. Timer ids stay dynamic; Flink already keeps them in a stable pending-timers map that is cleared on fire.

Residuals persisted by the previous code under per-residual ValueState tags are not migrated, so an in-flight residual from an older savepoint is dropped on this experimental portability path.

BundleCheckpointHandlersTest drives repeated self-checkpoints through the handler and asserts a single keyed-state descriptor is used regardless of the number of self-checkpoints; it fails on the previous per-residual behavior.

Fixes #27648

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 · 2026-07-01T15:54:06Z

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 addresses a memory leak in the portable Flink runner affecting Splittable DoFns (SDFs) that perform self-checkpoints. Previously, each self-checkpoint created a unique ValueState tag, leading to an accumulation of keyed-state descriptors and eventual OutOfMemory errors. The changes introduce a stable MapState to manage these residuals, ensuring the state remains bounded across checkpoints. Note that residuals from older savepoints using the previous mechanism are not migrated and will be dropped.

Highlights

State Management Optimization: Replaced per-residual ValueState tags with a single, stable MapState (sdf_checkpoint_residuals) to store SDF self-checkpoint residuals, preventing unbounded state growth.
Memory Leak Prevention: Fixed an OOM issue on the portable Flink runner where polling Splittable DoFns (SDFs) leaked keyed state descriptors on every self-checkpoint.
Regression Testing: Added BundleCheckpointHandlersTest to verify that repeated self-checkpoints now correctly use a single, bounded state descriptor.

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 addresses unbounded checkpoint state growth for splittable DoFns (SDFs) that self-checkpoint on the portable Flink runner by migrating from per-residual state tags to a single, stable MapState tag. It also ensures consumed residuals are removed to keep the state bounded. The review feedback highlights a critical issue where reading from older savepoints can return null, leading to potential NullPointerExceptions in both FlinkExecutableStageFunction and ExecutableStageDoFnOperator. Adding null checks in these locations is recommended to safely drop the missing residuals.

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-07-01T15:55:10Z

+            WindowedValue<InputT> stateValue = residualState.get(timer.getTimerId()).read();
+            residualState.remove(timer.getTimerId());
            residuals.add(stateValue);


If an older savepoint is used, the residual state won't be found under the new MapState tag, and residualState.get(timer.getTimerId()).read() will return null. Adding null to residuals will lead to a NullPointerException during processElements. We should add a null check to safely drop the residual as intended.

WindowedValue<InputT> stateValue = residualState.get(timer.getTimerId()).read(); residualState.remove(timer.getTimerId()); if (stateValue != null) { residuals.add(stateValue); }

gemini-code-assist · 2026-07-01T15:55:10Z

+        WindowedValue<InputT> stateValue = residualState.get(timerId).read();
        processElement(stateValue);


If an older savepoint is used, residualState.get(timerId).read() will return null. Passing null to processElement will cause a NullPointerException. We should add a null check to safely drop the residual as intended.

Suggested change

WindowedValue<InputT> stateValue = residualState.get(timerId).read();

processElement(stateValue);

WindowedValue<InputT> stateValue = residualState.get(timerId).read();

if (stateValue != null) {

processElement(stateValue);

}

Abacn

Also, as we confirmed SDK now functioning as expected, shall we remove these TODO? https://github.com/apache/beam/blob/master/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java#L735

Abacn · 2026-07-01T17:46:02Z

      return SDF_PREFIX + ":" + id + ":" + index;
    }

+    /** The state id under which all SDF self-checkpoint residuals for a key/window are stored. */


Just a note, it appears BundleCheckpointHandlers is only used by Flink runner, thus this bug only affects Flink runner as well

Abacn · 2026-07-01T18:02:13Z

  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
 })
 public class FlinkExecutableStageFunction<InputT> extends AbstractRichFunction
    implements MapPartitionFunction<WindowedValue<InputT>, RawUnionValue>,


This file is branched in runners/flink/2.0/src. Fix needs to be done in both place

github-actions · 2026-07-01T18:13:30Z

Assigning reviewers:

R: @chamikaramj added as fallback since no labels match configuration

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

Eliaaazzz · 2026-07-02T04:40:51Z

Thanks for the review!

On the TODO at L735, I removed // TODO: Fail on splittable DoFns. since SDFs now work as expected. The adjacent // TODO: Special-case single outputs to avoid multiplexing PCollections. is about output multiplexing rather than SDF support and is still not implemented, so I left it in. Happy to remove that one too if you meant both.

I also applied the same fix to the branched runners/flink/2.0 copy of FlinkExecutableStageFunction, and added the null checks Gemini flagged for the old-savepoint read path.

A splittable DoFn that self-checkpoints, for example a polling source using tracker.defer_remainder(), stored each residual under a new unique state tag, so every self-checkpoint registered another Flink keyed-state descriptor and the job leaked heap without bound. Store all residuals for a key and window under a single stable MapState keyed by the checkpoint id, and remove each entry once its timer fires and the residual is re-processed. The streaming and batch read paths look the residual up by timer id. Fixes apache#27648

github-actions Bot added runners flink fn-execution labels Jul 1, 2026

gemini-code-assist Bot reviewed Jul 1, 2026

View reviewed changes

Eliaaazzz force-pushed the flink-sdf-checkpoint-state-leak branch from 235b200 to a42f34d Compare July 1, 2026 16:03

Abacn reviewed Jul 1, 2026

View reviewed changes

github-actions Bot added the Next Action: Reviewers label Jul 1, 2026

Eliaaazzz force-pushed the flink-sdf-checkpoint-state-leak branch from a42f34d to 27aadfa Compare July 2, 2026 03:52

Eliaaazzz force-pushed the flink-sdf-checkpoint-state-leak branch from 27aadfa to 766000a Compare July 2, 2026 05:02

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Flink] Bound SDF self-checkpoint residual state on the portable runner#39191

[Flink] Bound SDF self-checkpoint residual state on the portable runner#39191
Eliaaazzz wants to merge 1 commit into
apache:masterfrom
Eliaaazzz:flink-sdf-checkpoint-state-leak

Eliaaazzz commented Jul 1, 2026

Uh oh!

gemini-code-assist Bot commented Jul 1, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 1, 2026

Uh oh!

gemini-code-assist Bot Jul 1, 2026

Uh oh!

Abacn left a comment

Uh oh!

Abacn Jul 1, 2026

Uh oh!

Abacn Jul 1, 2026

Uh oh!

github-actions Bot commented Jul 1, 2026

Uh oh!

Eliaaazzz commented Jul 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		WindowedValue<InputT> stateValue = residualState.get(timerId).read();
		processElement(stateValue);

Uh oh!

Conversation

Eliaaazzz commented Jul 1, 2026

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist Bot commented Jul 1, 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 Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

Abacn left a comment

Choose a reason for hiding this comment

Uh oh!

Abacn Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

Abacn Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 1, 2026

Uh oh!

Eliaaazzz commented Jul 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants