From 3715f6855af0bbc37330e57acba4635f57c5111b Mon Sep 17 00:00:00 2001 From: Marco Salazar Date: Fri, 20 Dec 2024 17:01:15 -0500 Subject: [PATCH] fix --- .../src/util/__tests__/weakMapMemoize.test.ts | 76 +++++++++++++++++++ .../ui-core/src/util/weakMapMemoize.ts | 46 +++++++---- 2 files changed, 107 insertions(+), 15 deletions(-) diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts index c99ca26e1f43e..40b754bd5975b 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts +++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts @@ -248,4 +248,80 @@ describe('weakMapMemoize', () => { expect(memoizedFn(1, 2)).toBe(3); expect(spy).toHaveBeenCalledTimes(6); }); + + // Test 14: Evicting a cache node without children deletes its result + it('should delete the result of an evicted cache node without children', () => { + const spy = jest.fn((a: number, b: number) => a + b); + const memoizedFn = weakMapMemoize(spy, {maxEntries: 2}); + + expect(memoizedFn(1, 1)).toBe(2); + expect(memoizedFn(2, 2)).toBe(4); + expect(spy).toHaveBeenCalledTimes(2); // Two unique calls + + // This should evict the least recently used entry (1,1) + expect(memoizedFn(3, 3)).toBe(6); + expect(spy).toHaveBeenCalledTimes(3); // New call for (3,3) + + // Accessing (1,1) again should trigger a new function call since it was evicted, and evict (2, 2) + expect(memoizedFn(1, 1)).toBe(2); + expect(spy).toHaveBeenCalledTimes(4); // New call for (1,1) + + // evicts 3,3 + expect(memoizedFn(2, 2)).toBe(4); + expect(spy).toHaveBeenCalledTimes(5); + + expect(memoizedFn(3, 3)).toBe(6); + expect(spy).toHaveBeenCalledTimes(6); + }); + + // Test 15: Evicting a cache node with children deletes only its result, children remain cached + it('should delete only the result of an evicted cache node with children and keep children cached', () => { + const spy = jest.fn((a: number, b: string, c?: number) => a + b.length + (c ?? 0)); + const memoizedFn = weakMapMemoize(spy, {maxEntries: 3}); + + // Creating cache nodes: + // (1, 'a') + // (1, 'a', 10) + // (2, 'b', 30) + + expect(memoizedFn(1, 'a')).toBe(1 + 1); + expect(memoizedFn(1, 'a', 10)).toBe(1 + 1 + 10); + expect(memoizedFn(2, 'b', 30)).toBe(2 + 1 + 30); + expect(spy).toHaveBeenCalledTimes(3); // Three unique calls + + // At this point, maxEntries is 3: + // Cache contains: + // (1, 'a'), (1, 'a', 10), (2, 'b', 30) + + // Adding a new entry should evict the least recently used, which is (1, 'a') + expect(memoizedFn(1, 'c', 40)).toBe(1 + 1 + 40); + expect(spy).toHaveBeenCalledTimes(4); + + // Now, cache contains: + // (1, 'a', 10), (2, 'b', 30), (1, 'c', 40) + expect(memoizedFn(1, 'a', 10)).toBe(1 + 1 + 10); + expect(spy).toHaveBeenCalledTimes(4); // no new call + + // Now, cache contains: + // (2, 'b', 30), (1, 'c', 40), (1, 'a', 10) + // Evicting (2, 'b', 30) + expect(memoizedFn(4, 'd', 50)).toBe(4 + 1 + 50); + expect(spy).toHaveBeenCalledTimes(5); + + // Now, cache contains: + // (1, 'c', 40), (1, 'a', 10), (4, 'd', 50) + + // Accessing (1, 'a', 10) again should trigger a new call since it was evicted and evict (2, b, 30) + expect(memoizedFn(2, 'b', 30)).toBe(2 + 1 + 30); + expect(spy).toHaveBeenCalledTimes(6); + + // Now, cache contains: + // (1, 'a', 10), (4, 'd', 50), (2, 'b', 30) + + expect(memoizedFn(1, 'a', 10)).toBe(12); + expect(memoizedFn(4, 'd', 50)).toBe(55); + expect(memoizedFn(2, 'b', 30)).toBe(33); + + expect(spy).toHaveBeenCalledTimes(6); + }); }); diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts index a7d79c71c7307..ec328e445d71d 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts +++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts @@ -25,6 +25,32 @@ function isObject(value: any): value is object { return value !== null && (typeof value === 'object' || typeof value === 'function'); } +/** + * Recursively deletes parent nodes if their childCount reaches zero. + * @param cacheNode - The cache node to start deletion from. + */ +function recursivelyDeleteParent(cacheNode: CacheNode) { + if (cacheNode.parent && cacheNode.parentKey !== undefined) { + const parent = cacheNode.parent; + const parentKey = cacheNode.parentKey; + + // Remove the current cacheNode from its parent + if (isObject(parentKey)) { + parent.weakMap.delete(parentKey); + } else { + parent.map.delete(parentKey); + } + + // Decrement the parent's child count + parent.childCount--; + + // If the parent's childCount reaches zero and it's not the root, recurse + if (parent.childCount === 0 && parent.parent) { + recursivelyDeleteParent(parent); + } + } +} + /** * Memoizes a function using nested Maps and WeakMaps based on the arguments. * Optionally limits the number of cached entries using an LRU cache. @@ -49,23 +75,13 @@ export function weakMapMemoize(fn: T, options?: WeakMapMe if (maxEntries) { lruCache = new LRU({ max: maxEntries, - dispose: (key, cacheNode) => { - // Remove the cached entry + dispose: (_key, cacheNode) => { + // Remove the cached result delete cacheNode.result; - // If there are no entries in this node then lets remove it. + // If there are no child nodes, proceed to remove this node from its parent if (cacheNode.childCount === 0 && cacheNode.parent && cacheNode.parentKey !== undefined) { - const parent = cacheNode.parent; - - // Decrement the parent's child count - parent.childCount--; - const parentKey = cacheNode.parentKey; - - if (isObject(parentKey)) { - parent.weakMap.delete(parentKey); - } else { - parent.map.delete(parentKey); - } + recursivelyDeleteParent(cacheNode); } }, noDisposeOnSet: false, // Ensure dispose is called on eviction @@ -131,7 +147,7 @@ export function weakMapMemoize(fn: T, options?: WeakMapMe // If LRU cache is enabled, manage the cache entries if (lruCache && !currentCache.lruKey) { - const cacheEntryKey: any[] = path.slice(); // Clone the path to ensure uniqueness + const cacheEntryKey = Symbol(); currentCache.lruKey = cacheEntryKey; // Associate the cache node with the LRU key lruCache.set(cacheEntryKey, currentCache); }