fix: bound prepared splat runtime cache

Author: Sportinger

src/engine/three/splatRuntimeCache.ts

@@ -15,6 +15,7 @@ const RUNTIME_MAGIC = 0x53475254; // SGRT
 const RUNTIME_VERSION = 2;
 const HEADER_BYTE_LENGTH = 22 * 4;
 export const DEFAULT_SPLAT_BASE_LOD_MAX_SPLATS = 65536;
+export const DEFAULT_PREPARED_SPLAT_RUNTIME_CACHE_MAX_BYTES = 256 * 1024 * 1024;
 
 export interface PreparedSplatRuntime {
   runtimeKey: string;
@@ -54,7 +55,10 @@ interface RuntimeRequestOptions extends RuntimeSourceOptions {
 const assetPromiseCache = new Map<string, Promise<GaussianSplatAsset>>();
 const runtimePromiseCache = new Map<string, Promise<PreparedSplatRuntime>>();
 const runtimeValueCache = new Map<string, PreparedSplatRuntime>();
+const runtimeValueSizeCache = new Map<string, number>();
 const idlePrewarmKeys = new Set<string>();
+let preparedRuntimeCacheMaxBytes = DEFAULT_PREPARED_SPLAT_RUNTIME_CACHE_MAX_BYTES;
+let preparedRuntimeCacheBytes = 0;
 
 function normalizeRequestedMaxSplats(requestedMaxSplats: number | undefined): number {
   const normalized = Math.floor(requestedMaxSplats ?? 0);
@@ -119,6 +123,76 @@ function cloneViewBytes(view: Float32Array): ArrayBuffer {
   return bytes.buffer as ArrayBuffer;
 }
 
+function getPreparedRuntimeByteLength(runtime: PreparedSplatRuntime): number {
+  return runtime.centers.byteLength
+    + runtime.centerOpacityTextureData.byteLength
+    + runtime.colorTextureData.byteLength
+    + runtime.axisXTextureData.byteLength
+    + runtime.axisYTextureData.byteLength
+    + runtime.axisZTextureData.byteLength
+    + runtime.orderTemplateData.byteLength;
+}
+
+function deletePreparedRuntime(runtimeKey: string): void {
+  const size = runtimeValueSizeCache.get(runtimeKey);
+  if (typeof size === 'number') {
+    preparedRuntimeCacheBytes = Math.max(0, preparedRuntimeCacheBytes - size);
+    runtimeValueSizeCache.delete(runtimeKey);
+  }
+  runtimeValueCache.delete(runtimeKey);
+}
+
+function touchPreparedRuntime(runtimeKey: string): PreparedSplatRuntime | null {
+  const runtime = runtimeValueCache.get(runtimeKey);
+  if (!runtime) {
+    return null;
+  }
+
+  runtimeValueCache.delete(runtimeKey);
+  runtimeValueCache.set(runtimeKey, runtime);
+  return runtime;
+}
+
+function trimPreparedRuntimeCache(): void {
+  while (preparedRuntimeCacheBytes > preparedRuntimeCacheMaxBytes && runtimeValueCache.size > 0) {
+    const oldestKey = runtimeValueCache.keys().next().value;
+    if (oldestKey === undefined) {
+      break;
+    }
+    deletePreparedRuntime(oldestKey);
+  }
+}
+
+function storePreparedRuntime(runtime: PreparedSplatRuntime): PreparedSplatRuntime {
+  const size = getPreparedRuntimeByteLength(runtime);
+  deletePreparedRuntime(runtime.runtimeKey);
+
+  if (size > preparedRuntimeCacheMaxBytes) {
+    log.debug('Skipping prepared gaussian splat runtime cache for oversize entry', {
+      runtimeKey: runtime.runtimeKey,
+      size,
+      preparedRuntimeCacheMaxBytes,
+    });
+    return runtime;
+  }
+
+  runtimeValueCache.set(runtime.runtimeKey, runtime);
+  runtimeValueSizeCache.set(runtime.runtimeKey, size);
+  preparedRuntimeCacheBytes += size;
+  trimPreparedRuntimeCache();
+  return touchPreparedRuntime(runtime.runtimeKey) ?? runtime;
+}
+
+function getRuntimeKeyForOptions(options: RuntimeRequestOptions): string {
+  const requestedMaxSplats = normalizeRequestedMaxSplats(options.requestedMaxSplats);
+  return buildRuntimeKey(
+    options.cacheKey,
+    options.variant,
+    requestedMaxSplats,
+    options.gaussianSplatSequence,
+  );
+}
+
 function gcd(a: number, b: number): number {
   let x = Math.abs(Math.floor(a));
   let y = Math.abs(Math.floor(b));
@@ -199,7 +273,11 @@ async function loadAsset(options: RuntimeSourceOptions): Promise<GaussianSplatAs
   })();
 
   assetPromiseCache.set(options.cacheKey, promise);
-  void promise.catch(() => {
+  void promise.then(() => {
+    if (assetPromiseCache.get(options.cacheKey) === promise) {
+      assetPromiseCache.delete(options.cacheKey);
+    }
+  }, () => {
     if (assetPromiseCache.get(options.cacheKey) === promise) {
       assetPromiseCache.delete(options.cacheKey);
     }
@@ -504,8 +582,7 @@ async function loadRuntimeFromProjectCache(
   try {
     const buffer = await file.arrayBuffer();
     const runtime = deserializeRuntime(runtimeKey, buffer);
-    runtimeValueCache.set(runtimeKey, runtime);
-    return runtime;
+    return storePreparedRuntime(runtime);
   } catch (error) {
     log.warn('Failed to read gaussian splat runtime cache file', {
       fileHash,
@@ -542,13 +619,8 @@ async function persistRuntimeToProjectCache(
 
 async function ensurePreparedRuntime(options: RuntimeRequestOptions): Promise<PreparedSplatRuntime> {
   const requestedMaxSplats = normalizeRequestedMaxSplats(options.requestedMaxSplats);
-  const runtimeKey = buildRuntimeKey(
-    options.cacheKey,
-    options.variant,
-    requestedMaxSplats,
-    options.gaussianSplatSequence,
-  );
-  const existing = runtimeValueCache.get(runtimeKey);
+  const runtimeKey = getRuntimeKeyForOptions(options);
+  const existing = touchPreparedRuntime(runtimeKey);
   if (existing) return existing;
 
   const existingPromise = runtimePromiseCache.get(runtimeKey);
@@ -574,13 +646,17 @@ async function ensurePreparedRuntime(options: RuntimeRequestOptions): Promise<Pr
       requestedMaxSplats,
       normalizationBounds,
     );
-    runtimeValueCache.set(runtimeKey, runtime);
+    storePreparedRuntime(runtime);
     void persistRuntimeToProjectCache(options.fileHash, runtime);
     return runtime;
   })();
 
   runtimePromiseCache.set(runtimeKey, promise);
-  void promise.catch(() => {
+  void promise.then(() => {
+    if (runtimePromiseCache.get(runtimeKey) === promise) {
+      runtimePromiseCache.delete(runtimeKey);
+    }
+  }, () => {
     if (runtimePromiseCache.get(runtimeKey) === promise) {
       runtimePromiseCache.delete(runtimeKey);
     }
@@ -591,10 +667,34 @@ async function ensurePreparedRuntime(options: RuntimeRequestOptions): Promise<Pr
 export function getPreparedSplatRuntimeSync(
   options: RuntimeRequestOptions,
 ): PreparedSplatRuntime | null {
-  const requestedMaxSplats = normalizeRequestedMaxSplats(options.requestedMaxSplats);
-  return runtimeValueCache.get(
-    buildRuntimeKey(options.cacheKey, options.variant, requestedMaxSplats, options.gaussianSplatSequence),
-  ) ?? null;
+  return touchPreparedRuntime(getRuntimeKeyForOptions(options));
+}
+
+export function clearPreparedSplatRuntimeCache(): void {
+  assetPromiseCache.clear();
+  runtimePromiseCache.clear();
+  runtimeValueCache.clear();
+  runtimeValueSizeCache.clear();
+  idlePrewarmKeys.clear();
+  preparedRuntimeCacheBytes = 0;
+}
+
+export function setPreparedSplatRuntimeCacheMaxBytes(maxBytes: number): void {
+  const normalized = Math.max(0, Math.floor(maxBytes));
+  preparedRuntimeCacheMaxBytes = normalized;
+  trimPreparedRuntimeCache();
+}
+
+export function getPreparedSplatRuntimeCacheStats(): {
+  runtimeCount: number;
+  totalBytes: number;
+  maxBytes: number;
+} {
+  return {
+    runtimeCount: runtimeValueCache.size,
+    totalBytes: preparedRuntimeCacheBytes,
+    maxBytes: preparedRuntimeCacheMaxBytes,
+  };
 }
 
 export async function resolvePreparedSplatRuntime(

tests/unit/splatRuntimeCache.test.ts

@@ -168,4 +168,60 @@ describe('splatRuntimeCache', () => {
     expect(runtime.centers[1]).toBeCloseTo(0);
     expect(runtime.centers[2]).toBeCloseTo(0);
   });
+
+  it('evicts prepared runtimes by byte budget and reloads assets after eviction', async () => {
+    const frameA = new File([new Uint8Array([1])], 'frame0000000.ply', {
+      type: 'application/octet-stream',
+    });
+    const frameB = new File([new Uint8Array([2])], 'frame0000001.ply', {
+      type: 'application/octet-stream',
+    });
+
+    loadGaussianSplatAssetMock.mockImplementation(async (file: File) => (
+      createAsset(file, {
+        center: file.name === frameA.name ? [1, 0, 0] : [2, 0, 0],
+      })
+    ));
+
+    const {
+      getPreparedSplatRuntimeSync,
+      setPreparedSplatRuntimeCacheMaxBytes,
+      waitForBasePreparedSplatRuntime,
+    } = await import('../../src/engine/three/splatRuntimeCache');
+
+    setPreparedSplatRuntimeCacheMaxBytes(120);
+
+    await waitForBasePreparedSplatRuntime({
+      cacheKey: 'Raw/frame0000000.ply',
+      file: frameA,
+      fileName: frameA.name,
+    });
+    expect(getPreparedSplatRuntimeSync({
+      cacheKey: 'Raw/frame0000000.ply',
+      file: frameA,
+      fileName: frameA.name,
+      variant: 'base',
+    })).not.toBeNull();
+
+    await waitForBasePreparedSplatRuntime({
+      cacheKey: 'Raw/frame0000001.ply',
+      file: frameB,
+      fileName: frameB.name,
+    });
+
+    expect(getPreparedSplatRuntimeSync({
+      cacheKey: 'Raw/frame0000000.ply',
+      file: frameA,
+      fileName: frameA.name,
+      variant: 'base',
+    })).toBeNull();
+
+    await waitForBasePreparedSplatRuntime({
+      cacheKey: 'Raw/frame0000000.ply',
+      file: frameA,
+      fileName: frameA.name,
+    });
+
+    expect(loadGaussianSplatAssetMock).toHaveBeenCalledTimes(3);
+  });
 });