Invalidate cached ContinuationIterable when its continuation fails

[addenergyx]

Fixes #39383

GlobalCachingStateHandler.blocking_get caches partially materialized bag state as a ContinuationIterable, keyed by the user-state cache token. The iterable's lazy tail holds the continuation token from the read that created it. If the runner permanently invalidates that token (Dataflow revokes them together with the work item that created them), the cached object fails on every later read of that state and the cache never recovers: each retry is served the same object and replays the dead token instead of re-fetching. See #39383 for the production symptoms.

This PR invalidates the cache entry when iterating the continuation fails, so the next read loads fresh state. A transient failure costs one extra reload; a permanent one self-heals instead of retrying indefinitely. The invalidation goes through a new identity-guarded StateCache.invalidate_if_value, so a stale iterable failing again later cannot evict a newer value cached under the same key, and a failure inside the invalidation itself is logged rather than allowed to mask the original exception. Backward compatible: ContinuationIterable.on_failure defaults to None.

Validation via new deterministic tests (no Dataflow required; the fake state handler uses epoch-scoped token revocation so revoked tokens fail while fresh reads succeed, matching real runner semantics): test_failed_continuation_invalidates_cached_iterable, test_failed_continuation_invalidates_after_partial_yield, test_stale_failed_iterable_does_not_evict_replacement, test_invalidate_if_value, test_invalidate_if_value_ignores_in_flight_load, and test_invalidation_failure_does_not_mask_continuation_error.

• Without the fix: the core regression test fails. The second bundle is served the cached iterable and raises RuntimeError: continuation token no longer valid from the dead token.
• With the fix: sdk_worker_test.py (including all pre-existing tests) and statecache_test.py pass, 40 tests total.

---

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

🤖 Generated with Claude Code

[gemini-code-assist]

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 critical issue in the Python SDK's state caching mechanism where permanently invalid continuation tokens could lead to infinite retry storms. By ensuring that failed state reads invalidate the corresponding cache entry, the system can now recover from transient or permanent runner-side token revocations instead of repeatedly attempting to use stale data.

Highlights

• Cache Invalidation: Implemented automatic cache invalidation for ContinuationIterable when continuation tokens fail, preventing infinite retry loops caused by stale, invalid tokens.
• Error Handling: Added an on_failure callback mechanism to ContinuationIterable to ensure that failed state reads trigger a cache eviction, allowing subsequent attempts to fetch fresh state.
• Regression Testing: Introduced a new unit test, test_failed_continuation_invalidates_cached_iterable, which simulates token revocation to verify that the cache correctly clears and recovers from poisoned state entries.

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.

Footnotes

1. 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 a mechanism to invalidate cached iterables in the SDK worker when a continuation fails, preventing permanently invalid continuation tokens from remaining in the cache. A unit test has been added to verify this behavior. The review feedback suggests wrapping the cache invalidation callback in a try-except block to ensure that any failure during invalidation does not mask the primary exception raised during iteration.

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]

If self.on_failure() raises an exception, it will mask the original exception raised by self.continue_iterator_fn(). To prevent masking the primary failure and ensure it is correctly propagated, we should wrap the on_failure() call in a try...except block and log any invalidation failures.

Suggested change

	if self.on_failure is not None:
	self.on_failure()
	if self.on_failure is not None:
	try:
	self.on_failure()
	except Exception:
	_LOGGER.warning('Failed to invalidate cache on continuation failure.', exc_info=True)

[addenergyx]

Closing for now, will reopen after further review.

[codecov]

Codecov Report

❌ Patch coverage is 90.47619% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 58.05%. Comparing base (0643150) to head (f2362fa).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
...ks/python/apache_beam/runners/worker/sdk_worker.py	83.33%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff            @@
##             master   #39384   +/-   ##
=========================================
  Coverage     58.04%   58.05%           
  Complexity    13051    13051           
=========================================
  Files          2518     2518           
  Lines        263804   263822   +18     
  Branches      10764    10764           
=========================================
+ Hits         153138   153160   +22     
+ Misses       104911   104907    -4     
  Partials       5755     5755           

Flag	Coverage Δ
python	79.75% <90.47%> (+<0.01%)	⬆️

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.