@navaramap/worker 0.2.1 → 0.3.0

package/src/pool.test.ts

@@ -3,7 +3,7 @@ import workerpool from "workerpool";
 import type Pool from "workerpool/types/Pool";
 
 import { createConcurrencyManager } from "./manager";
-import { initializeWorkerPool, queueTask } from "./pool";
+import { acquireWorkerPool, queueTask } from "./pool";
 
 const workPoolExecStates = (() => {
   const states = {
@@ -39,13 +39,20 @@ vi.spyOn(workerpool, "pool").mockImplementation(
         );
       },
       terminate: () => undefined,
+      stats: () => ({
+        totalWorkers: 1,
+        busyWorkers: 0,
+        idleWorkers: 1,
+        pendingTasks: 0,
+        activeTasks: 0,
+      }),
     }) as unknown as Pool,
 );
 
 // Worker pool is initialized globally.
 const manager = createConcurrencyManager(1);
 const incrementFnMock = vi.spyOn(manager, "increment");
-initializeWorkerPool("https://example.com", manager);
+acquireWorkerPool("https://example.com", manager);
 
 beforeEach(() => {
   workPoolExecStates.reset();

package/src/pool.ts

@@ -16,10 +16,27 @@ export type {
   WorkerPoolStats,
 } from "./recyclingPool";
 
+/** One holder's claim on the shared worker pool; see `acquireWorkerPool`. */
+export type WorkerPoolLease = {
+  /** The pool's concurrency manager, which may be another holder's. */
+  manager: ConcurrencyManager;
+  /** Whether this claim created the pool, so its holder warms it up. */
+  created: boolean;
+  /** Drops this claim; the last release terminates the pool. Idempotent. */
+  release: () => void;
+};
+
+const NOT_ACQUIRED =
+  "The worker pool is not acquired: call acquireWorkerPool() first, and do not use the pool after its last holder released it.";
+
 const {
-  initializeWorkerPool,
-  terminateWorkerPool,
+  acquireWorkerPool,
+  workerPoolHolders,
   warmUpWorkerPool,
+  canWorkerPoolAcceptTasks,
+  workerPoolLoad,
+  reserveWorkerCore,
+  workerPoolReservedSlots,
   worker,
   workerPoolStats,
   workerPoolHeapStats,
@@ -30,43 +47,105 @@ const {
     | {
         pool: RecyclingWorkerPool;
         manager: ConcurrencyManager;
+        maxWorkerHeapBytes: number | undefined;
       }
     | undefined;
+  /** Cores held by threads outside the pool; outlives the pool itself, so a
+   * reservation taken before initialization (or across a re-init) holds. */
+  let coreReservations = 0;
+  let holders = 0;
 
   return {
-    initializeWorkerPool: (
+    /**
+     * Claims the page-wide worker pool, creating it on the first claim. Every
+     * view shares the one pool so their tasks contend for the same cores;
+     * `url`, `manager` and `options` of later claims are ignored while the
+     * pool lives (a differing pool size or heap budget logs a warning).
+     */
+    acquireWorkerPool: (
       url: string,
       manager: ConcurrencyManager,
       options?: RecyclingWorkerPoolOptions,
-    ) => {
-      if (worker) {
-        throw new Error("Worker pool has already been initialized.");
+    ): WorkerPoolLease => {
+      const created = !worker;
+      if (!worker) {
+        worker = {
+          // Recycles worker threads to bound WASM linear memory growth; see
+          // recyclingPool.ts for the policy.
+          pool: new RecyclingWorkerPool(url, manager.total, options),
+          manager,
+          maxWorkerHeapBytes: options?.maxWorkerHeapBytes,
+        };
+        worker.pool.setReservedSlots(coreReservations);
+      } else if (
+        manager.total !== worker.manager.total ||
+        options?.maxWorkerHeapBytes !== worker.maxWorkerHeapBytes
+      ) {
+        console.warn(
+          `The shared worker pool keeps its size (${worker.manager.total}) and worker heap budget (${worker.maxWorkerHeapBytes ?? "default"}); this claim asked for ${manager.total} and ${options?.maxWorkerHeapBytes ?? "default"}.`,
+        );
       }
+      holders++;
 
-      worker = {
-        // Recycles worker threads to bound WASM linear memory growth; see
-        // recyclingPool.ts for the policy.
-        pool: new RecyclingWorkerPool(url, manager.total, options),
-        manager,
+      let released = false;
+      return {
+        manager: worker.manager,
+        created,
+        release: () => {
+          if (released) return;
+          released = true;
+          holders--;
+          if (holders === 0 && worker) {
+            worker.pool.terminate();
+            worker = undefined;
+          }
+        },
       };
     },
-    terminateWorkerPool: () => {
-      if (worker) {
-        worker.pool.terminate();
-        worker = undefined;
-      }
-    },
+    /** Live claims on the pool; see `acquireWorkerPool`. */
+    workerPoolHolders: () => holders,
     /**
      * Pre-initializes WASM on the workers so that upcoming tasks don't pay
      * the initialization latency. Warms a single worker and resolves once
      * it is ready; the rest warm in the background.
      */
     warmUpWorkerPool: () => {
-      invariant(worker, "initializeWorkerPool() must be invoked first.");
+      invariant(worker, NOT_ACQUIRED);
       return worker.pool.warmUp();
     },
+    /**
+     * Whether at least one worker is under its WASM heap budget. While false
+     * the engine's worker-task intake must pause so the draining workers
+     * reach a task boundary and recycle.
+     */
+    canWorkerPoolAcceptTasks: () => worker?.pool.canAcceptTasks() ?? true,
+    /**
+     * Tasks the pool's workers have running or queued, including tasks
+     * cancelled while running; see RecyclingWorkerPool.load.
+     */
+    workerPoolLoad: () => worker?.pool.load() ?? 0,
+    /**
+     * Hold one core for a thread outside the pool (a font worker, a splat
+     * decoder) while it works: one pool slot stops taking tasks and its
+     * worker idles once its queued tasks finish. Returns the release
+     * function; calling it again is a no-op. At least one slot always keeps
+     * taking tasks.
+     */
+    reserveWorkerCore: (): (() => void) => {
+      coreReservations++;
+      worker?.pool.setReservedSlots(coreReservations);
+      let released = false;
+      return () => {
+        if (released) return;
+        released = true;
+        coreReservations--;
+        worker?.pool.setReservedSlots(coreReservations);
+      };
+    },
+    /** Pool slots kept free for threads outside the pool right now. */
+    workerPoolReservedSlots: () => worker?.pool.reservedSlots() ?? 0,
     worker: () => {
-      invariant(worker, "initializeWorkerPool() must be invoked first.");
+      invariant(worker, NOT_ACQUIRED);
       return {
         pool: worker.pool,
         manager: worker.manager,
@@ -81,21 +160,18 @@ const {
 })();
 
 export {
-  initializeWorkerPool,
-  terminateWorkerPool,
+  acquireWorkerPool,
+  workerPoolHolders,
   warmUpWorkerPool,
+  canWorkerPoolAcceptTasks,
+  workerPoolLoad,
+  reserveWorkerCore,
+  workerPoolReservedSlots,
   workerPoolStats,
   workerPoolHeapStats,
   probeWorkerPoolHeap,
 };
 
-export const canWorkerProcessImmediately = () => {
-  const { pool, manager } = worker();
-  // Also pauses intake while every worker is over its WASM heap budget
-  // (draining); a recycled worker restores capacity moments later.
-  return manager.canIncrement() && pool.canAcceptTasks();
-};
-
 export type { ExecOptions } from "workerpool/types/types";
 
 type GetTaskName<Task> =
@@ -136,7 +212,9 @@ export function queueTask<
         // `finally` doesn't work, so use `then()` and `catch()` to ensure that `manager.decrement()` is invoked.
         .then((p) => {
           manager.decrement();
-          return p;
+          // workerpool types the result as the task's raw return type, but it
+          // resolves with the awaited value.
+          return p as MethodReturnType<Task, Name>;
         })
         .catch((e) => {
           manager.decrement();

package/src/poolLease.test.ts

@@ -0,0 +1,152 @@
+import { afterEach, beforeEach, expect, it, vi } from "vitest";
+
+import { createConcurrencyManager, type ConcurrencyManager } from "./manager";
+import {
+  acquireWorkerPool as acquireWorkerPoolRaw,
+  warmUpWorkerPool,
+  workerPoolHolders,
+  type WorkerPoolLease,
+} from "./pool";
+import type { RecyclingWorkerPoolOptions } from "./recyclingPool";
+
+const { created, terminated } = vi.hoisted(() => ({
+  created: { count: 0 },
+  terminated: { count: 0 },
+}));
+
+vi.mock("./recyclingPool", () => ({
+  RecyclingWorkerPool: class {
+    constructor() {
+      created.count++;
+    }
+    setReservedSlots() {}
+    warmUp() {
+      return Promise.resolve();
+    }
+    terminate() {
+      terminated.count++;
+    }
+  },
+}));
+
+// Every lease a test takes, released after it so a failed assertion cannot
+// leave the module-level pool alive for the next test.
+const leases: WorkerPoolLease[] = [];
+function acquireWorkerPool(
+  url: string,
+  manager: ConcurrencyManager,
+  options?: RecyclingWorkerPoolOptions,
+) {
+  const lease = acquireWorkerPoolRaw(url, manager, options);
+  leases.push(lease);
+  return lease;
+}
+
+beforeEach(() => {
+  created.count = 0;
+  terminated.count = 0;
+});
+
+afterEach(() => {
+  for (const lease of leases.splice(0)) lease.release();
+  vi.restoreAllMocks();
+});
+
+it("shares one pool and its manager between holders", () => {
+  const first = createConcurrencyManager(4);
+  const a = acquireWorkerPool("https://example.com", first);
+  const b = acquireWorkerPool(
+    "https://example.com",
+    createConcurrencyManager(2),
+  );
+
+  expect(created.count).toBe(1);
+  expect(a.manager).toBe(first);
+  expect(b.manager).toBe(first);
+
+  a.release();
+  b.release();
+});
+
+it("terminates the pool only when the last holder releases", () => {
+  const a = acquireWorkerPool(
+    "https://example.com",
+    createConcurrencyManager(1),
+  );
+  const b = acquireWorkerPool(
+    "https://example.com",
+    createConcurrencyManager(1),
+  );
+
+  a.release();
+  expect(terminated.count).toBe(0);
+  expect(() => warmUpWorkerPool()).not.toThrow();
+
+  b.release();
+  expect(terminated.count).toBe(1);
+  expect(() => warmUpWorkerPool()).toThrow();
+});
+
+it("ignores a repeated release so it cannot drop another holder's claim", () => {
+  const a = acquireWorkerPool(
+    "https://example.com",
+    createConcurrencyManager(1),
+  );
+  const b = acquireWorkerPool(
+    "https://example.com",
+    createConcurrencyManager(1),
+  );
+
+  a.release();
+  a.release();
+  expect(terminated.count).toBe(0);
+
+  b.release();
+  expect(terminated.count).toBe(1);
+});
+
+it("creates a fresh pool after every holder released", () => {
+  acquireWorkerPool(
+    "https://example.com",
+    createConcurrencyManager(1),
+  ).release();
+  const next = createConcurrencyManager(3);
+  const lease = acquireWorkerPool("https://example.com", next);
+
+  expect(created.count).toBe(2);
+  expect(lease.manager).toBe(next);
+
+  lease.release();
+});
+
+it("tells only the creating claim that it created the pool", () => {
+  const a = acquireWorkerPool(
+    "https://example.com",
+    createConcurrencyManager(1),
+  );
+  const b = acquireWorkerPool(
+    "https://example.com",
+    createConcurrencyManager(1),
+  );
+
+  expect(a.created).toBe(true);
+  expect(b.created).toBe(false);
+  expect(workerPoolHolders()).toBe(2);
+});
+
+it("warns when a later claim asks for different pool settings", () => {
+  const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined);
+  acquireWorkerPool("https://example.com", createConcurrencyManager(4), {
+    maxWorkerHeapBytes: 256,
+  });
+
+  acquireWorkerPool("https://example.com", createConcurrencyManager(4), {
+    maxWorkerHeapBytes: 256,
+  });
+  expect(warn).not.toHaveBeenCalled();
+
+  acquireWorkerPool("https://example.com", createConcurrencyManager(2), {
+    maxWorkerHeapBytes: 128,
+  });
+  expect(warn).toHaveBeenCalledOnce();
+});

package/src/recyclingPool.test.ts

@@ -17,6 +17,10 @@ type MockExecution = {
 
 type MockPool = {
   executions: MockExecution[];
+  /** Busy / queued counts the worker reports, independent of the slot's
+   * own `inflight` (e.g. a cancelled task still running). */
+  activeTasks: number;
+  pendingTasks: number;
   exec: ReturnType<typeof vi.fn>;
   terminate: ReturnType<typeof vi.fn>;
   stats: () => ReturnType<Pool["stats"]>;
@@ -29,6 +33,8 @@ let pools: MockPool[] = [];
 vi.spyOn(workerpool, "pool").mockImplementation(() => {
   const mock: MockPool = {
     executions: [],
+    activeTasks: 0,
+    pendingTasks: 0,
     exec: vi.fn(
       (method: string) =>
         new workerpool.Promise(
@@ -57,8 +63,8 @@ vi.spyOn(workerpool, "pool").mockImplementation(() => {
       totalWorkers: 1,
       busyWorkers: 0,
       idleWorkers: 1,
-      pendingTasks: 0,
-      activeTasks: 0,
+      pendingTasks: mock.pendingTasks,
+      activeTasks: mock.activeTasks,
     }),
   };
   pools.push(mock);
@@ -115,6 +121,91 @@ it("dispatches tasks to the least-loaded slot", () => {
   expect(pools[1].exec).toHaveBeenCalledTimes(1);
 });
 
+it("does not dispatch to a worker still running a cancelled task", async () => {
+  const pool = new RecyclingWorkerPool("https://example.com", 2);
+
+  const cancelled = pool.exec("test");
+  cancelled.cancel();
+  await cancelled.catch(() => undefined);
+  respondToProbe(pools[0], 0);
+  // The slot's own count is back to 0, but its worker is still busy.
+  pools[0].activeTasks = 1;
+
+  pool.exec("test");
+  const tests = (p: MockPool) =>
+    p.executions.filter((e) => e.method === "test").length;
+  expect(tests(pools[0])).toBe(1);
+  expect(tests(pools[1])).toBe(1);
+});
+
+it("sums the slots' load, counting a cancelled task still running", async () => {
+  const pool = new RecyclingWorkerPool("https://example.com", 2);
+  expect(pool.load()).toBe(0);
+
+  const cancelled = pool.exec("test");
+  pool.exec("test");
+  expect(pool.load()).toBe(2);
+
+  cancelled.cancel();
+  await cancelled.catch(() => undefined);
+  respondToProbe(pools[0], 0);
+  expect(pool.load()).toBe(1);
+  // Its worker is still running it.
+  pools[0].activeTasks = 1;
+  expect(pool.load()).toBe(2);
+});
+
+it("leaves the pool's own probes out of its load", async () => {
+  const pool = new RecyclingWorkerPool("https://example.com", 1);
+
+  pool.exec("test");
+  settleTask(pools[0]);
+  await vi.advanceTimersByTimeAsync(0);
+  expect(nextPending(pools[0], PROBE)).toBeDefined();
+  expect(pool.load()).toBe(0);
+});
+
+it("keeps reserved slots free of new tasks, leaving one slot open", () => {
+  const pool = new RecyclingWorkerPool("https://example.com", 3);
+  const tests = (p: MockPool) =>
+    p.executions.filter((e) => e.method === "test").length;
+
+  // Slot 0 is busy, so the idle slots are the ones reserved.
+  pool.exec("test");
+  pool.setReservedSlots(2);
+  expect(pool.reservedSlots()).toBe(2);
+  pool.exec("test");
+  pool.exec("test");
+  expect(tests(pools[0])).toBe(3);
+  expect(tests(pools[1])).toBe(0);
+  expect(tests(pools[2])).toBe(0);
+
+  pool.setReservedSlots(5);
+  expect(pool.reservedSlots()).toBe(2);
+
+  pool.setReservedSlots(0);
+  expect(pool.reservedSlots()).toBe(0);
+  pool.exec("test");
+  expect(tests(pools[1]) + tests(pools[2])).toBe(1);
+});
+
+it("reports no capacity while every unreserved slot is draining", async () => {
+  const pool = new RecyclingWorkerPool("https://example.com", 2, {
+    maxWorkerHeapBytes: 100,
+  });
+  pool.exec("test");
+  pool.setReservedSlots(1);
+  expect(pool.canAcceptTasks()).toBe(true);
+
+  // The open slot goes over budget while its next task still runs.
+  pool.exec("test");
+  settleTask(pools[0]);
+  await vi.advanceTimersByTimeAsync(0);
+  respondToProbe(pools[0], 200);
+  await vi.advanceTimersByTimeAsync(0);
+  expect(pool.canAcceptTasks()).toBe(false);
+});
+
 it("probes heap usage after a settled task and recycles once over budget", async () => {
   const pool = new RecyclingWorkerPool("https://example.com", 1, {
     maxWorkerHeapBytes: 100,
@@ -262,6 +353,17 @@ it("warmUp() warms a single worker first, then the rest in the background", asyn
   expect(nextPending(pools[1], "warmUp")).toBeDefined();
 });
 
+it("warmUp() starts no background warm-up once terminated", async () => {
+  const pool = new RecyclingWorkerPool("https://example.com", 2);
+
+  const warmed = pool.warmUp();
+  pool.terminate();
+  nextPending(pools[0], "warmUp")?.resolve();
+  await warmed;
+
+  expect(pools[1].executions).toHaveLength(0);
+});
+
 it("warmUp() falls back to the remaining workers when the first fails", async () => {
   const pool = new RecyclingWorkerPool("https://example.com", 2);
 

package/src/recyclingPool.ts

@@ -18,8 +18,9 @@ import type { ExecOptions } from "workerpool/types/types";
 // bounded even under sustained load (no pool-wide idle moment is required)
 // and never discards more than one warm WASM instance at a time. The
 // trade-off versus a single shared queue is possible head-of-line blocking
-// within a slot, but upstream backpressure (ConcurrencyManager) keeps
-// per-slot queues shallow.
+// within a slot, but upstream backpressure keeps per-slot queues shallow:
+// the engine admits only about two tasks per worker, and its intake pauses
+// while every slot is draining (see canAcceptTasks).
 //
 // A burst that stays under the heap budget would otherwise keep its heaps
 // forever, so an idle sweep additionally recycles every used slot once the
@@ -62,6 +63,12 @@ type Slot = {
   tasksSinceSpawn: number;
   /** Dispatched but not yet settled tasks, including warm-ups and probes. */
   inflight: number;
+  /** Warm-ups and probes in flight (the pool's own traffic, part of
+   * `inflight`). */
+  internal: number;
+  /** Kept free for a thread outside the pool: receives no new task, so its
+   * worker idles once the tasks it has finish (see setReservedSlots). */
+  reserved: boolean;
   /** Set once the slot exceeds its heap budget; drained slots recycle. */
   draining: boolean;
   /** Whether a heap probe is already in flight (probes are coalesced). */
@@ -84,6 +91,19 @@ export type WorkerPoolHeapStats = {
   maxWorkerHeapBytes: number;
 };
 
+/**
+ * Tasks a slot's worker still has ahead of a new one. `inflight` drops as
+ * soon as a task is cancelled, but a task cancelled while running keeps the
+ * worker busy until it finishes (the worker-side abort listener lets it run
+ * to completion), so the worker's own busy/queued counts are consulted too.
+ * Cancelled tasks still queued are skipped instantly once the worker frees;
+ * counting them only steers new work away from a busy worker.
+ */
+const slotLoad = (slot: Slot): number => {
+  const { activeTasks, pendingTasks } = slot.pool.stats();
+  return Math.max(slot.inflight, activeTasks + pendingTasks);
+};
+
 const createPool = (url: string) =>
   workerpool.pool(url, {
     // One worker per pool so each worker can be recycled independently.
@@ -116,6 +136,8 @@ export class RecyclingWorkerPool {
       pool: createPool(url),
       tasksSinceSpawn: 0,
       inflight: 0,
+      internal: 0,
+      reserved: false,
       draining: false,
       probing: false,
       lastHeapBytes: undefined,
@@ -189,7 +211,62 @@ export class RecyclingWorkerPool {
    * least-loaded draining slot).
    */
   canAcceptTasks(): boolean {
-    return this.slots.some((slot) => !slot.draining);
+    return this.slots.some((slot) => !slot.draining && !slot.reserved);
+  }
+
+  /**
+   * Keep `count` slots free for threads outside the pool (clamped so one
+   * slot always takes tasks): a reserved slot receives no new task, so its
+   * worker idles and leaves its core to them once its queued tasks finish.
+   * Newly reserved slots are the least loaded ones, which free soonest.
+   */
+  setReservedSlots(count: number): void {
+    const target = Math.max(0, Math.min(count, this.slots.length - 1));
+    let reserved = this.reservedSlots();
+    while (reserved < target) {
+      let least: Slot | undefined;
+      let leastLoad = Infinity;
+      for (const slot of this.slots) {
+        if (slot.reserved) continue;
+        const load = slotLoad(slot);
+        if (load < leastLoad) {
+          least = slot;
+          leastLoad = load;
+        }
+      }
+      if (!least) break;
+      least.reserved = true;
+      reserved++;
+    }
+    for (const slot of this.slots) {
+      if (reserved <= target) break;
+      if (slot.reserved) {
+        slot.reserved = false;
+        reserved--;
+      }
+    }
+  }
+
+  /** Slots currently kept free (see setReservedSlots). */
+  reservedSlots(): number {
+    let reserved = 0;
+    for (const slot of this.slots) {
+      if (slot.reserved) reserved++;
+    }
+    return reserved;
+  }
+
+  /**
+   * Tasks handed to the pool that the workers still have, summed over the
+   * slots: running and queued, including tasks cancelled while running (see
+   * slotLoad). The pool's own warm-ups and probes are left out.
+   */
+  load(): number {
+    let load = 0;
+    for (const slot of this.slots) {
+      load += Math.max(0, slotLoad(slot) - slot.internal);
+    }
+    return load;
   }
 
   stats(): WorkerPoolStats {
@@ -253,14 +330,23 @@ export class RecyclingWorkerPool {
   }
 
   private pickSlot(): Slot {
+    // Reserved slots never take a task (one slot always stays unreserved).
     // Prefer slots still under their heap budget; fall back to a draining
     // slot only when every slot is draining at once (it recycles once it
     // finally drains, so this cannot starve the recycling).
-    const candidates = this.slots.filter((slot) => !slot.draining);
-    const eligible = candidates.length > 0 ? candidates : this.slots;
-    return eligible.reduce((least, slot) =>
-      slot.inflight < least.inflight ? slot : least,
-    );
+    const open = this.slots.filter((slot) => !slot.reserved);
+    const candidates = open.filter((slot) => !slot.draining);
+    const eligible = candidates.length > 0 ? candidates : open;
+    let least = eligible[0];
+    let leastLoad = slotLoad(least);
+    for (const slot of eligible) {
+      const load = slotLoad(slot);
+      if (load < leastLoad) {
+        least = slot;
+        leastLoad = load;
+      }
+    }
+    return least;
   }
 
   private onSettled(slot: Slot): void {
@@ -297,9 +383,14 @@ export class RecyclingWorkerPool {
   // perpetuate itself through the work it triggers.
 
   private warmSlot(slot: Slot): Promise<void> {
+    // A terminated workerpool still spawns a worker on exec, and the
+    // background warm-ups can start after terminate().
+    if (this.terminated) return Promise.resolve();
     slot.inflight++;
+    slot.internal++;
     const settle = () => {
       slot.inflight--;
+      slot.internal--;
       // Same as the probe: if the slot drained while this warm-up was in
       // flight, its settlement is the last task boundary, so recycle here.
       if (!this.terminated && slot.draining && slot.inflight === 0) {
@@ -336,9 +427,11 @@ export class RecyclingWorkerPool {
     if (slot.probing) return;
     slot.probing = true;
     slot.inflight++;
+    slot.internal++;
     const settle = (overBudget: boolean) => {
       slot.probing = false;
       slot.inflight--;
+      slot.internal--;
       if (this.terminated) return;
       if (overBudget) {
         slot.draining = true;