feat: Add optional layer-owned runtime cleanup

src/chunk_manager/frontend.ts

@@ -332,7 +332,11 @@ export class ChunkQueueManager extends SharedObject {
 }
 
 function updateChunk(rpc: RPC, x: any) {
-  const source: ChunkSource = rpc.get(x.source);
+  const source = rpc.get(x.source) as ChunkSource | undefined;
+  if (source === undefined) {
+    // Source was removed while chunk update was in flight.
+    return;
+  }
   if (DEBUG_CHUNK_UPDATES) {
     console.log(
       `${Date.now()} Chunk.update received: ` +

src/datasource/catmaid/frontend.ts

@@ -476,6 +476,9 @@ export class CatmaidDataSourceProvider implements DataSourceProvider {
         id: "skeletons-chunked",
         default: true,
         subsource: { mesh: multiscaleSource },
+        layerRuntimeStateDisposal: {
+          kind: "spatiallyIndexedSkeleton",
+        },
       },
       {
         id: "skeletons",

src/datasource/index.ts

@@ -153,6 +153,10 @@ export interface CompleteUrlOptions extends CompleteUrlOptionsBase {
   signal: AbortSignal;
 }
 
+export interface LayerRuntimeStateDisposalRequest {
+  kind: string;
+}
+
 export interface DataSubsourceEntry {
   /**
    * Unique identifier (within the group) for this subsource.  Stored in the JSON state
@@ -182,6 +186,12 @@ export interface DataSubsourceEntry {
    * Specifies whether this associated data source is enabled by default.
    */
   default: boolean;
+
+  /**
+   * Optional layer-owned runtime cleanup requested when this active subsource's
+   * datasource is replaced or cleared.
+   */
+  layerRuntimeStateDisposal?: LayerRuntimeStateDisposalRequest;
 }
 
 export interface ChannelMetadata {

src/layer/index.ts

@@ -38,7 +38,10 @@ import type {
 } from "#src/datasource/index.js";
 import { makeEmptyDataSourceSpecification } from "#src/datasource/index.js";
 import type { DisplayContext, RenderedPanel } from "#src/display_context.js";
-import type { LoadedDataSubsource } from "#src/layer/layer_data_source.js";
+import type {
+  LayerDataSourceChangeRuntimeDisposalContext,
+  LoadedDataSubsource,
+} from "#src/layer/layer_data_source.js";
 import {
   LayerDataSource,
   layerDataSourceSpecificationFromJson,
@@ -429,6 +432,14 @@ export class UserLayer extends RefCounted {
     subsources;
   }
 
+  // Derived classes may override to clear layer-owned runtime state for active
+  // datasources that explicitly request cleanup on source change.
+  disposeLayerRuntimeStateForDataSourceChange(
+    _context: LayerDataSourceChangeRuntimeDisposalContext,
+  ) {
+    return false;
+  }
+
   updateDataSubsourceActivations() {
     function* getDataSubsources(
       this: UserLayer,

src/layer/layer_data_source.ts

@@ -31,6 +31,7 @@ import type {
   DataSourceWithRedirectInfo,
   DataSubsourceEntry,
   DataSubsourceSpecification,
+  LayerRuntimeStateDisposalRequest,
 } from "#src/datasource/index.js";
 import { makeEmptyDataSourceSpecification } from "#src/datasource/index.js";
 import type { UserLayer } from "#src/layer/index.js";
@@ -148,6 +149,7 @@ export class LoadedDataSubsource {
   enabled: boolean;
   activated: RefCounted | undefined = undefined;
   guardValues: any[] = [];
+  renderLayers = new Set<RenderLayer>();
   messages = new MessageList();
   isActiveChanged = new NullarySignal();
   constructor(
@@ -212,9 +214,13 @@ export class LoadedDataSubsource {
 
   addRenderLayer(renderLayer: Owned<RenderLayer>) {
     const activated = this.activated!;
-    activated.registerDisposer(
-      this.loadedDataSource.layer.addRenderLayer(renderLayer),
-    );
+    const removeRenderLayer =
+      this.loadedDataSource.layer.addRenderLayer(renderLayer);
+    this.renderLayers.add(renderLayer);
+    activated.registerDisposer(() => {
+      this.renderLayers.delete(renderLayer);
+      removeRenderLayer();
+    });
     activated.registerDisposer(this.messages.addChild(renderLayer.messages));
   }
 
@@ -301,6 +307,16 @@ export class LoadedLayerDataSource extends RefCounted {
   }
 }
 
+export type LayerDataSourceChangeReason = "replace" | "clear";
+
+export interface LayerDataSourceChangeRuntimeDisposalContext {
+  request: LayerRuntimeStateDisposalRequest;
+  reason: LayerDataSourceChangeReason;
+  layerDataSource: LayerDataSource;
+  loadedDataSource: LoadedLayerDataSource;
+  loadedSubsource: LoadedDataSubsource;
+}
+
 export type LayerDataSourceLoadState =
   | {
       error: Error;
@@ -368,10 +384,36 @@ export class LayerDataSource extends RefCounted {
     return this.loadState_;
   }
 
+  private disposeRuntimeStateForDataSourceChange(
+    reason: LayerDataSourceChangeReason,
+  ) {
+    const { loadState } = this;
+    if (loadState === undefined || loadState.error !== undefined) return false;
+    const handledRequestKinds = new Set<string>();
+    let changed = false;
+    for (const loadedSubsource of loadState.subsources) {
+      if (loadedSubsource.activated === undefined) continue;
+      const request = loadedSubsource.subsourceEntry.layerRuntimeStateDisposal;
+      if (request === undefined) continue;
+      if (handledRequestKinds.has(request.kind)) continue;
+      handledRequestKinds.add(request.kind);
+      changed =
+        this.layer.disposeLayerRuntimeStateForDataSourceChange({
+          request,
+          reason,
+          layerDataSource: this,
+          loadedDataSource: loadState,
+          loadedSubsource,
+        }) || changed;
+    }
+    return changed;
+  }
+
   set spec(spec: DataSourceSpecification) {
     const { layer } = this;
     this.messages.clearMessages();
     if (spec.url.length === 0) {
+      this.disposeRuntimeStateForDataSourceChange("clear");
       if (layer.dataSources.length !== 1) {
         const index = layer.dataSources.indexOf(this);
         if (index !== -1) {
@@ -395,6 +437,7 @@ export class LayerDataSource extends RefCounted {
       disposableOnce(layer.markLoading()),
     );
     if (this.refCounted_ !== undefined) {
+      this.disposeRuntimeStateForDataSourceChange("replace");
       this.refCounted_.dispose();
       this.loadState_ = undefined;
     }

src/layer/segmentation/index.ts

@@ -42,7 +42,10 @@ import {
   registerVolumeLayerType,
   UserLayer,
 } from "#src/layer/index.js";
-import type { LoadedDataSubsource } from "#src/layer/layer_data_source.js";
+import type {
+  LayerDataSourceChangeRuntimeDisposalContext,
+  LoadedDataSubsource,
+} from "#src/layer/layer_data_source.js";
 import { layerDataSourceSpecificationFromJson } from "#src/layer/layer_data_source.js";
 import * as json_keys from "#src/layer/segmentation/json_keys.js";
 import { registerLayerControls } from "#src/layer/segmentation/layer_controls.js";
@@ -788,6 +791,9 @@ function copyOptionalSpatialSkeletonPosition(
   return new Float32Array(Array.from(value, Number));
 }
 
+const SPATIALLY_INDEXED_SKELETON_RUNTIME_DISPOSAL_KIND =
+  "spatiallyIndexedSkeleton";
+
 const Base = UserLayerWithAnnotationsMixin(UserLayer);
 export class SegmentationUserLayer extends Base {
   sliceViewRenderScaleHistogram = new RenderScaleHistogram();
@@ -1535,6 +1541,31 @@ export class SegmentationUserLayer extends Base {
     this.spatialSkeletonState.markNodeDataChanged(options);
   }
 
+  disposeLayerRuntimeStateForDataSourceChange(
+    context: LayerDataSourceChangeRuntimeDisposalContext,
+  ) {
+    if (
+      context.request.kind !== SPATIALLY_INDEXED_SKELETON_RUNTIME_DISPOSAL_KIND
+    ) {
+      return super.disposeLayerRuntimeStateForDataSourceChange(context);
+    }
+    let changed = false;
+    const spatialSkeletonLayers = new Set<SpatiallyIndexedSkeletonLayer>();
+    for (const renderLayer of context.loadedSubsource.renderLayers) {
+      if (
+        renderLayer instanceof PerspectiveViewSpatiallyIndexedSkeletonLayer ||
+        renderLayer instanceof SliceViewPanelSpatiallyIndexedSkeletonLayer
+      ) {
+        spatialSkeletonLayers.add(renderLayer.base);
+      }
+    }
+    for (const spatialSkeletonLayer of spatialSkeletonLayers) {
+      changed = spatialSkeletonLayer.disposeRuntimeState() || changed;
+    }
+    changed = this.spatialSkeletonState.clearRuntimeState() || changed;
+    return changed;
+  }
+
   activateDataSubsources(subsources: Iterable<LoadedDataSubsource>) {
     const updatedSegmentPropertyMaps: SegmentPropertyMap[] = [];
     const isGroupRoot =

src/skeleton/command_history.ts

@@ -87,6 +87,10 @@ export class SpatialSkeletonCommandMappings {
   private nodeIdMappings = new Map<number, number>();
   private segmentIdMappings = new Map<number, number>();
 
+  get empty() {
+    return this.nodeIdMappings.size === 0 && this.segmentIdMappings.size === 0;
+  }
+
   clear() {
     this.nodeIdMappings.clear();
     this.segmentIdMappings.clear();
@@ -251,10 +255,15 @@ export class SpatialSkeletonCommandHistory extends RefCounted {
   }
 
   clear() {
+    const changed =
+      this.undoEntries.length !== 0 ||
+      this.redoEntries.length !== 0 ||
+      !this.mappings.empty;
     this.undoEntries = [];
     this.redoEntries = [];
     this.mappings.clear();
     this.updateState();
+    return changed;
   }
 
   setSource(source: unknown) {

src/skeleton/frontend.ts

@@ -1784,6 +1784,46 @@ export class SpatiallyIndexedSkeletonSource extends SliceViewChunkSource<
   }
 }
 
+export interface SpatiallyIndexedSkeletonSourceRuntimeDisposalOptions {
+  invalidateCache?: boolean;
+}
+
+export function disposeSpatiallyIndexedSkeletonSourceRuntimeState(
+  sources: Iterable<SpatiallyIndexedSkeletonSource>,
+  options: SpatiallyIndexedSkeletonSourceRuntimeDisposalOptions = {},
+) {
+  const uniqueSources = new Set(sources);
+  const invalidateCache = options.invalidateCache ?? true;
+  const chunkQueueManagersWithDeletedChunks = new Set<
+    ChunkManager["chunkQueueManager"]
+  >();
+  let changed = false;
+  for (const source of uniqueSources) {
+    if (source.chunks.size !== 0) {
+      for (const chunkKey of [...source.chunks.keys()]) {
+        source.deleteChunk(chunkKey);
+      }
+      chunkQueueManagersWithDeletedChunks.add(
+        source.chunkManager.chunkQueueManager,
+      );
+      changed = true;
+    }
+    if (
+      invalidateCache &&
+      source.wasDisposed !== true &&
+      source.rpc !== null &&
+      source.rpcId !== null
+    ) {
+      source.invalidateCache();
+      changed = true;
+    }
+  }
+  for (const chunkQueueManager of chunkQueueManagersWithDeletedChunks) {
+    chunkQueueManager.visibleChunksChanged.dispatch();
+  }
+  return changed;
+}
+
 // Options are provided by the SliceView framework for scale selection,
 // but spatial skeleton sources expose all grid levels unconditionally.
 // TODO (SKM): validate if this is an ok deviation from the SliceView
@@ -2100,9 +2140,61 @@ export class SpatiallyIndexedSkeletonLayer
   private cachedSelectedNodeOutlineColorGeneration = -1;
 
   private disposeOverlayChunk() {
+    const changed =
+      this.overlayChunk !== undefined || this.overlayGeometryKey !== undefined;
     this.overlayChunk?.dispose(this.gl);
     this.overlayChunk = undefined;
     this.overlayGeometryKey = undefined;
+    return changed;
+  }
+
+  getUniqueChunkSources() {
+    const sources = new Set<SpatiallyIndexedSkeletonSource>();
+    for (const sourceEntry of [...this.sources, ...this.sources2d]) {
+      sources.add(sourceEntry.chunkSource);
+    }
+    return sources;
+  }
+
+  private clearOverlayRuntimeState() {
+    let changed = this.disposeOverlayChunk();
+    if (this.pendingOverlaySegmentLoads.size !== 0) {
+      this.pendingOverlaySegmentLoads.clear();
+      changed = true;
+    }
+    if (this.editedSegmentIds.size !== 0) {
+      this.editedSegmentIds.clear();
+      changed = true;
+    }
+    if (this.retainedOverlaySegmentIds.length !== 0) {
+      this.retainedOverlaySegmentIds = [];
+      changed = true;
+    }
+    if (this.browseExcludedSegments.size !== 0) {
+      this.browseExcludedSegments.clear();
+      changed = true;
+    }
+    if (this.browseExcludedSegmentsKey !== undefined) {
+      this.browseExcludedSegmentsKey = undefined;
+      changed = true;
+    }
+    this.overlayRebuildFrame = -1;
+    return changed;
+  }
+
+  disposeRuntimeState(
+    options: SpatiallyIndexedSkeletonSourceRuntimeDisposalOptions = {},
+  ) {
+    const overlayChanged = this.clearOverlayRuntimeState();
+    const sourceChanged = disposeSpatiallyIndexedSkeletonSourceRuntimeState(
+      this.getUniqueChunkSources(),
+      options,
+    );
+    const changed = overlayChanged || sourceChanged;
+    if (changed) {
+      this.redrawNeeded.dispatch();
+    }
+    return changed;
   }
 
   private requestOverlaySegmentLoad(segmentId: number) {
@@ -2333,7 +2425,7 @@ export class SpatiallyIndexedSkeletonLayer
   ) {
     super();
     this.registerDisposer(() => {
-      this.disposeOverlayChunk();
+      this.disposeRuntimeState();
     });
     let sources3d: SpatiallyIndexedSkeletonSourceEntry[];
     let sources2d = options.sources2d ?? [];