diff --git a/src/data_manager/data_manager.ts b/src/data_manager/data_manager.ts new file mode 100644 index 00000000..5ce886aa --- /dev/null +++ b/src/data_manager/data_manager.ts @@ -0,0 +1,618 @@ +import type { + ChunkId, + ChunkPriority, + ChunkEntry, + DataManagerLimits, + LocalChunkId, + Chunk, + DeviceInterface, +} from "./types.js"; +import { + chunkIdToString, + ChunkState, + comparePriority, + reverseComparePriority, + getMinChunkPriority, + validateDataManagerLimits, + DEFAULT_DATA_MANAGER_LIMITS, + stringToChunkId, + requestLimitForPriority, + ChunkPriorityLevel, + deviceSizeLimitForPriority, +} from "./types.js"; +import type { NumberType, TypedArray } from "../types.js"; +import PriorityQueue from "./priority_queue.js"; +import { VolumeDims } from "../VolumeDims.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; + onChunkOnGpu?: (id: ChunkId, texture: Tex) => void; + // TODO events for when chunks are evicted? +} + +// TODO move to types file? along with above interfaces...? +type SourceEntry = { + source: IChunkSource; + subscribers: IDataSubscriber[]; +}; + +const swapRemove = (arr: T[], index: number) => { + if (index < 0) { + return; + } + + const { length } = arr; + const replace = arr.pop(); + if (replace !== undefined && index < length) { + arr[index] = replace; + } +}; + +const dataTypeToByteLength: { [T in NumberType]: number } = { + int8: 1, + int16: 2, + int32: 4, + uint8: 1, + uint16: 2, + uint32: 4, + float32: 4, + float64: 8, +}; + +type ChunkQueue = PriorityQueue; + +class ChunkQueues { + /** Chunks waiting to be loaded */ + load: ChunkQueue = new PriorityQueue(comparePriority); + /** Chunks in memory that may be evicted */ + evict: ChunkQueue = new PriorityQueue(reverseComparePriority); + /** Chunks in memory waiting to be uploaded to the GPU */ + deviceLoad: ChunkQueue = new PriorityQueue(comparePriority); + /** Chunks on the GPU that may be evicted */ + deviceEvict: ChunkQueue = new PriorityQueue(reverseComparePriority); +} + +export default class DataManager { + /** Data and current state for every chunk of data tracked and managed by this class. */ + private chunks = new Map>(); + private queues = new ChunkQueues(); + private sources: SourceEntry[] = []; + /** + * Information about each in-flight request. + * + * 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 memorySize = 0; + /** 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; + + constructor( + private deviceInterface: DeviceInterface, + limits?: Partial + ) { + this.limits = validateDataManagerLimits({ + ...DEFAULT_DATA_MANAGER_LIMITS, + ...(limits ?? {}), + }); + } + + /** 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; + } + return subscriber[SUBSCRIBER_ID]; + } + + /** + * Inserts a chunk into the appropriate queue, or updates its queue position. + * + * Assumes that the chunk is not in any queues that don't match its state. + */ + private updateChunkInQueue(key: string, entry: ChunkEntry) { + switch (entry.data.state) { + case ChunkState.QUEUED: + this.queues.load.insert(key, entry.memoryPriority); + break; + case ChunkState.MEMORY: + this.queues.deviceLoad.insert(key, entry.devicePriority); + this.queues.evict.insert(key, entry.memoryPriority); + break; + case ChunkState.DEVICE: + this.queues.deviceEvict.insert(key, entry.devicePriority); + break; + case ChunkState.LOADING: + if (entry.memoryPriority.level === ChunkPriorityLevel.RECENT) { + // TODO cancel request here + } + break; + case ChunkState.WORKER: + } + } + + /** Resolves a chunk's overall priority based on all requests for it, then updates its queue position. */ + private updateChunkPriority(key: string, entry: ChunkEntry) { + let nextMemory = getMinChunkPriority(); + let nextDevice = getMinChunkPriority(); + entry.subscriberPriorities.forEach(({ priority, device }) => { + if (comparePriority(priority, nextMemory) < 0) { + nextMemory = priority; + } + if (device && comparePriority(priority, nextDevice) < 0) { + nextDevice = priority; + } + }); + + if (nextMemory.level === ChunkPriorityLevel.RECENT) { + nextMemory.score = this.recentCounter; + this.recentCounter += 1; + } + + if ( + comparePriority(nextMemory, entry.memoryPriority) !== 0 || + comparePriority(nextDevice, entry.devicePriority) !== 0 + ) { + entry.memoryPriority = { ...nextMemory }; + entry.devicePriority = { ...nextDevice }; + this.updateChunkInQueue(key, entry); + } + } + + private getChunkDims(chunkId: ChunkId): { x: number; y: number; z: number; dataType: NumberType } | undefined { + const sourceEntry = this.sources[chunkId.source]; + if (sourceEntry === undefined) { + return undefined; + } + + const multiscale = sourceEntry.source.getDims()[chunkId.multiscale]; + if (multiscale === undefined) { + return undefined; + } + + const { shape, chunkShape, dataType } = multiscale; + const [, , iz, iy, ix] = chunkId.tczyx; + const [, , cz, cy, cx] = chunkShape; + const [, , sz, sy, sx] = shape; + // Chunks that abut the max edge of the volume in any dimension may be smaller in that dimension. + // e.g. for a volume with chunk size 10 & volume size 35 in x, chunks at x = 3 have size 5 in x. + const z = Math.min(cz, sz - cz * iz); + const y = Math.min(cy, sy - cy * iy); + const x = Math.min(cx, sx - cx * ix); + return { x, y, z, dataType }; + } + + private estimateChunkSize(chunkId: ChunkId): number { + const dims = this.getChunkDims(chunkId); + if (dims === undefined) { + return 0; + } + const { x, y, z, dataType } = dims; + return x * y * z * dataTypeToByteLength[dataType]; + } + + /** + * Resolves which chunks should be uploaded to the GPU, and which should be evicted from it. + * + * In other words, this function drives the `deviceLoad` and `deviceEvict` queues. + */ + private updateDeviceData(deviceHandle: Dev) { + // STEP 1: pull eligible chunks out of the `deviceLoad` queue + const loads: [string, ChunkEntry][] = []; + // This one's easier if we just loop unconditionally and `break` when done. + while (true) { + const nextLoad = this.queues.deviceLoad.peek(); + if (nextLoad === undefined) { + // all out of queued device loads + break; + } + const [loadPriority, loadKey] = nextLoad; + + const loadEntry = this.chunks.get(loadKey); + // type safety requires these next two `if`s, but neither should be true unless this class has a bug + if (loadEntry === undefined) { + console.error(`chunk ${loadKey} queued for device upload without data`); + this.queues.deviceLoad.pop(); + this.queues.evict.remove(loadKey); + continue; + } + if (loadEntry.data.state !== ChunkState.MEMORY) { + console.error( + `chunk ${loadKey} queued for device upload in invalid state (expected "memory", found ${loadEntry.data.state})` + ); + this.queues.deviceLoad.pop(); + this.queues.evict.remove(loadKey); + this.updateChunkInQueue(loadKey, loadEntry); + continue; + } + + const chunkSize = loadEntry.data.memory.byteLength; + const nextEvictPriority = this.queues.deviceEvict.peek()?.[0] ?? getMinChunkPriority(); + if ( + // queue this chunk if there's space for it... + !(this.deviceSize + chunkSize <= deviceSizeLimitForPriority(this.limits, loadPriority)) && + // ...or if its priority is greater than the lowest-priority chunk that's already on the GPU. + // (this implies that at least one chunk will be evicted from the GPU in the next step) + !(loadPriority.level !== ChunkPriorityLevel.RECENT && comparePriority(loadPriority, nextEvictPriority) < 0) + ) { + // otherwise, we're done + break; + } + + this.queues.deviceLoad.pop(); + this.queues.evict.remove(loadKey); + this.queues.deviceEvict.insert(loadKey, loadPriority); + this.deviceSize += chunkSize; + loads.push([loadKey, loadEntry]); + } + + // STEP 2: evict chunks that have been pushed off the GPU + let nextEvict = this.queues.deviceEvict.peek(); + while (nextEvict !== undefined && this.deviceSize >= deviceSizeLimitForPriority(this.limits, nextEvict[0])) { + const [evictPriority, evictKey] = nextEvict; + this.queues.deviceEvict.pop(); + this.queues.deviceLoad.insert(evictKey, evictPriority); + this.queues.evict.insert(evictKey, evictPriority); + + const evictEntry = this.chunks.get(evictKey); + if (evictEntry !== undefined) { + if (evictEntry.data.state === ChunkState.DEVICE) { + // this chunk was previously on the GPU; demote to `MEMORY` and destroy its texture + const { memory, dtype } = evictEntry.data; + this.deviceSize -= memory?.byteLength; + this.deviceInterface.destroyTexture(evictEntry.data.texture); + evictEntry.data = { state: ChunkState.MEMORY, memory, dtype }; + } else if (evictEntry.data.state === ChunkState.MEMORY) { + // this chunk was promoted in the previous step; put it back + this.deviceSize -= evictEntry.data.memory.byteLength; + const index = loads.findIndex(([_, entry]) => entry === evictEntry); + swapRemove(loads, index); + } else { + // again, if either of the following errors are ever printed, there's a bug somewhere in this file + console.error( + `chunk ${evictKey} queued for device evict in invalid state (expected "device" or "memory", found ${evictEntry.data.state})` + ); + } + } else { + console.error(`chunk ${evictKey} queued for device evict without data`); + } + + nextEvict = this.queues.deviceEvict.peek(); + } + + // STEP 3: create textures for newly-promoted chunks + for (const [loadKey, loadEntry] of loads) { + const loadId = stringToChunkId(loadKey); + const dims = this.getChunkDims(loadId); + + if (dims === undefined) { + console.error( + `chunk ${loadKey} queued for device upload with invalid source id or multiscale index (id ${loadId.source}, multiscale index ${loadId.multiscale})` + ); + this.queues.deviceEvict.remove(loadKey); + this.chunks.delete(loadKey); + continue; + } + + const { x, y, z, dataType } = dims; + const { memory } = loadEntry.data as { memory: TypedArray }; + const texture = this.deviceInterface.createTexture(memory, dataType, [x, y, z], deviceHandle); + + this.sources[loadId.source].subscribers.forEach((s) => s.onChunkOnGpu?.(loadId, texture)); + + loadEntry.data = { state: ChunkState.DEVICE, memory, texture, dtype: dataType }; + } + + this.deviceInterface.finishUpdate(); + } + + /** + * Evicts data cached in memory to stay under the size limit. + * + * In other words, this function drives the `evict` queue. + */ + private evictCachedData() { + while (this.memorySize > this.limits.size) { + const nextEvict = this.queues.evict.pop(); + if (nextEvict === undefined) { + // should never get here (checked above that the chunk isn't larger than the entire cache) + break; + } + const evictKey = nextEvict[1]; + + this.queues.deviceLoad.remove(evictKey); + const evictEntry = this.chunks.get(evictKey); + // yup... below `if` makes the type system happy, but if they ever evaluate to `true` that's a bug in this class. + if (evictEntry === undefined) { + console.error(`chunk ${evictKey} queued for eviction without data`); + continue; + } + + switch (evictEntry.data.state) { + case ChunkState.MEMORY: + // If we find any state other than this one, it's a bug! + this.memorySize -= evictEntry.data.memory.byteLength; + break; + case ChunkState.DEVICE: + console.error(`chunk ${evictKey} queued for eviction while in the "device" state`); + this.deviceSize -= evictEntry.data.memory.byteLength; + this.memorySize -= evictEntry.data.memory.byteLength; + this.deviceInterface.destroyTexture(evictEntry.data.texture); + this.queues.deviceEvict.remove(evictKey); + // TODO if subscribers get "chunk removed from GPU" events, one should go here + break; + case ChunkState.LOADING: + console.error(`chunk ${evictKey} queued for eviction while in the "loading" state`); + // TODO cancel request + break; + default: + // TODO `WORKER` state may be weird to manage here + console.error( + `chunk ${evictKey} queued for eviction in an invalid state (expected "memory", found ${evictEntry.data.state})` + ); + } + + this.chunks.delete(evictKey); + } + } + + /** + * Submits requests for chunks that are queued to load. + * + * In other words, this function drives the `load` queue. + */ + private submitRequests() { + const nextRequest = this.queues.load.peek(); + if (nextRequest === undefined) { + return; + } + let [requestPriority, requestKey] = nextRequest; + let requestId = 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 || + comparePriority(requestPriority, nextEvictPriority) < 0) + ) { + this.queues.load.pop(); + + if (requestPriority.level === ChunkPriorityLevel.RECENT) { + // Ignore requests at the `RECENT` level. That's supposed to be for chunks that are already loaded! + 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]; + if (sourceEntry === undefined) { + console.error(`chunk ${requestKey} queued for load with invalid source id ${requestId.source}`); + this.chunks.delete(requestKey); + continue; + } + const abortController = new AbortController(); + this.requests.set(requestKey, abortController); + + // Make a stable by-reference copy of `requestId` for the closures below + const requestedId = requestId; + sourceEntry.source + .getChunk(requestedId, abortController.signal) + .then((data) => this.onChunkLoad(requestedId, data)) + .catch(() => { + this.chunks.delete(requestKey); + this.requests.delete(requestKey); + }); + + const nextRequest = this.queues.load.peek(); + if (nextRequest === undefined) { + return; + } + [requestPriority, requestKey] = nextRequest; + requestId = stringToChunkId(requestKey); + } + } + + /** + * Drives progress on all queued requests, including creating textures, evicting data, and submitting requests. + * + * Request queueing and submission happen in separate steps to allow all requests to be queued and sorted into + * priority order before submitting. If requests were submitted immediately upon being added to the queue, the first + * requests in a batch would submit in the order they were queued, not priority order. + */ + update(deviceHandle?: Dev) { + if (this.deviceInterface.isDeviceHandle(deviceHandle)) { + this.updateDeviceData(deviceHandle); + } + this.evictCachedData(); + this.submitRequests(); + } + + private onChunkLoad(id: ChunkId, chunk: Chunk) { + const key = chunkIdToString(id); + 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]; + if (sourceEntry === undefined) { + console.error(`received chunk ${key} with invalid source id ${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; + } else { + chunkEntry.data = data; + } + + this.queues.deviceLoad.insert(key, chunkEntry.memoryPriority); + this.queues.evict.insert(key, chunkEntry.memoryPriority); + + this.memorySize += memory.byteLength; + sourceEntry.subscribers.forEach((s) => s.onChunkLoaded?.(id, chunk)); + + this.update(); + } + + /** + * Declares that subscriber `subscriber` requires chunk `chunkId` at priority level `priority`. + * + * If the chunk is not already in memory, this will queue the chunk to be loaded. Chunk load requests are not + * submitted until the next call to `update`. + */ + addChunkRequest(subscriber: IDataSubscriber, chunkId: ChunkId, priority: ChunkPriority, device: boolean) { + const subscriberId = this.getIdForSubscriber(subscriber); + const chunkIdString = chunkIdToString(chunkId); + const chunkEntry = this.chunks.get(chunkIdString); + + if (chunkEntry === undefined) { + const newChunkEntry: ChunkEntry = { + data: { state: ChunkState.QUEUED }, + subscriberPriorities: [{ subscriberId, priority: { ...priority }, device }], + memoryPriority: { ...priority }, + devicePriority: device ? { ...priority } : getMinChunkPriority(), + }; + this.chunks.set(chunkIdString, newChunkEntry); + this.queues.load.insert(chunkIdString, priority); + } else { + const subscriberPriorityIndex = chunkEntry.subscriberPriorities.findIndex( + (entry) => entry.subscriberId === subscriberId + ); + if (subscriberPriorityIndex === -1) { + chunkEntry.subscriberPriorities.push({ subscriberId, priority: { ...priority }, device }); + } else { + chunkEntry.subscriberPriorities[subscriberPriorityIndex][1] = priority; + } + + this.updateChunkPriority(chunkIdString, chunkEntry); + } + } + + /** + * Declares that subscriber `subscriber` no longer requires chunk `chunkId`. + * + * 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 chunkKey = chunkIdToString(chunkId); + const chunkEntry = this.chunks.get(chunkKey); + if (chunkEntry === undefined) { + return; + } + + const index = chunkEntry.subscriberPriorities.findIndex((entry) => entry.subscriberId === subscriberId); + if (index < 0) { + return; + } + + swapRemove(chunkEntry.subscriberPriorities, index); + this.updateChunkPriority(chunkKey, chunkEntry); + } + + /** Get a chunk's data buffer, if it is in memory. */ + getChunkBuffer(chunkId: ChunkId): TypedArray | undefined { + const key = chunkIdToString(chunkId); + const entry = this.chunks.get(key); + if (entry === undefined) { + return undefined; + } + + if (entry.data.state === ChunkState.MEMORY || entry.data.state === ChunkState.DEVICE) { + return entry.data.memory; + } + + return undefined; + } + + /** Get a chunk's texture, if it is on the GPU. */ + getChunkTexture(chunkId: ChunkId): Tex | undefined { + const key = chunkIdToString(chunkId); + const entry = this.chunks.get(key); + if (entry?.data.state === ChunkState.DEVICE) { + return entry.data.texture; + } + return undefined; + } + + /** + * Add a chunk source to this data manager. + * + * Returns the ID that will be used to identify this source in chunk requests and event subscriptions. + */ + addSource(source: IChunkSource): number { + const id = this.sources.length; + this.sources.push({ source, subscribers: [] }); + return id; + } + + /** Returns the dimensions of the given source. */ + getSourceDims(sourceId: number): ExtVolumeDims[] | undefined { + return this.sources[sourceId]?.source.getDims(); + } + + /** Subscribes `subscriber` to data events from the source with id `sourceId` */ + subscribeToSource(subscriber: IDataSubscriber, sourceId: number): boolean { + const sourceEntry = this.sources[sourceId]; + if (sourceEntry === undefined) { + return false; + } + + if (sourceEntry.subscribers.findIndex((s) => s === subscriber) === -1) { + sourceEntry.subscribers.push(subscriber); + } + + return true; + } + + /** Unsubscribes `subscriber` from data events from the source with id `sourceId` */ + unsubscribeFromSource(subscriber: IDataSubscriber, sourceId: number): boolean { + const sourceEntry = this.sources[sourceId]; + if (sourceEntry === undefined) { + return false; + } + + const subscriberIndex = sourceEntry.subscribers.findIndex((s) => s === subscriber); + swapRemove(sourceEntry.subscribers, subscriberIndex); + return true; + } +} diff --git a/src/data_manager/priority_queue.ts b/src/data_manager/priority_queue.ts index e512dc70..d5942975 100644 --- a/src/data_manager/priority_queue.ts +++ b/src/data_manager/priority_queue.ts @@ -9,7 +9,7 @@ const entryToTuple = (entry?: Entry): [P, V] | undefined => entry && const parentIndexOf = (i: number) => Math.ceil(i / 2) - 1; const childIndexOf = (i: number) => i * 2 + 1; -export class PriorityQueue { +export default class PriorityQueue { private keys: Map> = new Map(); private heap: Entry[] = []; diff --git a/src/data_manager/types.ts b/src/data_manager/types.ts new file mode 100644 index 00000000..39125fcd --- /dev/null +++ b/src/data_manager/types.ts @@ -0,0 +1,179 @@ +import type { NumberType, TypedArray } from "../types.ts"; + +/** Settings for the maximum resources a `DataManager` may use (memory, request concurrency, etc.) */ +export type DataManagerLimits = { + /** The maximum amount of data that can be kept in memory, in bytes. */ + size: number; + /** + * The maximum amount of data that can be kept on the GPU, in bytes. + * + * Cannot be set higher than `size`. + */ + deviceSize: number; + /** + * The maximum amount of data that can be kept on the GPU if at least some GPU data has the `PREFETCH` priority. + * + * Setting this lower than `deviceSize` keeps the data manager from prematurely filling up GPU memory with prefetched + * data that may never be used. Cannot be set higher than `deviceSize`. + */ + devicePrefetchSize: number; + /** The maximum number of concurrent requests. */ + concurrentRequests: number; + /** + * The maximum number of concurrent requests if at least one request has the `PREFETCH` priority. + * + * Setting this lower than `concurrentRequests` ensures that some concurrency always remains available for requests + * with the `VISIBLE` priority. Cannot be set higher than `concurrentRequests`. + */ + concurrentPrefetches: number; +}; + +const ONE_GIGABYTE = 1024 * 1024 * 1024; + +export const DEFAULT_DATA_MANAGER_LIMITS: DataManagerLimits = { + size: ONE_GIGABYTE, + deviceSize: ONE_GIGABYTE, + devicePrefetchSize: ONE_GIGABYTE / 2, + concurrentRequests: 10, + concurrentPrefetches: 4, +}; + +/** Enforces validity requirements on `DataManagerLimits`, e.g. `deviceSize` must not be greater than `size`. */ +export const validateDataManagerLimits = (limits: DataManagerLimits): DataManagerLimits => ({ + ...limits, + deviceSize: Math.min(limits.size, limits.deviceSize), + devicePrefetchSize: Math.min(limits.size, limits.deviceSize, limits.devicePrefetchSize), + concurrentPrefetches: Math.min(limits.concurrentRequests, limits.concurrentPrefetches), +}); + +/** Abstraction over the graphics API or library used by `DataManager` to move chunks to/from the GPU. */ +export interface DeviceInterface { + /** + * Type predicate for this interface's *device handle*, for APIs (like WebGPU) that require explicit access to a + * shared resource to perform GPU operations. + */ + isDeviceHandle(val: unknown): val is Dev; + /** Create a texture with the given `data`, `dtype`, and `size` using `deviceHandle`. */ + createTexture(data: TypedArray, dtype: NumberType, size: [number, number, number], deviceHandle: Dev): Tex; + /** Destroy a texture and clean up its GPU resources. */ + destroyTexture(tex: Tex): void; + /** Called when `DataManager` is done with an update cycle, in case the interface has state to clean up. */ + finishUpdate(): void; +} + +/** Reasons that a chunk is tracked by `DataManager`, in increasing order of priority. */ +export const enum ChunkPriorityLevel { + /** + * The chunk was recently loaded but is not currently needed. + * + * A chunk is automatically assigned a priority at this level when one or more subscribers request it (at the + * `VISIBLE` or `PREFETCH` level), its data is loaded, and then all requests are removed. Subscribers shouldn't use + * this level explicitly when making requests: if the chunk isn't already cached, the request will be ignored! + */ + RECENT = 0, + /** + * The chunk isn't currently visible, but is expected to become necessary in the future. + * + * Besides being loaded after `VISIBLE` chunks, chunks prioritized at this level may be subject to lower request + * quotas and/or memory limits, depending on the `DataManager`'s configuration. + */ + PREFETCH = 1, + /** The chunk is currently needed for rendering. */ + VISIBLE = 2, +} + +/** + * The priority at which a chunk is required. + * + * Chunks are prioritized by their `ChunkPriorityLevel`, then within their level by a numeric score. + */ +export type ChunkPriority = { + level: ChunkPriorityLevel; + score: number; +}; + +export const requestLimitForPriority = (limits: DataManagerLimits, { level }: ChunkPriority): number => + level === ChunkPriorityLevel.VISIBLE + ? limits.concurrentRequests + : level === ChunkPriorityLevel.PREFETCH + ? limits.concurrentPrefetches + : 1; + +export const deviceSizeLimitForPriority = (limits: DataManagerLimits, { level }: ChunkPriority): number => + level === ChunkPriorityLevel.VISIBLE + ? limits.deviceSize + : level === ChunkPriorityLevel.PREFETCH + ? limits.devicePrefetchSize + : 0; + +export const getMinChunkPriority = (): ChunkPriority => ({ level: ChunkPriorityLevel.RECENT, score: 0 }); +/** Result is less than `0` if `a > b`, greater than `0` if `a < b`, or `0` if `a === b` */ +export const comparePriority = (a: ChunkPriority, b: ChunkPriority): number => + a.level !== b.level ? b.level - a.level : b.score - a.score; +/** Result is greater than `0` if `a > b`, less than `0` if `a < b`, or `0` if `a === b` */ +export const reverseComparePriority = (a: ChunkPriority, b: ChunkPriority): number => + a.level !== b.level ? a.level - b.level : a.score - b.score; + +export const enum ChunkState { + /** Chunk is queued to be loaded. */ + QUEUED = "queued", + /** Chunk is currently loading. */ + LOADING = "loading", + /** Chunk is cached in memory. */ + MEMORY = "memory", + /** Chunk is cached in memory *and* uploaded to the GPU. */ + DEVICE = "device", + /** Chunk has been temporarily handed off to a worker. Unused, for now. */ + WORKER = "worker", +} + +export type ChunkData = + | { state: ChunkState.QUEUED | ChunkState.WORKER | ChunkState.LOADING } + | { state: ChunkState.MEMORY; memory: TypedArray; dtype: NumberType } + | { state: ChunkState.DEVICE; memory: TypedArray; dtype: NumberType; texture: Tex }; + +export type SubscriberPriority = { + subscriberId: number; + priority: ChunkPriority; + device: boolean; +}; + +/** `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; + /** 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. */ + memoryPriority: ChunkPriority; + /** + * Priority at which the chunk is required in GPU memory. Derived from `subscriberPriorities` on request add/remove. + * + * Must never be greater than `memoryPriority`. + */ + devicePriority: ChunkPriority; +}; + +export type LocalChunkId = { + multiscale: number; + tczyx: [number, number, number, number, number]; +}; + +export type ChunkId = LocalChunkId & { source: number }; + +export const chunkIdToString = (id: ChunkId): string => `${id.source}:${id.multiscale}:${id.tczyx.join(",")}`; + +export const stringToChunkId = (id: string): ChunkId => { + const [sourceId, multiscaleIndex, tczyx] = id.split(":"); + const [t, c, z, y, x] = tczyx.split(","); + return { + source: parseInt(sourceId), + multiscale: parseInt(multiscaleIndex), + tczyx: [parseInt(t), parseInt(c), parseInt(z), parseInt(y), parseInt(x)], + }; +}; + +export type Chunk = { + data: TypedArray; + dtype: T; +}; diff --git a/src/test/priority_queue.test.ts b/src/test/priority_queue.test.ts index cafa0ce1..fe3dee4a 100644 --- a/src/test/priority_queue.test.ts +++ b/src/test/priority_queue.test.ts @@ -1,4 +1,4 @@ -import { PriorityQueue } from "../data_manager/priority_queue.js"; +import PriorityQueue from "../data_manager/priority_queue.js"; const ENTRY_COUNT = 100;