Worker pool (Data Manager 7/9)

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<Dev, Tex> {
           this.tryAbortRequest(entry.data.requestId);
         }
         break;
-      case ChunkState.WORKER:
     }
   }
 
@@ -335,8 +335,14 @@ export default class DataManager<Dev, Tex> {
       }
 
       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<Dev, Tex> {
           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<Dev, Tex> {
   private onChunkLoad(source: number, chunk: Chunk<NumberType>) {
     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<Dev, Tex> {
       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<Dev, Tex> {
     this.updateChunkPriority(chunkKey, chunkEntry);
   }
 
-  /** Get a chunk's data buffer, if it is in memory. */
-  getChunkBuffer(chunkId: ChunkId): TypedArray<NumberType> | 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<Dev, Tex> {
     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);

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] };
+  }
+);

src/data_manager/slot_map.ts

@@ -6,6 +6,7 @@ const isFreeSlot = <T>(slot: Slot<T>): 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<T> {
   private slots: Slot<T>[] = [];
   private head = 0;

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<T extends TypedArray = TypedArray> extends EventDispatcher<BorrowGuardEvents> {
+  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<T> {
+    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 = <T extends TypedArray<NumberType>>(value: T): { [borrowedMarker]: T } => ({
+  [borrowedMarker]: value,
+});
+
+export const getBorrowed = (value: unknown): BorrowedArray<NumberType> | undefined => {
+  if (typeof value === "object" && value !== null && Object.hasOwn(value, borrowedMarker)) {
+    const result = value[borrowedMarker];
+    if (ArrayBuffer.isView(result)) {
+      return result as BorrowedArray<NumberType>;
+    }
+  }
+  return undefined;
+};

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";

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<AddTask>("add");
+ * const pool = new TaskPool(new URL(./worker.ts, import.meta.url));
+ * const three = await pool.runTask(addTask, 1, 2);
+ * ```
+ */
+export const task = <Id extends string, In extends unknown[], Out>(
+  id: Id,
+  handler: (...args: In) => { result: Out; transfer: Transferable[] }
+): Task<Id, In, Out> => {
+  const task = handler as Task<Id, In, Out>;
+  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<NumberType>[] = [];
+    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<WorkerRequest>) => {
+    const [result, transfer] = runTask(data);
+    self.postMessage(result, { transfer });
+  };
+
+  self.postMessage(null);
+};