From 5806855bb549c5b70b78e9d46a77af9f5aa7112b Mon Sep 17 00:00:00 2001 From: Atirna <288419661+atirna@users.noreply.github.com> Date: Sat, 15 Aug 2026 20:16:14 +0530 Subject: [PATCH] [Web] Retain evicted shape tuples until teardown Signed-off-by: Atirna <288419661+atirna@users.noreply.github.com> --- web/src/cache_state.ts | 15 +++++++-- web/tests/node/test_cache_state.js | 54 ++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 2 deletions(-) create mode 100644 web/tests/node/test_cache_state.js diff --git a/web/src/cache_state.ts b/web/src/cache_state.ts index c571a071b838..0baaed6d3251 100644 --- a/web/src/cache_state.ts +++ b/web/src/cache_state.ts @@ -145,11 +145,18 @@ export class CacheState { * Invalidation rule: None required — shape tuples are immutable. */ readonly shapeCache: LRUCache; + /** + * Shape tuples evicted from the LRU cache but retained until instance teardown. + * + * `makeShapeTuple()` returns borrowed references, so an evicted tuple can still + * be in use by a pending GPU dispatch. + */ + private readonly evictedShapeTuples = new Set(); constructor(shapeCacheSize: number = 256) { this.shapeCache = new LRUCache( shapeCacheSize, - (_key, value) => value.dispose() + (_key, value) => this.evictedShapeTuples.add(value) ); } @@ -168,8 +175,12 @@ export class CacheState { */ dispose(): void { for (const obj of this.shapeCache.values()) { - obj.dispose(); + this.evictedShapeTuples.add(obj); } this.shapeCache.invalidate(); + for (const obj of this.evictedShapeTuples) { + obj.dispose(); + } + this.evictedShapeTuples.clear(); } } diff --git a/web/tests/node/test_cache_state.js b/web/tests/node/test_cache_state.js new file mode 100644 index 000000000000..b07c180c804c --- /dev/null +++ b/web/tests/node/test_cache_state.js @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +const fs = require("fs"); +const ts = require("typescript"); + +const previousTypeScriptLoader = require.extensions[".ts"]; +require.extensions[".ts"] = (module, filename) => { + const source = fs.readFileSync(filename, "utf8"); + const output = ts.transpileModule(source, { + compilerOptions: { module: ts.ModuleKind.CommonJS, target: ts.ScriptTarget.ES2018 }, + fileName: filename, + }).outputText; + module._compile(output, filename); +}; + +const { CacheState } = require("../../src/cache_state.ts"); + +if (previousTypeScriptLoader === undefined) { + delete require.extensions[".ts"]; +} else { + require.extensions[".ts"] = previousTypeScriptLoader; +} + +test("keeps an evicted shape tuple alive until cache state disposal", () => { + const cacheState = new CacheState(1); + const first = { dispose: jest.fn() }; + const second = { dispose: jest.fn() }; + + cacheState.shapeCache.get("first", () => first); + cacheState.shapeCache.get("second", () => second); + + expect(first.dispose).not.toHaveBeenCalled(); + + cacheState.dispose(); + + expect(first.dispose).toHaveBeenCalledTimes(1); + expect(second.dispose).toHaveBeenCalledTimes(1); +});