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.
| function estimateSize(value: any): number { | ||
| if (value === null || value === undefined) { | ||
| return 8; | ||
| } | ||
|
|
||
| const type = typeof value; | ||
|
|
||
| if (type === "boolean") { | ||
| return 4; | ||
| } | ||
| if (type === "number") { | ||
| return 8; | ||
| } | ||
| if (type === "string") { | ||
| // Each character is 2 bytes in JavaScript (UTF-16) + overhead | ||
| return 40 + value.length * 2; | ||
| } | ||
| if (value instanceof Uint8Array || value instanceof Buffer) { | ||
| return 40 + value.length; | ||
| } | ||
| if (Array.isArray(value)) { | ||
| let size = 40; // Array overhead | ||
| for (const item of value) { | ||
| size += estimateSize(item); | ||
| } | ||
| return size; | ||
| } | ||
| if (type === "object") { | ||
| let size = 40; // Object overhead | ||
| for (const key of Object.keys(value)) { | ||
| size += estimateSize(key) + estimateSize(value[key]); | ||
| } | ||
| return size; | ||
| } | ||
|
|
||
| // Default for unknown types | ||
| return 64; | ||
| } |
There was a problem hiding this comment.
The recursive implementation of estimateSize for objects and arrays does not handle circular references. If a cached object contains a cycle (e.g., a.b = a), this function will enter an infinite loop, leading to a stack overflow and crashing the worker. This is a critical stability risk.
function estimateSize(value: any, visited = new Set<any>()): number {
if (value === null || value === undefined) {
return 8;
}
if (typeof value === "object") {
if (visited.has(value)) {
return 8; // Account for reference size, not the full object again.
}
visited.add(value);
}
const type = typeof value;
if (type === "boolean") {
return 4;
}
if (type === "number") {
return 8;
}
if (type === "string") {
// Each character is 2 bytes in JavaScript (UTF-16) + overhead
return 40 + value.length * 2;
}
if (value instanceof Uint8Array || value instanceof Buffer) {
return 40 + value.length;
}
if (Array.isArray(value)) {
let size = 40; // Array overhead
for (const item of value) {
size += estimateSize(item, visited);
}
return size;
}
if (type === "object") {
let size = 40; // Object overhead
for (const key of Object.keys(value)) {
size += estimateSize(key, visited) + estimateSize(value[key], visited);
}
return size;
}
// Default for unknown types
return 64;
}| // Evict if needed before adding new entry | ||
| this.currentWeight += weight; | ||
| this.evictIfNeeded(); | ||
| this.cache.set(cacheKey, { entry: result, weight }); |
There was a problem hiding this comment.
There is a bug in the ordering of operations for a synchronous cache miss. evictIfNeeded() is called before the new item is added to the cache via this.cache.set(). If the cache is empty and a new item's weight exceeds maxCacheWeight, evictIfNeeded() will not run (as this.cache.size is 0). The oversized item is then added, causing the cache to violate its weight limit until the next eviction event. The item should be added to the cache before eviction is checked.
// Add new entry and then evict if needed
this.currentWeight += weight;
this.cache.set(cacheKey, { entry: result, weight });
this.evictIfNeeded();| private evictIfNeeded() { | ||
| while (this.currentWeight > this.maxCacheWeight && this.cache.size > 0) { | ||
| // Remove the first (oldest) entry | ||
| const firstKey = this.cache.keys().next().value; | ||
| if (firstKey !== undefined) { | ||
| const evicted = this.cache.get(firstKey); | ||
| if (evicted !== undefined) { | ||
| this.currentWeight -= evicted.weight; | ||
| } | ||
| this.cache.delete(firstKey); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
The evictIfNeeded method can be made more efficient and readable. Currently, it gets the first key using this.cache.keys().next().value and then performs a separate this.cache.get(firstKey) lookup. You can achieve this in a single operation by iterating over this.cache.entries().
private evictIfNeeded() {
while (this.currentWeight > this.maxCacheWeight && this.cache.size > 0) {
// JavaScript Maps preserve insertion order, so the first entry is the oldest.
const [firstKey, evicted] = this.cache.entries().next().value;
this.currentWeight -= evicted.weight;
this.cache.delete(firstKey);
}
}| } | ||
| } | ||
|
|
||
| describe("CachingStateProvider", function () { |
There was a problem hiding this comment.
The test suite is quite thorough. To make it even more robust, consider adding a test case for the edge case where a single item, whose weight is larger than maxCacheWeight, is added to an empty cache. This would verify that the cache's weight limit is strictly enforced. With the recommended fix in state.ts, such an item should be added and then immediately evicted, leaving the cache empty and its weight at zero.
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.