diff --git a/src/data_manager/data_manager.ts b/src/data_manager/data_manager.ts index b853c1f5..52594a35 100644 --- a/src/data_manager/data_manager.ts +++ b/src/data_manager/data_manager.ts @@ -18,6 +18,7 @@ import type { DeviceInterface } from "./device_interface.js"; import { dataTypeToSize } from "./utils.js"; import { ChunkSource, type VolumeMetadata, type ExtVolumeDims } from "./sources/chunk_source.js"; import SlotMap from "./slot_map.js"; +import { BorrowGuard } from "./task_pool/borrow_guard.js"; const SUBSCRIBER_ID = Symbol("DataManager.subscriberId"); @@ -146,7 +147,6 @@ export default class DataManager { this.tryAbortRequest(entry.data.requestId); } break; - case ChunkState.WORKER: } } @@ -335,8 +335,14 @@ export default class DataManager { } const { x, y, z, dataType } = dims; - const { memory } = loadEntry.data as { memory: TypedArray }; - const texture = this.deviceInterface.createTexture(memory, dataType, [x, y, z], deviceHandle); + const { memory } = loadEntry.data as { memory: BorrowGuard }; + + const array = memory.get(); + if (array === undefined) { + continue; + } + + const texture = this.deviceInterface.createTexture(array, dataType, [x, y, z], deviceHandle); this.sources[loadId.source].subscribers.forEach((s) => s.onChunkOnGpu?.(loadId, texture)); @@ -386,7 +392,6 @@ export default class DataManager { this.tryAbortRequest(evictEntry.data.requestId); 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})` ); @@ -493,8 +498,8 @@ export default class DataManager { private onChunkLoad(source: number, chunk: Chunk) { const globalId = { ...chunk.id, source }; const key = chunkIdToString(globalId); - const { data: memory, dtype } = chunk; - if (memory.byteLength > this.limits.size) { + const { data: array, dtype } = chunk; + if (array.byteLength > this.limits.size) { console.error(`received chunk ${key} which is larger than the cache limit`); return; } @@ -503,6 +508,7 @@ export default class DataManager { console.error(`received chunk ${key} with invalid source id ${source}`); return; } + const memory = new BorrowGuard(array); const data = { state: ChunkState.MEMORY as const, memory, dtype }; let chunkEntry = this.chunks.get(key); @@ -580,8 +586,8 @@ export default class DataManager { this.updateChunkPriority(chunkKey, chunkEntry); } - /** Get a chunk's data buffer, if it is in memory. */ - getChunkBuffer(chunkId: ChunkId): TypedArray | undefined { + /** Get a chunk's data buffer wrapped in a `BorrowGuard`, if it is in memory. */ + getChunkGuard(chunkId: ChunkId): BorrowGuard | undefined { const key = chunkIdToString(chunkId); const entry = this.chunks.get(key); if (entry === undefined) { @@ -595,6 +601,15 @@ export default class DataManager { return undefined; } + /** + * Get a chunk's data buffer, if it is in memory and not in use on a worker. + * + * Shorthand for `dataManager.getChunkGuard(chunkId)?.get();` + */ + getChunkBuffer(chunkId: ChunkId): TypedArray | undefined { + return this.getChunkGuard(chunkId)?.get(); + } + /** Get a chunk's texture, if it is on the GPU. */ getChunkTexture(chunkId: ChunkId): Tex | undefined { const key = chunkIdToString(chunkId); diff --git a/src/data_manager/indexing_task.ts b/src/data_manager/indexing_task.ts new file mode 100644 index 00000000..00a3c4d5 --- /dev/null +++ b/src/data_manager/indexing_task.ts @@ -0,0 +1,19 @@ +import type { NumberType, TypedArray } from "../types.js"; +import { copyChunk, reorderChunk, type CopyChunkParams } from "./indexing.js"; +import { type BorrowedArray, task } from "./task_pool/index.js"; + +export const copyChunkTask = task( + "copyChunk", + (src: TypedArray | BorrowedArray, dest: BorrowedArray, params: CopyChunkParams) => { + copyChunk(src, dest, params); + return { result: undefined, transfer: [] }; + } +); + +export const reorderChunkTask = task( + "reorderChunk", + (src: TypedArray | BorrowedArray, shape: number[], order: number[], dtype: NumberType) => { + const result = reorderChunk(src, shape, order, dtype); + return { result, transfer: [result.buffer] }; + } +); diff --git a/src/data_manager/slot_map.ts b/src/data_manager/slot_map.ts index cf2703ed..4a7bb10a 100644 --- a/src/data_manager/slot_map.ts +++ b/src/data_manager/slot_map.ts @@ -6,6 +6,7 @@ const isFreeSlot = (slot: Slot): slot is { [NEXT_FREE]: number } => { return slot !== undefined && slot !== null && Object.hasOwn(slot, NEXT_FREE); }; +/** A "map" that assigns values a stable numeric key on insertion that can be used to get or remove the value later. */ export default class SlotMap { private slots: Slot[] = []; private head = 0; diff --git a/src/data_manager/task_pool/borrow_guard.ts b/src/data_manager/task_pool/borrow_guard.ts new file mode 100644 index 00000000..fa191541 --- /dev/null +++ b/src/data_manager/task_pool/borrow_guard.ts @@ -0,0 +1,82 @@ +import type { NumberType, TypedArray } from "../../types.js"; +import EventDispatcher from "../../EventDispatcher.js"; +import { BorrowedArray } from "./types.js"; + +const guardMarker = Symbol.for("TaskPool.borrowGuard"); + +type BorrowGuardEvents = { restored: void }; + +/** + * Wraps a `TypedArray` that may be temporarily transferred to a worker for processing. + * + * A `TypedArray` guarded by this class must be accessed through the `get` method, which returns `undefined` when the + * buffer is on a worker. The guard will trigger the `"restored"` event when it is returned to the main thread. + */ +export class BorrowGuard extends EventDispatcher { + private [guardMarker] = true as const; + private restored = false; + public readonly byteLength: number; + + constructor(private array: T) { + super(); + this.byteLength = array.byteLength; + } + + get borrowed(): boolean { + return this.array.buffer.detached; + } + + get(): T | undefined { + if (this.borrowed) { + return undefined; + } + return this.array; + } + + getAsync(): Promise { + if (this.borrowed) { + return new Promise((resolve) => { + const listener = () => { + if (!this.borrowed) { + resolve(this.array); + } + this.removeEventListener("restored", listener); + }; + this.addEventListener("restored", listener); + }); + } + return Promise.resolve(this.array); + } + + restore(array: T) { + this.array = array; + this.restored = true; + } + + afterRestore() { + if (!this.borrowed && this.restored) { + this.restored = false; + this.dispatchEvent({ type: "restored" }); + } + } + + static isBorrowGuard(value: unknown): value is BorrowGuard { + return typeof value === "object" && value !== null && value[guardMarker]; + } +} + +const borrowedMarker = "TaskPool.borrowedArray"; + +export const markBorrowed = >(value: T): { [borrowedMarker]: T } => ({ + [borrowedMarker]: value, +}); + +export const getBorrowed = (value: unknown): BorrowedArray | undefined => { + if (typeof value === "object" && value !== null && Object.hasOwn(value, borrowedMarker)) { + const result = value[borrowedMarker]; + if (ArrayBuffer.isView(result)) { + return result as BorrowedArray; + } + } + return undefined; +}; diff --git a/src/data_manager/task_pool/index.ts b/src/data_manager/task_pool/index.ts new file mode 100644 index 00000000..7e9c9ee7 --- /dev/null +++ b/src/data_manager/task_pool/index.ts @@ -0,0 +1,4 @@ +export { BorrowGuard } from "./borrow_guard.js"; +export { task, initWorker } from "./runner.js"; +export { taskHandle, TaskPool } from "./task_pool.js"; +export type { Task, TaskHandle, BorrowedArray } from "./types.js"; diff --git a/src/data_manager/task_pool/runner.ts b/src/data_manager/task_pool/runner.ts new file mode 100644 index 00000000..f8b24fc2 --- /dev/null +++ b/src/data_manager/task_pool/runner.ts @@ -0,0 +1,84 @@ +import { serializeError } from "serialize-error"; + +import type { TypedArray, NumberType } from "../../types.js"; +import { getBorrowed } from "./borrow_guard.js"; +import type { Task, WorkerRequest, WorkerResponse } from "./types.js"; + +/** + * Combines a `string` task `id` and a `handler` into a `Task` that can be registered and called with a `TaskPool`. + * + * Modules that define tasks that can be run on a worker should export the task handler wrapped in this function: + * + * ``` + * // add_task.ts + * export const addTask = task("add", (a: number, b: number) => a + b); + * ``` + * + * To create a task pool and run this task, first create a worker module that passes the task to `initWorker`: + * + * ``` + * // worker.ts + * import { addTask } from "./add_task.js"; + * initWorker([addTask]); + * ``` + * + * Passing the url of this module to a new `TaskPool` creates a pool that can recognize and run this task. Main thread + * code should import *only* the task's type and use it to create a `TaskHandle`, which can be used to run the task: + * + * ``` + * // main.ts + * import type { AddTask } from "./add_task.js"; + * const addTask = taskHandle("add"); + * const pool = new TaskPool(new URL(./worker.ts, import.meta.url)); + * const three = await pool.runTask(addTask, 1, 2); + * ``` + */ +export const task = ( + id: Id, + handler: (...args: In) => { result: Out; transfer: Transferable[] } +): Task => { + const task = handler as Task; + task.taskId = id; + return task; +}; + +export const initWorker = (tasks: Task[]) => { + const handlers = new Map(tasks.map((task) => [task.taskId, task])); + + const runTask = ({ id, task, args }: WorkerRequest): [WorkerResponse, Transferable[]] => { + const handler = handlers.get(task); + + // Extract borrowed `TypedArray`s, which must be returned to the main thread on completion + const borrows: TypedArray[] = []; + const borrowedBuffers: ArrayBuffer[] = []; + const processedArgs = args.map((arg) => { + const borrowed = getBorrowed(arg); + if (borrowed !== undefined) { + borrows.push(borrowed); + borrowedBuffers.push(borrowed.buffer); + return borrowed; + } + return arg; + }); + + if (handler === undefined) { + const result = serializeError(new Error(`TaskPool: tried to run nonexistent task ${task}`)); + return [{ id, task, borrows, result, error: true }, borrowedBuffers]; + } + + try { + const { result, transfer } = handler(...processedArgs); + return [{ id, task, borrows, result, error: false }, [...borrowedBuffers, ...transfer]]; + } catch (e) { + const result = serializeError(e); + return [{ id, task, borrows, result, error: true }, borrowedBuffers]; + } + }; + + self.onmessage = ({ data }: MessageEvent) => { + const [result, transfer] = runTask(data); + self.postMessage(result, { transfer }); + }; + + self.postMessage(null); +}; diff --git a/src/data_manager/task_pool/task_pool.ts b/src/data_manager/task_pool/task_pool.ts new file mode 100644 index 00000000..da250702 --- /dev/null +++ b/src/data_manager/task_pool/task_pool.ts @@ -0,0 +1,137 @@ +import { deserializeError } from "serialize-error"; +import SlotMap from "../slot_map.js"; +import type { Task, TaskHandle, WorkerResponse, WorkerRequest, TaskArgs } from "./types.js"; +import { BorrowGuard, markBorrowed } from "./borrow_guard.js"; + +export const taskHandle = ( + id: T["taskId"], + transfer?: (...args: Parameters) => Transferable[] +): TaskHandle, ReturnType["result"]> => { + return { id, transfer }; +}; + +type StoredPromise = { + resolve: (value: unknown) => void; + reject: (reason: unknown) => void; + borrows: BorrowGuard[]; +}; + +type WaitingTask = { + promiseId: number; + taskHandle: TaskHandle; + args: In; +}; + +export class TaskPool { + private workerCount = 0; + private idleWorkers: Worker[] = []; + private promises: SlotMap = new SlotMap(); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + private queue: WaitingTask[] = []; + + constructor( + private readonly workerPath: URL, + public maxWorkers = navigator.hardwareConcurrency ?? 4 + ) {} + + private handleWorkerMessage(message: WorkerResponse>) { + const promise = this.promises.remove(message.id); + if (promise === undefined) { + console.error( + `TaskPool: received worker response for nonexistent task ID ${message.id} with type "${message.task}"` + ); + return; + } + + const borrowsExpected = promise.borrows.length; + const borrowsReceived = message.borrows.length; + if (borrowsExpected === borrowsReceived) { + for (let i = 0; i < borrowsReceived; i++) { + promise.borrows[i].restore(message.borrows[i]); + } + } else { + console.error( + `TaskPool: number of borrows returned from task "${message.task}" with ID ${message.id} did not match (expected ${borrowsExpected}, received ${borrowsReceived})` + ); + } + + if (message.error) { + promise.reject(deserializeError(message.result)); + } else { + promise.resolve(message.result); + } + + this.submitTasks(); + promise.borrows.forEach((borrow) => borrow.afterRestore()); + } + + private addWorker() { + const worker = new Worker(this.workerPath, { type: "module" }); + let ready = false; + worker.onmessage = (message) => { + this.idleWorkers.push(worker); + if (!ready) { + ready = true; + this.submitTasks(); + } else { + this.handleWorkerMessage(message.data); + } + }; + this.workerCount += 1; + } + + private submitTasks() { + let i = 0; + while (i < this.queue.length) { + const task = this.queue[i]; + const borrowedBuffers: ArrayBuffer[] = []; + const args: unknown[] = []; + // Handle borrowed buffers + for (const arg of task.args) { + if (BorrowGuard.isBorrowGuard(arg)) { + const array = arg.get(); + // Task is not ready -- some resource it depends on is in use by another task + if (array === undefined) { + i++; + continue; + } + args.push(markBorrowed(array)); + borrowedBuffers.push(array.buffer); + } else { + args.push(arg); + } + } + + const worker = this.idleWorkers.pop(); + if (worker === undefined) { + // We have a ready task but we're out of workers to run it on + if (this.workerCount < this.maxWorkers) { + this.addWorker(); + } + return; + } + + this.queue.splice(i, 1); + const transfer = [...borrowedBuffers, ...(task.taskHandle.transfer?.(...task.args) ?? [])]; + worker.postMessage( + { + task: task.taskHandle.id, + id: task.promiseId, + args, + } as WorkerRequest, + { transfer } + ); + } + } + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + public runTask(taskHandle: TaskHandle, ...args: TaskArgs): Promise { + const borrows = args.filter(BorrowGuard.isBorrowGuard); + const promise = new Promise((resolve, reject) => { + const promiseId = this.promises.insert({ resolve, reject, borrows }); + this.queue.push({ taskHandle, promiseId, args }); + }); + this.submitTasks(); + return promise as Promise; + } +} diff --git a/src/data_manager/task_pool/types.ts b/src/data_manager/task_pool/types.ts new file mode 100644 index 00000000..7448226a --- /dev/null +++ b/src/data_manager/task_pool/types.ts @@ -0,0 +1,69 @@ +import type { ErrorObject } from "serialize-error"; + +import type { TypedArray, NumberType } from "../../types.js"; +import type { BorrowGuard } from "./borrow_guard.js"; + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +export type Task = { + taskId: Id; + (...args: In): { result: Out; transfer: Transferable[] }; +}; + +export type TaskHandle = { + id: string; + transfer?: (...args: In) => Transferable[]; + marker?: Out; +}; + +/** + * Marker type for a `TypedArray` that has been *borrowed* from the main thread by a task handler. + * + * This type is functionally identical to the corresponding `TypedArray`. However, when it is used as an argument to a + * *task handler*, the task must be invoked from the main thread using a matching `BorrowGuard`: + * + * ```typescript + * // worker code + * const fooTask = task("foo", (borrow: BorrowedArray<"uint8">) => { ... }); + * export type FooTask = typeof fooTask; + * registerTask(fooTask); + * + * // main thread code + * const fooTask = taskHandle("foo"); + * const pool = new TaskPool(); + * // Because the handler takes a `BorrowedArray`, the task must be invoked on the main thread with a `BorrowGuard`: + * const guard = new BorrowGuard(new Uint8Array([1, 2, 3, 4, 5])); + * const result = await pool.runTask("foo", guard); + * ``` + * + * When the task runs, the array inside the `BorrowGuard` is automatically *transferred* to the assigned worker, where + * it can be safely modified in-place. When the task completes, the array is transferred back into the `BorrowGuard`. + */ +export type BorrowedArray = TypedArray & { borrowed: true }; + +/** Converts all `BorrowedArray`s in a task handler's argument list to the corresponding `BorrowGuard`s */ +export type TaskArgs = T extends [infer E, ...infer R] + ? E extends BorrowedArray + ? [BorrowGuard>, ...TaskArgs] + : [E, ...TaskArgs] + : []; + +export type WorkerRequest = { + id: number; + task: T["taskId"]; + args: Parameters; +}; + +export type WorkerResponse = { + id: number; + task: T["taskId"]; + borrows: TypedArray[]; +} & ( + | { + error: false; + result: ReturnType["result"]; + } + | { + error: true; + result: ErrorObject; + } +); diff --git a/src/data_manager/types.ts b/src/data_manager/types.ts index c5375341..883c40f2 100644 --- a/src/data_manager/types.ts +++ b/src/data_manager/types.ts @@ -1,4 +1,5 @@ import type { NumberType, TypedArray } from "../types.ts"; +import type { BorrowGuard } from "./task_pool/borrow_guard.js"; /** Settings for the maximum resources a `DataManager` may use (memory, request concurrency, etc.) */ export type DataManagerLimits = { @@ -108,15 +109,13 @@ export const enum ChunkState { 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", } type ChunkEntryData = - | { state: ChunkState.QUEUED | ChunkState.WORKER } + | { state: ChunkState.QUEUED } | { state: ChunkState.LOADING; requestId: number } - | { state: ChunkState.MEMORY; memory: TypedArray; dtype: NumberType } - | { state: ChunkState.DEVICE; memory: TypedArray; dtype: NumberType; texture: Tex }; + | { state: ChunkState.MEMORY; memory: BorrowGuard; dtype: NumberType } + | { state: ChunkState.DEVICE; memory: BorrowGuard; dtype: NumberType; texture: Tex }; export type SubscriberPriority = { subscriberId: number; diff --git a/tsconfig.json b/tsconfig.json index 02a5c51b..abd0857c 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -4,7 +4,7 @@ "allowSyntheticDefaultImports": true, "declaration": true, "declarationDir": "es/types", - "lib": ["DOM", "ES2022", "ES6", "DOM.Iterable", "ScriptHost", "ES2016.Array.Include"], + "lib": ["DOM", "ES2022", "ES6", "DOM.Iterable", "ScriptHost", "ES2016.Array.Include", "ES2024.ArrayBuffer"], "module": "Node16", "moduleResolution": "node16", "noImplicitAny": false,