fix

dagster-io · Dec 20, 2024 · 80e4ac3 · 80e4ac3
1 parent f40c8e4
commit 80e4ac3
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 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
@@ -174,8 +174,8 @@ describe('weakMapMemoize', () => {
     const object2 = {b: 10};
     const object3 = {c: 15};
     const object4 = {a: 30};
-    const object5 = {b: 20}; // Corrected to have 'b' property
-    const object6 = {c: 25}; // Corrected to have 'c' property
+    const object5 = {b: 20};
+    const object6 = {c: 25};
 
     expect(memoizedFn(object1, 'test', object2, true, object3, 20)).toBe(5 + 4 + 10 + 1 + 15 + 20); // 55
     expect(memoizedFn(object4, 'test', object2, true, object3, 20)).toBe(30 + 4 + 10 + 1 + 15 + 20); // 80

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
@@ -13,6 +13,7 @@ interface CacheNode {
   parent?: CacheNode;
   parentKey?: any;
   lruKey?: any; // Reference to the key in the LRU cache
+  childCount: number; // Number of child nodes
 }
 
 /**
@@ -39,6 +40,7 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
   const cacheRoot: CacheNode = {
     map: new Map(),
     weakMap: new WeakMap(),
+    childCount: 0,
   };
 
   // Initialize LRU Cache if maxEntries is specified
@@ -48,9 +50,15 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
     lruCache = new LRU<any, CacheNode>({
       max: maxEntries,
       dispose: (key, cacheNode) => {
-        // When an entry is evicted from the LRU cache, remove it from its parent
-        if (cacheNode.parent && cacheNode.parentKey !== undefined) {
+        // Remove the cached entry
+        delete cacheNode.result;
+
+        // If there are no entries in this node then lets remove it.
+        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)) {
@@ -82,8 +90,10 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
             weakMap: new WeakMap(),
             parent: currentCache,
             parentKey: arg,
+            childCount: 0,
           };
           currentCache.weakMap.set(arg, newCacheNode);
+          currentCache.childCount++;
         }
         nextCacheNode = currentCache.weakMap.get(arg);
       } else {
@@ -93,8 +103,10 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
             weakMap: new WeakMap(),
             parent: currentCache,
             parentKey: arg,
+            childCount: 0,
           };
           currentCache.map.set(arg, newCacheNode);
+          currentCache.childCount++;
         }
         nextCacheNode = currentCache.map.get(arg);
       }
@@ -118,10 +130,9 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
     currentCache.result = result;
 
     // If LRU cache is enabled, manage the cache entries
-    if (lruCache) {
+    if (lruCache && !currentCache.lruKey) {
       const cacheEntryKey: any[] = path.slice(); // Clone the path to ensure uniqueness
       currentCache.lruKey = cacheEntryKey; // Associate the cache node with the LRU key
-
       lruCache.set(cacheEntryKey, currentCache);
     }