Invalidate shard index cache when chunk source cache is invalidated

src/chunk_manager/backend.ts

@@ -295,6 +295,13 @@ export class ChunkSourceBase extends SharedObject {
    */
   sourceQueueLevel = 0;
 
+  /**
+   * Optional hook called when the source's cache is invalidated.
+   * Subclasses can override this to invalidate related caches
+   * (e.g. shard index caches for sharded zarr3 sources).
+   */
+  onInvalidateCache?(): void;
+
   constructor(public chunkManager: Borrowed<ChunkManager>) {
     super();
     chunkManager.queueManager.sources.add(this);
@@ -1111,6 +1118,7 @@ export class ChunkQueueManager extends SharedObjectCounterpart {
   }
 
   invalidateSourceCache(source: ChunkSource) {
+    source.onInvalidateCache?.();
     for (const chunk of source.chunks.values()) {
       switch (chunk.state) {
         case ChunkState.DOWNLOADING:

src/datasource/zarr/backend.ts

@@ -46,6 +46,11 @@ export class ZarrVolumeChunkSource extends WithParameters(
     this.sharedKvStoreContext.kvStoreContext.getKvStore(this.parameters.url),
   );
 
+  onInvalidateCache() {
+    const { kvStore } = this.chunkKvStore;
+    kvStore.invalidateIndexCache?.();
+  }
+
   async download(chunk: VolumeChunk, signal: AbortSignal) {
     chunk.chunkDataSize = this.spec.chunkDataSize;
     const { parameters } = this;

src/datasource/zarr/codec/decode.ts

@@ -140,19 +140,23 @@ export async function decodeArray(
   return decoded;
 }
 
+export interface ShardedKvStoreWithInvalidation {
+  invalidateIndexCache?: () => void;
+}
+
 export function applySharding(
   chunkManager: ChunkManager,
   codecs: CodecChainSpec,
   baseKvStore: KvStoreWithPath,
 ): {
-  kvStore: ReadableKvStore<unknown>;
+  kvStore: ReadableKvStore<unknown> & Partial<ShardedKvStoreWithInvalidation>;
   getChunkKey: (
     chunkGridPosition: ArrayLike<number>,
     baseKey: string,
   ) => unknown;
   decodeCodecs: CodecChainSpec;
 } {
-  let kvStore: ReadableKvStore<unknown> = baseKvStore.store;
+  let kvStore: ReadableKvStore<unknown> & Partial<ShardedKvStoreWithInvalidation> = baseKvStore.store;
   let curCodecs = codecs;
   while (true) {
     const { shardingInfo } = curCodecs;

src/datasource/zarr/codec/sharding_indexed/decode.ts

@@ -15,6 +15,7 @@
  */
 
 import type { ChunkManager } from "#src/chunk_manager/backend.js";
+import { ChunkState } from "#src/chunk_manager/base.js";
 import { SimpleAsyncCache } from "#src/chunk_manager/generic_file_source.js";
 import {
   decodeArray,
@@ -180,6 +181,16 @@ class ShardedKvStore<BaseKey>
     return `subchunk ${JSON.stringify(key.subChunk)} within shard ${this.base.getUrl(key.base)}`;
   }
 
+  invalidateIndexCache() {
+    const { indexCache } = this;
+    for (const chunk of indexCache.chunks.values()) {
+      if (chunk.state === ChunkState.SYSTEM_MEMORY_WORKER) {
+        chunk.freeSystemMemory();
+      }
+      indexCache.chunkManager.queueManager.updateChunkState(chunk, ChunkState.QUEUED);
+    }
+  }
+
   get supportsOffsetReads() {
     return true;
   }