diff --git a/src/data_manager/data_manager.ts b/src/data_manager/data_manager.ts index fc05ec85..21998f11 100644 --- a/src/data_manager/data_manager.ts +++ b/src/data_manager/data_manager.ts @@ -1,4 +1,4 @@ -import type { ChunkId, ChunkPriority, ChunkEntry, DataManagerLimits, LocalChunkId, Chunk } from "./types.js"; +import type { ChunkId, ChunkPriority, ChunkEntry, DataManagerLimits, Chunk } from "./types.js"; import { chunkIdToString, ChunkState, @@ -14,20 +14,12 @@ import { } from "./types.js"; import type { NumberType, TypedArray } from "../types.js"; import PriorityQueue from "./priority_queue.js"; -import { VolumeDims } from "../VolumeDims.js"; import type { DeviceInterface } from "./device_interface.js"; +import { ChunkSource, type VolumeMetadata, type ExtVolumeDims } from "./sources/chunk_source.js"; +import SlotMap from "./slot_map.js"; const SUBSCRIBER_ID = Symbol("DataManager.subscriberId"); -// TODO not modifying original `VolumeDims` for compatibility, but at some point this will either need to be -// incorporated into `VolumeDims` or replaced with an entirely new type -export type ExtVolumeDims = VolumeDims & { chunkShape: [number, number, number, number, number] }; - -export interface IChunkSource { - getDims(): ExtVolumeDims[]; - getChunk(id: LocalChunkId, signal?: AbortSignal): Promise>; -} - export interface IDataSubscriber { [SUBSCRIBER_ID]?: number; onChunkLoaded?: (id: ChunkId, chunk: Chunk) => void; @@ -35,12 +27,28 @@ export interface IDataSubscriber { // TODO events for when chunks are evicted? } -// TODO move to types file? along with above interfaces...? +// Subscriber IDs increment globally, like timeout IDs +let subscriberCount = 0; + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +const getIdForSubscriber = (subscriber: IDataSubscriber): number => { + if (subscriber[SUBSCRIBER_ID] === undefined) { + subscriber[SUBSCRIBER_ID] = subscriberCount; + subscriberCount += 1; + } + return subscriber[SUBSCRIBER_ID]; +}; + type SourceEntry = { - source: IChunkSource; + source: ChunkSource; subscribers: IDataSubscriber[]; }; +type RequestEntry = { + chunkKeys: string[]; + controller: AbortController; +}; + const swapRemove = (arr: T[], index: number) => { if (index < 0) { return; @@ -88,15 +96,13 @@ export default class DataManager { * * This could be stored in `chunks`, but requests don't necessarily correspond one-to-one with chunks. */ - private requests = new Map(); - /** The amount of chunk data currently cached in memory, in bytes */ + private requests = new SlotMap(); + /** The amount of chunk data currently cached in memory, in bytes. */ private memorySize = 0; - /** The amount of GPU texture memory managed by this class, in bytes */ + /** The amount of GPU texture memory managed by this class, in bytes. */ private deviceSize = 0; /** A counter for assigning chunks a priority at the `RECENT` level. */ private recentCounter = 0; - /** A counter for assigning subscriber IDs. */ - private subscriberCount = 0; public limits: DataManagerLimits; @@ -110,13 +116,23 @@ export default class DataManager { }); } - /** Gets the id of data subscriber `subscriber`, assigning it one if it doesn't have one. */ - private getIdForSubscriber(subscriber: IDataSubscriber): number { - if (subscriber[SUBSCRIBER_ID] === undefined) { - subscriber[SUBSCRIBER_ID] = this.subscriberCount; - this.subscriberCount += 1; + // MARK: Helpers + + private tryAbortRequest(id: number): boolean { + const request = this.requests.get(id); + if (request === undefined) { + return false; + } + + const canAbort = request.chunkKeys.every((key) => { + const chunk = this.chunks.get(key); + return chunk === undefined || chunk.memoryPriority.level === ChunkPriorityLevel.RECENT; + }); + + if (canAbort) { + request.controller.abort(); } - return subscriber[SUBSCRIBER_ID]; + return canAbort; } /** @@ -124,7 +140,7 @@ export default class DataManager { * * Assumes that the chunk is not in any queues that don't match its state. */ - private updateChunkInQueue(key: string, entry: ChunkEntry) { + private updateChunkInQueue(key: string, entry: ChunkEntry, noAbort = false) { switch (entry.data.state) { case ChunkState.QUEUED: this.queues.load.insert(key, entry.memoryPriority); @@ -137,8 +153,8 @@ export default class DataManager { this.queues.deviceEvict.insert(key, entry.devicePriority); break; case ChunkState.LOADING: - if (entry.memoryPriority.level === ChunkPriorityLevel.RECENT) { - // TODO cancel request here + if (entry.memoryPriority.level === ChunkPriorityLevel.RECENT && !noAbort) { + this.tryAbortRequest(entry.data.requestId); } break; case ChunkState.WORKER: @@ -173,6 +189,24 @@ export default class DataManager { } } + /** + * Adds an entry for a chunk that has no requests. + * + * This happens when fetching a requested chunk also requires getting a chunk that no one asked for. + */ + private insertChunkUnprioritized(key: string, data: ChunkEntry["data"]): ChunkEntry { + const entry = { + data, + subscriberPriorities: [], + memoryPriority: { level: ChunkPriorityLevel.RECENT, score: this.recentCounter }, + devicePriority: getMinChunkPriority(), + }; + this.chunks.set(key, entry); + this.recentCounter += 1; + this.updateChunkInQueue(key, entry, true); + return entry; + } + private getChunkDims(chunkId: ChunkId): { x: number; y: number; z: number; dataType: NumberType } | undefined { const sourceEntry = this.sources[chunkId.source]; if (sourceEntry === undefined) { @@ -197,6 +231,7 @@ export default class DataManager { } private estimateChunkSize(chunkId: ChunkId): number { + // TODO account for edge chunks const dims = this.getChunkDims(chunkId); if (dims === undefined) { return 0; @@ -205,6 +240,8 @@ export default class DataManager { return x * y * z * dataTypeToByteLength[dataType]; } + // MARK: Update cycle + /** * Resolves which chunks should be uploaded to the GPU, and which should be evicted from it. * @@ -357,7 +394,7 @@ export default class DataManager { break; case ChunkState.LOADING: console.error(`chunk ${evictKey} queued for eviction while in the "loading" state`); - // TODO cancel request + this.tryAbortRequest(evictEntry.data.requestId); break; default: // TODO `WORKER` state may be weird to manage here @@ -381,47 +418,63 @@ export default class DataManager { return; } let [requestPriority, requestKey] = nextRequest; - let requestId = stringToChunkId(requestKey); + let chunkId = stringToChunkId(requestKey); const nextEvictPriority = this.queues.evict.peek()?.[0] ?? getMinChunkPriority(); while ( // keep submitting requests while concurrency is available and... this.requests.size < requestLimitForPriority(this.limits, requestPriority) && // ...either space will be available or this chunk has higher priority than an already cached chunk - (this.memorySize + this.estimateChunkSize(requestId) < this.limits.size || + (this.memorySize + this.estimateChunkSize(chunkId) < this.limits.size || comparePriority(requestPriority, nextEvictPriority) < 0) ) { - this.queues.load.pop(); - + // Ignore requests at the `RECENT` level. That's supposed to be for chunks that are already loaded! if (requestPriority.level === ChunkPriorityLevel.RECENT) { - // Ignore requests at the `RECENT` level. That's supposed to be for chunks that are already loaded! + this.queues.load.pop(); this.chunks.delete(requestKey); continue; } - const chunkEntry = this.chunks.get(requestKey); - if (chunkEntry === undefined) { - continue; - } - chunkEntry.data = { state: ChunkState.LOADING }; - - const sourceEntry = this.sources[requestId.source]; + // Get the source that this chunk comes from + const sourceId = chunkId.source; + const sourceEntry = this.sources[sourceId]; if (sourceEntry === undefined) { - console.error(`chunk ${requestKey} queued for load with invalid source id ${requestId.source}`); + console.error(`chunk ${requestKey} queued for load with invalid source id ${sourceId}`); + this.queues.load.pop(); this.chunks.delete(requestKey); continue; } - const abortController = new AbortController(); - this.requests.set(requestKey, abortController); + const { source } = sourceEntry; + + // Determine which chunks will be fetched with this request, and set up bookkeeping for the request + // Accounts for sources that store multiple chunks at the same storage id + const storageId = source.chunkIdToStorageId(chunkId); + const chunkIdsAtKey = source.storageIdToChunkIds(storageId); + const chunkKeys = chunkIdsAtKey.map((id) => chunkIdToString({ ...id, source: sourceId })); + const controller = new AbortController(); + const requestId = this.requests.insert({ chunkKeys, controller }); + + for (const key of chunkKeys) { + this.queues.load.remove(key); + const entry = this.chunks.get(key); + const data = { state: ChunkState.LOADING as const, requestId }; + if (entry === undefined) { + this.insertChunkUnprioritized(key, data); + } else if (entry.data.state === ChunkState.QUEUED) { + entry.data = data; + } + } - // Make a stable by-reference copy of `requestId` for the closures below - const requestedId = requestId; + // Request the storage key sourceEntry.source - .getChunk(requestedId, abortController.signal) - .then((data) => this.onChunkLoad(requestedId, data)) + .getKey(storageId, controller.signal) + .then((chunks) => { + this.requests.remove(requestId); + chunks.map((chunk) => this.onChunkLoad(sourceId, chunk)); + }) .catch(() => { - this.chunks.delete(requestKey); - this.requests.delete(requestKey); + this.requests.remove(requestId); + chunkKeys.map((key) => this.chunks.delete(key)); }); const nextRequest = this.queues.load.peek(); @@ -429,7 +482,7 @@ export default class DataManager { return; } [requestPriority, requestKey] = nextRequest; - requestId = stringToChunkId(requestKey); + chunkId = stringToChunkId(requestKey); } } @@ -448,32 +501,25 @@ export default class DataManager { this.submitRequests(); } - private onChunkLoad(id: ChunkId, chunk: Chunk) { - const key = chunkIdToString(id); + private onChunkLoad(source: number, chunk: Chunk) { + const globalId = { ...chunk.id, source }; + const key = chunkIdToString(globalId); const { data: memory, dtype } = chunk; - this.requests.delete(key); if (memory.byteLength > this.limits.size) { console.error(`received chunk ${key} which is larger than the cache limit`); return; } - const sourceEntry = this.sources[id.source]; + const sourceEntry = this.sources[source]; if (sourceEntry === undefined) { - console.error(`received chunk ${key} with invalid source id ${id.source}`); + console.error(`received chunk ${key} with invalid source id ${source}`); return; } const data = { state: ChunkState.MEMORY as const, memory, dtype }; let chunkEntry = this.chunks.get(key); if (chunkEntry === undefined) { - console.error(`received chunk ${key} without data manager entry`); - chunkEntry = { - data, - subscriberPriorities: [], - memoryPriority: { level: ChunkPriorityLevel.RECENT, score: this.recentCounter }, - devicePriority: { level: ChunkPriorityLevel.RECENT, score: this.recentCounter }, - }; - this.chunks.set(key, chunkEntry); - this.recentCounter += 1; + console.warn(`received chunk ${key} without data manager entry`); + chunkEntry = this.insertChunkUnprioritized(key, data); } else { chunkEntry.data = data; } @@ -482,11 +528,13 @@ export default class DataManager { this.queues.evict.insert(key, chunkEntry.memoryPriority); this.memorySize += memory.byteLength; - sourceEntry.subscribers.forEach((s) => s.onChunkLoaded?.(id, chunk)); + sourceEntry.subscribers.forEach((s) => s.onChunkLoaded?.(globalId, chunk)); this.update(); } + // MARK: Public interface + /** * Declares that subscriber `subscriber` requires chunk `chunkId` at priority level `priority`. * @@ -494,7 +542,7 @@ export default class DataManager { * submitted until the next call to `update`. */ addChunkRequest(subscriber: IDataSubscriber, chunkId: ChunkId, priority: ChunkPriority, device: boolean) { - const subscriberId = this.getIdForSubscriber(subscriber); + const subscriberId = getIdForSubscriber(subscriber); const chunkIdString = chunkIdToString(chunkId); const chunkEntry = this.chunks.get(chunkIdString); @@ -527,7 +575,7 @@ export default class DataManager { * The `DataManager` will not respond to this change until the next call to `update`. */ removeChunkRequest(subscriber: IDataSubscriber, chunkId: ChunkId) { - const subscriberId = this.getIdForSubscriber(subscriber); + const subscriberId = getIdForSubscriber(subscriber); const chunkKey = chunkIdToString(chunkId); const chunkEntry = this.chunks.get(chunkKey); if (chunkEntry === undefined) { @@ -573,7 +621,7 @@ export default class DataManager { * * Returns the ID that will be used to identify this source in chunk requests and event subscriptions. */ - addSource(source: IChunkSource): number { + addSource(source: ChunkSource): number { const id = this.sources.length; this.sources.push({ source, subscribers: [] }); return id; @@ -584,6 +632,11 @@ export default class DataManager { return this.sources[sourceId]?.source.getDims(); } + /** Returns the metadata of the given source. */ + getSourceMeta(sourceId: number): VolumeMetadata | undefined { + return this.sources[sourceId]?.source.getMeta(); + } + /** Subscribes `subscriber` to data events from the source with id `sourceId` */ subscribeToSource(subscriber: IDataSubscriber, sourceId: number): boolean { const sourceEntry = this.sources[sourceId]; diff --git a/src/data_manager/slot_map.ts b/src/data_manager/slot_map.ts new file mode 100644 index 00000000..cf2703ed --- /dev/null +++ b/src/data_manager/slot_map.ts @@ -0,0 +1,61 @@ +const NEXT_FREE = Symbol("SlotMap.nextFree"); + +type Slot = T | { [NEXT_FREE]: number }; + +const isFreeSlot = (slot: Slot): slot is { [NEXT_FREE]: number } => { + return slot !== undefined && slot !== null && Object.hasOwn(slot, NEXT_FREE); +}; + +export default class SlotMap { + private slots: Slot[] = []; + private head = 0; + private count = 0; + + public get size(): number { + return this.count; + } + + public insert(value: T): number { + this.count += 1; + + if (this.head < this.slots.length) { + const freeSlot = this.slots[this.head]; + if (!isFreeSlot(freeSlot)) { + // should never happen + throw new Error(`SlotMap: free list head (${this.head}) points to occupied slot`); + } + + const index = this.head; + this.head = freeSlot[NEXT_FREE]; + this.slots[index] = value; + return index; + } + + const index = this.slots.length; + this.slots.push(value); + this.head = this.slots.length; + return index; + } + + public get(index: number): T | undefined { + if (index >= this.slots.length) { + return undefined; + } + const slot = this.slots[index]; + if (isFreeSlot(slot)) { + return undefined; + } + return slot; + } + + public remove(index: number): T | undefined { + const value = this.get(index); + if (value === undefined) { + return undefined; + } + this.slots[index] = { [NEXT_FREE]: this.head }; + this.head = index; + this.count -= 1; + return value; + } +} diff --git a/src/data_manager/sources/chunk_source.ts b/src/data_manager/sources/chunk_source.ts new file mode 100644 index 00000000..2c0a51f7 --- /dev/null +++ b/src/data_manager/sources/chunk_source.ts @@ -0,0 +1,45 @@ +import type { NumberType } from "../../types.js"; +import type { VolumeDims } from "../../VolumeDims.js"; +import { type Chunk, type LocalChunkId } from "../types.js"; + +// TODO not modifying original `VolumeDims` for compatibility, but at some point this will either need to be +// incorporated into `VolumeDims` or replaced with an entirely new type +export type ExtVolumeDims = VolumeDims & { chunkShape: [number, number, number, number, number] }; + +// This is more or less the subset of `ImageInfo` that is *not* derivable from `[Ext]VolumeDims` and is therefore *not* +// essential to handling volume data and rendering it at the correct scale. +export type VolumeMetadata = { + name?: string; + channelNames: string[]; + channelColors?: ([number, number, number] | undefined)[]; + transform: { + translation: [number, number, number]; + rotation: [number, number, number]; + scale: [number, number, number]; + }; + userData?: Record; +}; + +export abstract class ChunkSource { + /** + * Maps a chunk `id` to the *storage key* that contains it. + * + * In most cases, a chunk id and its storage key are the same. However, `DataManager` requires that chunks extend + * over exactly one channel and exactly one time point, and some data formats may allow multiple channels and/or time + * points to be stored at the same storage key. + */ + chunkIdToStorageId(id: LocalChunkId): LocalChunkId { + return id; + } + + /** Maps a storage key to a list of the chunk ids it contains. */ + storageIdToChunkIds(id: LocalChunkId): LocalChunkId[] { + return [id]; + } + + abstract getMeta(): VolumeMetadata; + + abstract getDims(): ExtVolumeDims[]; + + abstract getKey(key: LocalChunkId, signal?: AbortSignal): Promise[]>; +} diff --git a/src/data_manager/sources/ome_zarr_source.ts b/src/data_manager/sources/ome_zarr_source.ts index ceb2ccd0..66c43c82 100644 --- a/src/data_manager/sources/ome_zarr_source.ts +++ b/src/data_manager/sources/ome_zarr_source.ts @@ -1,13 +1,19 @@ import * as zarr from "zarrita"; -import type { ExtVolumeDims, IChunkSource } from "../data_manager.js"; +import { type ExtVolumeDims, ChunkSource, type VolumeMetadata } from "./chunk_source.js"; import { assertMetadataHasMultiscales, toOMEZarrMetaV4, validateOMEZarrMetadata, } from "../../loaders/zarr_utils/validation.js"; import type { OMEMultiscale, OmeroTransitionalMetadata } from "../../loaders/zarr_utils/types.js"; -import { getScale, orderByDimension, orderByTCZYX, remapAxesToTCZYX } from "../../loaders/zarr_utils/utils.js"; +import { + getScale, + getSourceChannelMeta, + orderByDimension, + orderByTCZYX, + remapAxesToTCZYX, +} from "../../loaders/zarr_utils/utils.js"; import { unitNameToSymbol } from "../../loaders/VolumeLoaderUtils.js"; import type { Chunk, LocalChunkId } from "../types.js"; import type { NumberType, TypedArray } from "../../types.js"; @@ -15,13 +21,20 @@ import type { NumberType, TypedArray } from "../../types.js"; const PLACEHOLDER_NAME = "zarr source"; const PLACEHOLDER_SCENE_INDEX = 0; -export class OMEZarrSource implements IChunkSource { +export class OMEZarrSource extends ChunkSource { + // Magic number that makes `getSourceChannelMeta` (below) work. + // TODO if/when this class becomes the primary adapter for zarrs, move that util over here and adapt it such that + // this property can be removed and the remaining properties can be made `private`. + public readonly channelOffset = 0; + private constructor( - private scaleLevels: zarr.Array[], - private multiscaleMetadata: OMEMultiscale, - private omeroMetadata: OmeroTransitionalMetadata | undefined, - private axesTCZYX: [number, number, number, number, number] - ) {} + public readonly scaleLevels: zarr.Array[], + public readonly multiscaleMetadata: OMEMultiscale, + public readonly omeroMetadata: OmeroTransitionalMetadata | undefined, + public readonly axesTCZYX: [number, number, number, number, number] + ) { + super(); + } static async new(url: string): Promise { const store = new zarr.FetchStore(url); @@ -57,6 +70,47 @@ export class OMEZarrSource implements IChunkSource { return [spaceUnitSymbol, timeUnitSymbol]; } + chunkIdToStorageId({ multiscale, tczyx: [t, c, z, y, x] }: LocalChunkId): LocalChunkId { + const level = this.scaleLevels[multiscale]; + const [shapeT, shapeC] = orderByTCZYX(level.chunks, this.axesTCZYX, 1); + const keyT = Math.floor(t / shapeT); + const keyC = Math.floor(c / shapeC); + return { multiscale, tczyx: [keyT, keyC, z, y, x] }; + } + + storageIdToChunkIds({ multiscale, tczyx: [t, c, z, y, x] }: LocalChunkId): LocalChunkId[] { + const level = this.scaleLevels[multiscale]; + const [shapeT, shapeC] = orderByTCZYX(level.chunks, this.axesTCZYX, 1); + + const minT = t * shapeT; + const maxT = minT + shapeT; + const minC = c * shapeC; + const maxC = minC + shapeC; + + const chunksInKey: LocalChunkId[] = []; + for (let chunkT = minT; chunkT < maxT; chunkT++) { + for (let chunkC = minC; chunkC < maxC; chunkC++) { + chunksInKey.push({ multiscale, tczyx: [chunkT, chunkC, z, y, x] }); + } + } + return chunksInKey; + } + + getMeta(): VolumeMetadata { + const { names: channelNames, colors: channelColors } = getSourceChannelMeta(this); + + return { + name: this.omeroMetadata?.name ?? PLACEHOLDER_NAME, + channelNames, + channelColors, + transform: { + translation: [0, 0, 0], + rotation: [0, 0, 0], + scale: [1, 1, 1], + }, + }; + } + getDims(): ExtVolumeDims[] { const [spaceUnit, timeUnit] = this.getUnitSymbols(); return this.scaleLevels.map((level, i) => { @@ -72,10 +126,10 @@ export class OMEZarrSource implements IChunkSource { }); } - async getChunk(chunkId: LocalChunkId): Promise> { - const multiscale = this.scaleLevels[chunkId.multiscale]; - const coords = orderByDimension(chunkId.tczyx, this.axesTCZYX); - const { data } = await multiscale.getChunk(coords); - return { data: data as TypedArray, dtype: multiscale.dtype }; + async getKey(id: LocalChunkId, signal?: AbortSignal): Promise[]> { + const multiscale = this.scaleLevels[id.multiscale]; + const coords = orderByDimension(id.tczyx, this.axesTCZYX); + const { data } = await multiscale.getChunk(coords, { signal }); + return [{ data: data as TypedArray, dtype: multiscale.dtype, id }]; } } diff --git a/src/data_manager/types.ts b/src/data_manager/types.ts index a20dfb7e..c5375341 100644 --- a/src/data_manager/types.ts +++ b/src/data_manager/types.ts @@ -112,8 +112,9 @@ export const enum ChunkState { WORKER = "worker", } -export type ChunkData = - | { state: ChunkState.QUEUED | ChunkState.WORKER | ChunkState.LOADING } +type ChunkEntryData = + | { state: ChunkState.QUEUED | ChunkState.WORKER } + | { state: ChunkState.LOADING; requestId: number } | { state: ChunkState.MEMORY; memory: TypedArray; dtype: NumberType } | { state: ChunkState.DEVICE; memory: TypedArray; dtype: NumberType; texture: Tex }; @@ -126,7 +127,7 @@ export type SubscriberPriority = { /** `DataManager`-internal type containing full state of a single chunk. */ export type ChunkEntry = { /** The current lifecycle state of a chunk, and the data associated with it. */ - data: ChunkData; + data: ChunkEntryData; /** A record of every request for this chunk. */ subscriberPriorities: SubscriberPriority[]; /** Priority at which the chunk is required in memory. Derived from `subscriberPriorities` on request add/remove. */ @@ -158,7 +159,8 @@ export const stringToChunkId = (id: string): ChunkId => { }; }; -export type Chunk = { +export type Chunk = { + id: LocalChunkId; data: TypedArray; dtype: T; }; diff --git a/src/loaders/zarr_utils/utils.ts b/src/loaders/zarr_utils/utils.ts index 2f8736f7..366c4c81 100644 --- a/src/loaders/zarr_utils/utils.ts +++ b/src/loaders/zarr_utils/utils.ts @@ -28,7 +28,9 @@ export function parseHexColor(color: string | undefined): [number, number, numbe } /** Extracts channel names from a `ZarrSource`. Handles missing `omeroMetadata`. Does *not* resolve name collisions. */ -export function getSourceChannelMeta(src: ZarrSource): { +export function getSourceChannelMeta( + src: Pick +): { names: string[]; colors: ([number, number, number] | undefined)[]; } { diff --git a/src/test/data_manager.test.ts b/src/test/data_manager.test.ts index 47b4fb52..9858fca9 100644 --- a/src/test/data_manager.test.ts +++ b/src/test/data_manager.test.ts @@ -1,12 +1,13 @@ -import type { IChunkSource, IDataSubscriber } from "../data_manager/data_manager.js"; +import type { IDataSubscriber } from "../data_manager/data_manager.js"; import DataManager from "../data_manager/data_manager.js"; import type { DeviceInterface } from "../data_manager/device_interface.js"; +import type { ChunkSource } from "../data_manager/sources/chunk_source.js"; import { type ChunkPriority, ChunkPriorityLevel } from "../data_manager/types.js"; import type { TypedArray, NumberType } from "../types.js"; const mockSource = () => { return { - getDims: vi.fn(() => [ + getDims: vi.fn(() => [ { shape: [7, 7, 7, 7, 7], spacing: [1, 1, 1, 1, 1], @@ -16,13 +17,23 @@ const mockSource = () => { timeUnit: "second", }, ]), - getChunk: vi.fn(async ({ tczyx }) => { - const length = tczyx.slice(2).reduce((len, coord) => (coord >= 3 ? len : len * 2), 1); + getKey: vi.fn(async (id) => { + const length = id.tczyx.slice(2).reduce((len, coord) => (coord >= 3 ? len : len * 2), 1); const data = new Uint8Array(length).fill(0); - data.set(tczyx); - return { data, dtype: "uint8" }; + data.set(id.tczyx); + return [{ data, dtype: "uint8", id }]; }), - } satisfies IChunkSource; + getMeta: vi.fn(() => ({ + channelNames: ["one", "two", "three"], + transform: { + translation: [0, 0, 0], + rotation: [0, 0, 0], + scale: [1, 1, 1], + }, + })), + chunkIdToStorageId: vi.fn((id) => id), + storageIdToChunkIds: vi.fn((id) => [id]), + } satisfies ChunkSource; }; type MockTex = { @@ -69,7 +80,7 @@ describe("request priority", () => { queue(2, 2); manager.update(); - const getCalls = sourceMock.getChunk.mock.calls; + const getCalls = sourceMock.getKey.mock.calls; expect(getCalls[0][0].tczyx[4]).toEqual(2); expect(getCalls[1][0].tczyx[4]).toEqual(0); expect(getCalls[2][0].tczyx[4]).toEqual(1); @@ -87,7 +98,7 @@ describe("request priority", () => { queue(1, { level: ChunkPriorityLevel.VISIBLE, score: 0 }); manager.update(); - const getCalls = sourceMock.getChunk.mock.calls; + const getCalls = sourceMock.getKey.mock.calls; expect(getCalls[0][0].tczyx[4]).toEqual(1); expect(getCalls[1][0].tczyx[4]).toEqual(0); });