Skip to content
Open
15 changes: 13 additions & 2 deletions web/src/cache_state.ts
Original file line number Diff line number Diff line change
Expand Up @@ -145,11 +145,18 @@ export class CacheState {
* Invalidation rule: None required — shape tuples are immutable.
*/
readonly shapeCache: LRUCache<string, Disposable>;
/**
* 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<Disposable>();

constructor(shapeCacheSize: number = 256) {
this.shapeCache = new LRUCache<string, Disposable>(
shapeCacheSize,
(_key, value) => value.dispose()
(_key, value) => this.evictedShapeTuples.add(value)
);
}

Expand All @@ -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();
}
}
35 changes: 35 additions & 0 deletions web/tests/node/test_cache_state.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
/*
* 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 { CacheState } = require("../../src/cache_state.ts");

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);
});