Add LRU cache eviction to CachingStateProvider #37214
Conversation
Fixes apache#37213 Implements LRU (Least Recently Used) cache eviction to prevent unbounded memory growth in long-running workers. Adds configurable maxCacheSize parameter (default: 1000 entries) and maintains LRU order using JavaScript Map's insertion order. - Add maxCacheSize constructor parameter with default value of 1000 - Implement evictIfNeeded() to remove oldest entry when cache is full - Implement touchCacheEntry() to move accessed items to end (LRU) - Add comprehensive test coverage in state_provider_test.ts This addresses the TODO comment in the code and improves reliability for production workloads.
Summary of ChangesHello @junaiddshaukat, I'm Gemini Code Assist1! 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 memory management concern in the Highlights
🧠 New Feature in Public Preview: You can now enable Memory 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 AssistThe 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
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 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. Footnotes
|
|
@damccorm, Please have a look on this PR. Thanks |
|
Assigning reviewers: R: @chamikaramj added as fallback since no labels match configuration Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a much-needed LRU cache eviction mechanism to CachingStateProvider, preventing unbounded memory growth. The implementation correctly uses JavaScript's Map to maintain insertion order for LRU logic, and the accompanying tests are thorough. I've identified a bug in the promise-handling logic that could lead to incorrect cache eviction, along with a minor code style improvement. A single code suggestion is provided to address both issues. Overall, this is a solid contribution.
…ider - Fixed bug: removed incorrect evictIfNeeded() call in promise callback - Removed unnecessary this_ variable (arrow functions capture this) - Changed from count-based to size-based eviction (similar to Python statecache.py) - Added estimateSize() to calculate memory weight of cached values - Default cache weight: 100MB - Updated tests to work with weight-based eviction
| * Estimates the memory size of a value in bytes. | ||
| * This is a simplified estimation - actual memory usage may vary. | ||
| */ | ||
| function estimateSize(value: any): number { |
There was a problem hiding this comment.
Can we just use https://www.npmjs.com/package/object-sizeof?activeTab=readme ?
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a much-needed LRU cache eviction mechanism to CachingStateProvider, preventing unbounded memory growth. The implementation uses a weight-based approach with estimateSize to approximate memory usage and relies on JavaScript's Map insertion order to manage LRU logic. The changes are well-structured and include a comprehensive set of new tests.
My review includes a few key points:
- A critical fix for handling circular references in the
estimateSizefunction to prevent potential crashes. - A fix for a bug in the cache-miss logic that could cause the cache to exceed its maximum weight limit.
- A suggestion to improve the efficiency of the eviction logic.
- A recommendation to enhance the test suite to cover an important edge case.
Overall, this is a solid contribution that addresses a significant performance and stability concern.
|
Reminder, please take a look at this pr: @chamikaramj |
|
waiting on author |
github-actions
bot
added
Next Action: Author
and removed
Next Action: Reviewers
slow-review
labels
I am working on it, will update you soon. Thanks for understanding |
|
No problem, there's no rush. My comment just tells the bot not to assume the review is slow and reassign the PR |
- Fixed sizeof function to handle circular references using visited Set - Fixed eviction ordering: add to cache first, then evict (fixes edge case) - Added test for oversized item that exceeds maxCacheWeight - Implemented custom sizeof instead of object-sizeof package (has Node.js compatibility issues)
|
Hi @damccorm, I tried using object-sizeof as you suggested, but it has Node.js compatibility issues (util.isPrimitive was removed in Node v22+) and doesn't handle circular references. Instead, I implemented a custom sizeof function with a visited Set to handle circular references (addressing gemini's critical concern), fixed the eviction ordering bug (now add-then-evict), and added the oversized item test case. All 292 tests pass locally. Let me know if you'd prefer a different approach! |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a much-needed LRU cache eviction mechanism to CachingStateProvider, preventing unbounded memory growth. The implementation is based on a configurable maximum cache weight and leverages the insertion order of JavaScript Maps for LRU logic. The new sizeof utility for estimating object memory size is a nice addition, and the accompanying tests are thorough, covering various scenarios including cache hits/misses, LRU behavior, promise-based fetches, and edge cases.
I've found a critical race condition in the promise resolution logic that could lead to incorrect cache weight accounting. I've also suggested a minor refactoring to improve the efficiency and readability of the eviction logic. Overall, this is a great contribution that significantly improves the robustness of state management.
|
Thanks - just a few more gemini comments, otherwise this looks good to me |
- Fixed critical race condition in promise callback: only update cache if the entry is still the same promise we're resolving - Optimized evictIfNeeded: use entries() iterator and removed redundant checks
|
@damccorm, I have addressed Gemini comments: fix race condition, optimize evictIfNeeded
|
2 participants
Summary
Fixes #37213
Implements LRU (Least Recently Used) cache eviction for
CachingStateProviderto prevent unbounded memory growth in long-running workers.Problem
The
CachingStateProvidercurrently has unbounded cache growth, which can lead to memory issues. There's a TODO comment indicating this needs to be addressed:// TODO: (Perf) Cache eviction.## Solution
maxCacheSizeparameter (default: 1000 entries)Changes
Modified
sdks/typescript/src/apache_beam/worker/state.ts:maxCacheSizeconstructor parameterevictIfNeeded()methodtouchCacheEntry()method for LRU orderinggetState()to use eviction and LRU trackingAdded
sdks/typescript/test/state_provider_test.ts:Testing
npm testinsdks/typescriptCachingStateProviderThank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
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, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.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.