refactor: expose memory stats accessor on idetik (#679)

summary:

add a read-only memoryStats accessor on idetik for baseline memory
instrumentation ahead of experimenting with shared residency.

the accessor reports:
- cpu-resident chunk bytes + chunk count
- gpu texture bytes + texture count
- js heap usage (chromium performance.memory, optional)

Author: shlomnissan

src/core/renderer.ts

@@ -60,6 +60,10 @@ export abstract class Renderer {
     return this.renderedObjects_;
   }
 
+  public abstract get gpuTextureBytes(): number;
+
+  public abstract get gpuTextureCount(): number;
+
   public set backgroundColor(color: ColorLike) {
     this.backgroundColor_ = Color.from(color);
   }

src/data/chunk_manager.ts

@@ -1,5 +1,6 @@
 import { ChunkSource } from "./chunk";
 import { ChunkQueue } from "./chunk_queue";
+import { chunkMemoryStats } from "./chunk_memory";
 
 export type QueueStats = {
   pending: number;
@@ -21,6 +22,10 @@ export class ChunkManager {
     };
   }
 
+  public get memoryStats() {
+    return chunkMemoryStats();
+  }
+
   public async addView(
     source: ChunkSource,
     policy: ImageSourcePolicy

src/data/chunk_memory.ts

@@ -0,0 +1,33 @@
+import { Chunk, ChunkData } from "./chunk";
+
+// Running totals for CPU-resident chunk voxel data. All mutations of
+// `chunk.data` must route through setChunkData/clearChunkData so these stay
+// accurate. Module-level (not per-store) because chunk loaders are created
+// per-source and have no back-reference to their store; this means the totals
+// are summed across all Idetik instances in the page, which is fine for the
+// single-instance debug/instrumentation use case.
+
+let cpuChunkBytes = 0;
+let cpuChunkCount = 0;
+
+export function setChunkData(chunk: Chunk, data: ChunkData): void {
+  if (chunk.data) {
+    cpuChunkBytes -= chunk.data.byteLength;
+    cpuChunkCount -= 1;
+  }
+  chunk.data = data;
+  cpuChunkBytes += data.byteLength;
+  cpuChunkCount += 1;
+}
+
+export function clearChunkData(chunk: Chunk): void {
+  if (chunk.data) {
+    cpuChunkBytes -= chunk.data.byteLength;
+    cpuChunkCount -= 1;
+  }
+  chunk.data = undefined;
+}
+
+export function chunkMemoryStats() {
+  return { cpuChunkBytes, cpuChunkCount };
+}

src/data/chunk_store.ts

@@ -1,4 +1,5 @@
 import { Chunk, SourceDimensionMap, ChunkLoader } from "./chunk";
+import { clearChunkData } from "./chunk_memory";
 import { almostEqual } from "../utilities/almost_equal";
 import { Logger } from "../utilities/logger";
 import { ChunkStoreView } from "./chunk_store_view";
@@ -223,7 +224,7 @@ export class ChunkStore {
         );
       }
 
-      chunk.data = undefined;
+      clearChunkData(chunk);
       chunk.state = "unloaded";
       chunk.orderKey = null;
       Logger.debug(

src/data/ome_zarr/image_loader.ts

@@ -1,6 +1,7 @@
 import * as zarr from "zarrita";
 
 import { Chunk, SourceDimension, SourceDimensionMap } from "../chunk";
+import { setChunkData } from "../chunk_memory";
 import { isTextureUnpackRowAlignment } from "../../objects/textures/texture";
 
 import { Image as OmeZarrImage } from "./0.5/image";
@@ -96,7 +97,7 @@ export class OmeZarrImageLoader {
       );
     }
     chunk.rowAlignmentBytes = rowAlignment;
-    chunk.data = data;
+    setChunkData(chunk, data);
   }
 }
 

src/idetik.ts

@@ -27,6 +27,17 @@ export type IdetikContext = {
   chunkManager: ChunkManager;
 };
 
+// JS heap usage from the non-standard performance.memory API.
+// Chromium only, undefined where the API is unavailable.
+type MemoryStats = {
+  cpuChunkBytes: number;
+  cpuChunkCount: number;
+  gpuTextureBytes: number;
+  gpuTextureCount: number;
+  jsHeapUsedBytes?: number;
+  jsHeapLimitBytes?: number;
+};
+
 export class Idetik {
   private readonly chunkManager_: ChunkManager;
   private readonly context_: IdetikContext;
@@ -137,6 +148,22 @@ export class Idetik {
     return this.chunkManager_.queueStats;
   }
 
+  public get memoryStats(): MemoryStats {
+    const perf = (
+      performance as Performance & {
+        memory?: { usedJSHeapSize: number; jsHeapSizeLimit: number };
+      }
+    ).memory;
+
+    return {
+      ...this.chunkManager_.memoryStats,
+      gpuTextureBytes: this.renderer_.gpuTextureBytes,
+      gpuTextureCount: this.renderer_.gpuTextureCount,
+      jsHeapUsedBytes: perf?.usedJSHeapSize,
+      jsHeapLimitBytes: perf?.jsHeapSizeLimit,
+    };
+  }
+
   public get renderedObjects() {
     return this.renderer_.renderedObjects;
   }

src/renderers/webgl_renderer.ts

@@ -61,6 +61,14 @@ export class WebGLRenderer extends Renderer {
     this.resize(this.canvas.width, this.canvas.height);
   }
 
+  public get gpuTextureBytes() {
+    return this.textures_.gpuTextureBytes;
+  }
+
+  public get gpuTextureCount() {
+    return this.textures_.textureCount;
+  }
+
   public render(viewport: Viewport) {
     this.renderedObjects_ = 0;
     this.renderedObjectsPerFrame_ = 0;

src/renderers/webgl_textures.ts

@@ -30,6 +30,14 @@ export class WebGLTextures {
     this.maxTextureUnits_ = gl.getParameter(gl.MAX_TEXTURE_IMAGE_UNITS);
   }
 
+  public get gpuTextureBytes() {
+    return this.gpuTextureBytes_;
+  }
+
+  public get textureCount() {
+    return this.textureCount_;
+  }
+
   public bindTexture(texture: Texture, index: number) {
     if (this.alreadyActive(texture)) return;
 

src/renderers/webgpu/webgpu_renderer.ts

@@ -441,6 +441,15 @@ class WebGPURenderer extends Renderer {
       projectionMatrix
     );
   }
+
+  // The experimental WebGPU texture pool does not track byte usage yet.
+  public get gpuTextureBytes() {
+    return 0;
+  }
+
+  public get gpuTextureCount() {
+    return 0;
+  }
 }
 
 function shaderNameForTexture(texture: Texture): ShaderName {