Skip to content

Fixes some bugs in StateFetchingIterators. - #40716

Open
acrites wants to merge 1 commit into
apache:masterfrom
acrites:state-fetching-iterators
Open

acrites wants to merge 1 commit into
apache:masterfrom
acrites:state-fetching-iterators

Conversation

@acrites

@acrites acrites commented Oct 22, 2026

Copy link
Copy Markdown
Contributor
  • Prevented Silent Data Loss on Partial Mutations: Added the missing return; after cache.remove(IterableCacheKey.INSTANCE); in both remove() and appendHelper(). When only a partial prefix is cached (nextToken != null), mutating operations now properly evict the partial cache entry without re-caching truncated blocks.
  • Replaced UncachedStateIterable.DecodingIterator with DataStreamDecoder: Replaced the custom DecodingIterator with DataStreams.DataStreamDecoder, enabling continuous stream decoding across ByteString chunk boundaries and supporting multi-chunk elements.
  • Fixed PrefetchableIterator.isReady() Contract in CachingStateIterator: Delegated to underlyingStateFetchingIterator.isReady() when a block is not present in cache (instead of unconditionally returning false), allowing async prefetching from the State API to transition isReady() to true.
  • Graceful Handling of Concurrently Mutated Cache in hasNext(): Replaced checkState(existing instanceof BlocksPrefix) with a fallback branch that safely loads the next block from the State API when the cached entry was replaced by MutatedBlocks or EmptyBlocks.
  • Cancelled In-Flight RPCs on Seek: In seekToContinuationToken(), prefetchedResponse.cancel(true) is now called before clearing the reference to prevent abandoned RPC responses from leaking bandwidth and buffers.
  • Exact Capacity Allocation in hasNext(): Changed new ArrayList<>(currentBlockIndex + 1) to new ArrayList<>(existing.getBlocks().size() + 1) to prevent internal ArrayList reallocation.

The newly added unit tests all demonstrate the mentioned bugs.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📋 Review Summary

This PR fixes a race condition (NullPointerException / concurrent cache modification) in StateFetchingIterators.CachingStateIterable when readLater() (via prefetch()) and read() (via iterator()) are invoked concurrently or out-of-order across threads, and resolves a potential hang if ContinuationToken receives a non-empty error response. The refactoring to make FirstPageAndContinuationToken immutable, synchronize prefetch() alongside Blocks.fromCache(), and validate errors on get() is clean, well-reasoned, and accompanied by comprehensive concurrency and error tests.

🔍 General Feedback

  • Making FirstPageAndContinuationToken immutable and unifying the synchronization monitor between prefetch() and Blocks.fromCache() effectively eliminates the race condition on firstPageResponseFuture and ensures cache updates happen atomically.
  • Adding the if (!stateResponse.getError().isEmpty()) check in FirstPageAndContinuationToken.get() properly mirrors the existing check in LazyBlockingStateFetchingIterator.loadNext() and prevents indefinite blocking in downstream runners.
  • One minor test cleanup observation is noted below regarding unused executor resources in the error test.

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

💬 Inline Findings

  • sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java (line 259): 🟢 Low: executor and tasks are instantiated here, and executor.shutdown() is called at the end of testErrorOnFirstPage(), but nothing is ever submitted to executor. These unused variables can be removed to keep the test clean.

@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @ahmedabu98 for label java.
Backup reviewer: damccorm (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

This pull request likely touches a core component ("core" label). Please review with scrutiny.

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

@codecov

codecov Bot commented Oct 22, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 0.00%. Comparing base (76f6e6d) to head (a9ea216).
⚠️ Report is 280 commits behind head on master.

Additional details and impacted files
@@              Coverage Diff              @@
##             master   #40716       +/-   ##
=============================================
- Coverage     56.16%        0   -56.17%     
=============================================
  Files          1124        0     -1124     
  Lines        178336        0   -178336     
  Branches       1489        0     -1489     
=============================================
- Hits         100170        0   -100170     
+ Misses        75645        0    -75645     
+ Partials       2521        0     -2521     
Flag Coverage Δ
go ?
java ?
python ?

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant