@atolljs/core 0.1.0

package/src/pool/workerClient.ts

@@ -0,0 +1,181 @@
+import type { ServiceMethod, TaskRunner } from '../service';
+import type { SharedAccess, SharedMemory, SharedSpec } from '../contract/sharedMemory';
+import type { MemoryConfig, Prettify } from '../contract/types';
+import { WorkerPool, type RunOptions } from './workerPool';
+import type {
+  WorkerDefinition,
+  WorkerMethodMap,
+} from '../worker/defineWorker';
+
+/** Maps a worker method's declared type to its async client signature. */
+type Sig<F> = F extends ServiceMethod<infer A extends any[], infer R>
+  ? (...args: A) => Promise<R>
+  : F extends { run: (...args: infer A extends any[]) => infer R }
+    ? (...args: A) => Promise<Awaited<R>>
+    : F extends (...args: infer A extends any[]) => infer R
+      ? (...args: A) => Promise<Awaited<R>>
+      : never;
+
+/** Flat method map → client signatures. */
+export type ClientMethods<M extends WorkerMethodMap> = { [K in keyof M]: Sig<M[K]> };
+
+/**
+ * The typed client surface a {@link WorkerDefinition} exposes: flat methods
+ * plus one nested client per service namespace.
+ */
+export type WorkerMethods<W extends WorkerDefinition> = ClientMethods<W['methods']> & {
+  [K in keyof W['services']]: W['services'][K] extends infer M extends WorkerMethodMap
+    ? ClientMethods<M>
+    : never;
+} & {
+  /** Per-call options — `client.with({ signal, timeout }).method(args)`. */
+  with(options: RunOptions): WorkerMethods<W>;
+};
+
+const makeClientProxy = (
+  path: string[],
+  resolveRunner: () => TaskRunner,
+  reserved?: Record<string, unknown>,
+  options?: RunOptions,
+): unknown => {
+  // Cache nested proxies so `client.method` is referentially stable across
+  // accesses — matters for memoized bindings (e.g. useTask's [source] deps).
+  const children = new Map<string, unknown>();
+  return new Proxy(() => {}, {
+    get(_, prop) {
+      // `then` and symbols return undefined so the proxy is never mistaken
+      // for a thenable (await client, Promise.resolve(client), assimilation).
+      if (prop === 'then' || typeof prop === 'symbol') return undefined;
+      if (prop === 'with') {
+        // client.with({ signal, timeout }) — same surface, calls dispatch
+        // with RunOptions. Root-level only makes sense, but works anywhere.
+        return (opts: RunOptions) =>
+          makeClientProxy(path, resolveRunner, reserved, opts);
+      }
+      if (path.length === 0 && reserved && prop in reserved) {
+        return Reflect.get(reserved, prop);
+      }
+      let child = children.get(prop);
+      if (!child) {
+        child = makeClientProxy([...path, prop], resolveRunner, reserved, options);
+        children.set(prop, child);
+      }
+      return child;
+    },
+    apply(_, __, args: unknown[]) {
+      const runner = resolveRunner();
+      const contract = { taskId: path.join('.') };
+      return runner.dispatch
+        ? runner.dispatch(contract, args, options)
+        : runner.runTask(contract, ...args);
+    },
+  });
+};
+
+/**
+ * Typed Proxy over any {@link TaskRunner} (WorkerPool, SharedWorkerClient,
+ * test stub). The path joins with '.', so `client.query(q)` dispatches
+ * taskId "query" and `client.pricing.reprice(x)` dispatches
+ * "pricing.reprice" — no runtime knowledge of the shape needed.
+ */
+export function workerClient<W extends WorkerDefinition>(
+  runner: TaskRunner | (() => TaskRunner),
+): WorkerMethods<W> {
+  let resolved: TaskRunner | undefined;
+  const resolveRunner =
+    typeof runner === 'function'
+      ? () => (resolved ??= runner())
+      : () => runner;
+  return makeClientProxy([], resolveRunner) as WorkerMethods<W>;
+}
+
+export interface ConnectWorkerConfig<S extends SharedSpec> {
+  /** Omit for a message-only pool — the SharedArrayBuffer check is skipped. */
+  sharedMemory?: SharedMemory<S> & Prettify<SharedAccess<S>>;
+  /**
+   * Bundler-detectable factory `() => new Worker(new URL('./x.worker.ts', import.meta.url), { type: 'module' })`,
+   * or a URL.
+   */
+  worker: (() => Worker) | URL;
+  poolSize?: number | 'auto';
+  memory?: MemoryConfig;
+  /** Max in-flight tasks per worker (default 1); excess calls queue FIFO. */
+  concurrency?: number;
+  /** Max queued calls (default Infinity); a full queue rejects with PoolQueueFullError. */
+  maxQueue?: number;
+  /** Default true — a crashed worker is replaced and its in-flight calls reject. */
+  respawn?: boolean;
+  /** Default per-call timeout ms (`client.with({ timeout })` overrides). */
+  taskTimeout?: number;
+  /** Default true: pool spawns on first method call (or start()). False spawns immediately. */
+  lazy?: boolean;
+}
+
+/** A typed worker client plus pool lifecycle controls. */
+export type WorkerClient<W extends WorkerDefinition, S extends SharedSpec> = WorkerMethods<W> & {
+  /** Spawn the pool now (idempotent). */
+  start(): void;
+  /** Terminate the pool; the next method call re-spawns it. */
+  terminate(): void;
+  readonly pool: WorkerPool<S> | null;
+  readonly sharedMemory: ConnectWorkerConfig<S>['sharedMemory'] | undefined;
+};
+
+/**
+ * One-line worker hookup for the main thread — the WorkerPool boilerplate
+ * (worker factory, shared memory, sizing) collapses into the client itself:
+ *
+ *   import type { IncidentsWorker } from './worker/incidents.worker';
+ *
+ *   export const incidents = connectWorker<IncidentsWorker>({
+ *     sharedMemory: incidentsMemory,
+ *     worker: () => new Worker(new URL('./worker/incidents.worker.ts', import.meta.url), { type: 'module' }),
+ *     poolSize: 1,
+ *   });
+ *
+ *   await incidents.queryIncidents(q);   // typed dispatch to the pool
+ *   incidents.terminate();               // next call lazily re-spawns
+ *
+ * Lazy by default: no Worker is touched until the first method call (or
+ * `start()`), so importing the client is SSR-safe.
+ */
+export function connectWorker<
+  W extends WorkerDefinition,
+  S extends SharedSpec = W extends WorkerDefinition<infer X> ? X : SharedSpec,
+>(config: ConnectWorkerConfig<S>): WorkerClient<W, S> {
+  let pool: WorkerPool<S> | null = null;
+  const spawn = (): WorkerPool<S> =>
+    (pool ??= new WorkerPool<S>({
+      sharedMemory: config.sharedMemory,
+      createWorker:
+        typeof config.worker === 'function'
+          ? config.worker
+          : () => new Worker(config.worker as URL, { type: 'module' }),
+      poolSize: config.poolSize,
+      memory: config.memory,
+      concurrency: config.concurrency,
+      maxQueue: config.maxQueue,
+      respawn: config.respawn,
+      taskTimeout: config.taskTimeout,
+    }));
+
+  if (config.lazy === false) spawn();
+
+  const reserved = {
+    start: () => {
+      spawn();
+    },
+    terminate: () => {
+      pool?.terminate();
+      pool = null;
+    },
+    get pool() {
+      return pool;
+    },
+    get sharedMemory() {
+      return config.sharedMemory;
+    },
+  };
+
+  return makeClientProxy([], spawn, reserved) as WorkerClient<W, S>;
+}

package/src/pool/workerPool.ts

@@ -0,0 +1,438 @@
+import { MemoryManager } from './memory';
+import { PoolQueueFullError, TaskAbortedError, TaskTimeoutError, WorkerCrashedError } from './errors';
+import { SharedAccess, SharedMemory, SharedSpec } from '../contract/sharedMemory';
+import { PoolTasks, TaskContract, TaskMap, TaskResult, WorkerPoolConfig } from '../contract/types';
+import { fmtBytes, scoped } from '../log';
+
+const poolLog = scoped('pool');
+const taskLog = scoped('task');
+
+/** Per-call dispatch controls. */
+export interface RunOptions {
+  /**
+   * Cancels the call: a queued task is removed and rejects with
+   * TaskAbortedError; an in-flight task rejects immediately and the worker's
+   * eventual reply is discarded. JS cannot be interrupted — the worker keeps
+   * running, and the worker's slot is released only when its reply arrives.
+   * Callers needing a hard cancel should `terminate()`.
+   */
+  signal?: AbortSignal;
+  /**
+   * Per-call timeout in ms — overrides the pool's `taskTimeout`. The timer
+   * starts at enqueue, so it bounds queue-wait + execution time together.
+   */
+  timeout?: number;
+}
+
+interface RunningAgg {
+  count: number;
+  mean: number;
+  max: number;
+  add(v: number): void;
+}
+const agg = (): RunningAgg => ({
+  count: 0,
+  mean: 0,
+  max: 0,
+  add(v: number) {
+    this.count++;
+    this.mean += (v - this.mean) / this.count;
+    if (v > this.max) this.max = v;
+  },
+});
+
+export interface PoolStats {
+  workers: number;
+  idle: number;
+  inFlight: number;
+  queued: number;
+  completed: number;
+  failed: number;
+  aborted: number;
+  /** queue→dispatch timings. */
+  waitMs: { count: number; mean: number; max: number };
+  /** dispatch→reply timings. */
+  runMs: { count: number; mean: number; max: number };
+}
+
+/** One dispatched call — first queued, then in-flight on a worker slot. */
+interface TaskEntry {
+  contract: TaskContract<any[], any>;
+  args: any[];
+  enqueuedAt: number;
+  resolve: (v: any) => void;
+  reject: (e: unknown) => void;
+  /** Finalized — settled (or reply already discarded); no further transitions. */
+  done: boolean;
+  /** Set once handed to a worker. */
+  messageId?: number;
+  slot?: WorkerSlot;
+  sentAt?: number;
+  timer?: ReturnType<typeof setTimeout>;
+  signal?: AbortSignal;
+  onSignalAbort?: () => void;
+}
+
+interface WorkerSlot {
+  worker: Worker;
+  inFlight: number;
+}
+
+class WorkerPoolImpl<S extends SharedSpec = SharedSpec> {
+  private slots: WorkerSlot[] = [];
+  private pending = new Map<number, TaskEntry>();
+  private nextMessageId = 0;
+  private queue: TaskEntry[] = [];
+  private memoryManager?: MemoryManager;
+  private workerUrl?: URL;
+  private spawn: () => Worker;
+  private readonly concurrency: number;
+  private readonly maxQueue: number;
+  private readonly respawn: boolean;
+  private readonly taskTimeout?: number;
+  private closed = false;
+  private drainWaiters: Array<() => void> = [];
+  private statsAgg = {
+    completed: 0,
+    failed: 0,
+    /** aborts + timeouts. */
+    aborted: 0,
+    waitMs: agg(),
+    runMs: agg(),
+  };
+  public readonly sharedMemory: (SharedMemory<S> & SharedAccess<S>) | undefined;
+
+  constructor(config: WorkerPoolConfig<S, TaskMap>) {
+    if (config.sharedMemory) {
+      if (typeof SharedArrayBuffer === 'undefined' || (typeof crossOriginIsolated !== 'undefined' && !crossOriginIsolated)) {
+        throw new Error(
+          'WorkerPool requires SharedArrayBuffer in a cross-origin isolated context — serve the page with COOP/COEP headers.'
+        );
+      }
+    }
+    if (!config.workerUrl && !config.createWorker) {
+      throw new Error('WorkerPool requires either `workerUrl` or `createWorker` in its config.');
+    }
+    this.workerUrl = config.workerUrl;
+    this.spawn = config.createWorker ?? (() => new Worker(this.workerUrl!, { type: 'module' }));
+    this.concurrency = config.concurrency ?? 1;
+    this.maxQueue = config.maxQueue ?? Infinity;
+    this.respawn = config.respawn ?? true;
+    this.taskTimeout = config.taskTimeout;
+
+    if (config.sharedMemory) {
+      this.sharedMemory = config.sharedMemory;
+      this.memoryManager = new MemoryManager(config.memory);
+      this.memoryManager.ensureCapacity(this.sharedMemory.totalBytes);
+      this.sharedMemory.bind(this.memoryManager.getBuffer());
+    }
+
+    const requestedSize = config.poolSize ?? 'auto';
+    const poolSize =
+      requestedSize === 'auto' ? (navigator.hardwareConcurrency ?? 4) : requestedSize;
+
+    poolLog.info(
+      this.sharedMemory
+        ? `spawning ${poolSize} worker(s) with ${fmtBytes(this.sharedMemory.totalBytes)} shared buffer`
+        : `spawning ${poolSize} worker(s), message-only`,
+      { workerUrl: config.workerUrl ? String(config.workerUrl) : 'createWorker()' },
+    );
+    poolLog.debug(`tip: setLogLevel('trace') for per-task detail`);
+
+    for (let i = 0; i < poolSize; i++) {
+      this.addWorker(this.spawn());
+    }
+
+    // Install first-class task methods: pool.queryIncidents(q) →
+    // runTask(QueryIncidents, q). Typed through PoolTasks<T> at the
+    // WorkerPool constructor signature.
+    const tasks = config.tasks ?? {};
+    for (const [name, contract] of Object.entries(tasks)) {
+      if (name in this) {
+        poolLog.warn(`task "${name}" collides with a WorkerPool member — not installed`);
+        continue;
+      }
+      (this as Record<string, unknown>)[name] = (...args: unknown[]) =>
+        this.runTask(contract as TaskContract<any[], any>, ...(args as any[]));
+    }
+    if (Object.keys(tasks).length > 0) {
+      poolLog.info(`task methods: ${Object.keys(tasks).join(', ')}`);
+    }
+  }
+
+  /** Adds a spawned worker: one message + one error listener, then the INIT handshake. */
+  private addWorker(worker: Worker): WorkerSlot {
+    const slot: WorkerSlot = { worker, inFlight: 0 };
+    worker.addEventListener('message', (event: MessageEvent) => {
+      this.handleReply(slot, event.data as TaskResult & { messageId?: number });
+    });
+    worker.addEventListener('error', (event: unknown) => {
+      this.handleWorkerError(slot, event);
+    });
+    if (this.memoryManager) {
+      worker.postMessage({ type: 'INIT_MEMORY', memory: this.memoryManager.memory });
+    } else {
+      worker.postMessage({ type: 'INIT' });
+    }
+    this.slots.push(slot);
+    poolLog.debug(`worker spawned, ${this.memoryManager ? 'INIT_MEMORY' : 'INIT'} sent`);
+    return slot;
+  }
+
+  private handleReply(slot: WorkerSlot, data: TaskResult & { messageId?: number }) {
+    const entry = data.messageId !== undefined ? this.pending.get(data.messageId) : undefined;
+    if (!entry) return; // stray reply — discard
+    this.pending.delete(data.messageId!);
+    slot.inFlight--;
+    this.detach(entry);
+    // Orphaned (aborted/timed-out) entries free the slot here but skip
+    // stats and settlement — the caller was already rejected.
+    if (!entry.done) {
+      entry.done = true;
+      this.statsAgg.runMs.add(performance.now() - (entry.sentAt ?? entry.enqueuedAt));
+      if (data.success) {
+        this.statsAgg.completed++;
+        taskLog.debug(`← ${entry.contract.taskId} ok`);
+        entry.resolve(data.result);
+      } else {
+        this.statsAgg.failed++;
+        taskLog.warn(`← ${entry.contract.taskId} failed: ${data.error}`);
+        entry.reject(new Error(data.error));
+      }
+    }
+    this.drainQueue();
+    this.checkDrained();
+  }
+
+  private handleWorkerError(slot: WorkerSlot, event: unknown) {
+    if (this.closed) return; // pool terminated — a late error event can't respawn
+    const e = event as { message?: string; error?: Error } | undefined;
+    const message = e?.message ?? (e?.error ? String(e.error) : 'worker error');
+    poolLog.warn(`worker error: ${message}`);
+
+    // Reject the dead worker's in-flight tasks. Orphaned (already
+    // aborted/timed-out) entries are dropped without re-counting.
+    for (const [id, entry] of [...this.pending]) {
+      if (entry.slot === slot) {
+        this.pending.delete(id);
+        this.detach(entry);
+        if (entry.done) continue;
+        entry.done = true;
+        this.statsAgg.failed++;
+        entry.reject(new WorkerCrashedError(message));
+      }
+    }
+
+    const idx = this.slots.indexOf(slot);
+    if (idx >= 0) this.slots.splice(idx, 1);
+    slot.worker.terminate();
+
+    if (this.respawn) {
+      poolLog.info('respawning crashed worker');
+      this.addWorker(this.spawn());
+    } else if (this.slots.length === 0) {
+      // No workers left to ever serve the queue.
+      for (const entry of this.queue.splice(0)) {
+        entry.done = true;
+        this.detach(entry);
+        this.statsAgg.failed++;
+        entry.reject(new WorkerCrashedError(message));
+      }
+    }
+    this.drainQueue();
+    this.checkDrained();
+  }
+
+  /** Least-busy worker below its concurrency cap, if any. */
+  private pickWorker(): WorkerSlot | undefined {
+    let best: WorkerSlot | undefined;
+    for (const slot of this.slots) {
+      if (slot.inFlight >= this.concurrency) continue;
+      if (!best || slot.inFlight < best.inFlight) best = slot;
+    }
+    return best;
+  }
+
+  private drainQueue() {
+    let slot: WorkerSlot | undefined;
+    while (this.queue.length > 0 && (slot = this.pickWorker())) {
+      const entry = this.queue.shift()!;
+      if (entry.done) continue;
+      this.send(slot, entry);
+    }
+  }
+
+  private send(slot: WorkerSlot, entry: TaskEntry) {
+    const messageId = this.nextMessageId++;
+    entry.messageId = messageId;
+    entry.slot = slot;
+    entry.sentAt = performance.now();
+    this.statsAgg.waitMs.add(performance.now() - entry.enqueuedAt);
+    slot.inFlight++;
+    this.pending.set(messageId, entry);
+    taskLog.debug(`→ ${entry.contract.taskId}`, { worker: this.slots.indexOf(slot), args: entry.args });
+    slot.worker.postMessage({
+      type: 'EXECUTE_TASK',
+      messageId,
+      taskId: entry.contract.taskId,
+      args: entry.args,
+    });
+  }
+
+  /** Detaches signal/timeout listeners from a finalized entry. */
+  private detach(entry: TaskEntry) {
+    if (entry.timer) clearTimeout(entry.timer);
+    if (entry.signal && entry.onSignalAbort) {
+      entry.signal.removeEventListener('abort', entry.onSignalAbort);
+    }
+  }
+
+  /**
+   * Abort/timeout path: a queued task is removed outright; an in-flight task
+   * stays in `pending` as an orphan — the worker is still running it, so the
+   * slot remains occupied until the reply arrives (handleReply frees it).
+   */
+  private cancel(entry: TaskEntry, err: Error) {
+    if (entry.done) return;
+    const qi = this.queue.indexOf(entry);
+    if (qi >= 0) this.queue.splice(qi, 1);
+    entry.done = true;
+    this.detach(entry);
+    this.statsAgg.aborted++;
+    entry.reject(err);
+    this.drainQueue();
+    this.checkDrained();
+  }
+
+  private checkDrained() {
+    if (this.queue.length === 0 && this.pending.size === 0) {
+      const waiters = this.drainWaiters.splice(0);
+      for (const w of waiters) w();
+    }
+  }
+
+  /**
+   * Full dispatch: least-busy scheduling, `concurrency` cap per worker,
+   * FIFO queue bounded by `maxQueue`, abort/timeout, and queue→reply timing.
+   */
+  public dispatch<Args extends any[], Result>(
+    contract: TaskContract<Args, Result>,
+    args: Args,
+    options?: RunOptions,
+  ): Promise<Result> {
+    if (this.closed || this.slots.length === 0) {
+      return Promise.reject(new WorkerCrashedError('WorkerPool has been terminated — no workers available.'));
+    }
+    if (options?.signal?.aborted) {
+      return Promise.reject(new TaskAbortedError());
+    }
+    return new Promise<Result>((resolve, reject) => {
+      const entry: TaskEntry = {
+        contract,
+        args: args as any[],
+        enqueuedAt: performance.now(),
+        resolve,
+        reject,
+        done: false,
+        signal: options?.signal,
+      };
+      if (options?.signal) {
+        entry.onSignalAbort = () => this.cancel(entry, new TaskAbortedError());
+        options.signal.addEventListener('abort', entry.onSignalAbort, { once: true });
+      }
+      const timeout = options?.timeout ?? this.taskTimeout;
+      if (timeout !== undefined) {
+        entry.timer = setTimeout(() => this.cancel(entry, new TaskTimeoutError()), timeout);
+      }
+
+      const slot = this.pickWorker();
+      if (slot) {
+        this.send(slot, entry);
+      } else if (this.queue.length >= this.maxQueue) {
+        entry.done = true;
+        this.detach(entry);
+        reject(new PoolQueueFullError()); // never dispatched — stats untouched
+        return;
+      } else {
+        this.queue.push(entry);
+      }
+    });
+  }
+
+  /**
+   * Dispatches a task to the least-busy worker in the pool.
+   * The task contract guarantees the argument and result types across the thread boundary.
+   */
+  public runTask<Args extends any[], Result>(contract: TaskContract<Args, Result>, ...args: Args): Promise<Result> {
+    return this.dispatch(contract, args);
+  }
+
+  /** Running aggregates — no histograms, just counts and means. */
+  public stats(): PoolStats {
+    return {
+      workers: this.slots.length,
+      idle: this.slots.filter((s) => s.inFlight === 0).length,
+      inFlight: this.pending.size,
+      queued: this.queue.length,
+      completed: this.statsAgg.completed,
+      failed: this.statsAgg.failed,
+      aborted: this.statsAgg.aborted,
+      waitMs: { count: this.statsAgg.waitMs.count, mean: this.statsAgg.waitMs.mean, max: this.statsAgg.waitMs.max },
+      runMs: { count: this.statsAgg.runMs.count, mean: this.statsAgg.runMs.mean, max: this.statsAgg.runMs.max },
+    };
+  }
+
+  /** Resolves once the queue is empty and nothing is in flight, then terminates. */
+  public close(): Promise<void> {
+    if (this.queue.length === 0 && this.pending.size === 0) {
+      this.terminate();
+      return Promise.resolve();
+    }
+    return new Promise((resolve) => {
+      this.drainWaiters.push(() => {
+        this.terminate();
+        resolve();
+      });
+    });
+  }
+
+  public terminate(): void {
+    this.closed = true;
+    poolLog.info(`terminating ${this.slots.length} worker(s)`);
+    const crashed = new WorkerCrashedError('pool terminated');
+    for (const [, entry] of this.pending) {
+      this.detach(entry);
+      if (entry.done) continue; // orphaned — already counted as aborted
+      entry.done = true;
+      this.statsAgg.failed++;
+      entry.reject(crashed);
+    }
+    this.pending.clear();
+    for (const entry of this.queue.splice(0)) {
+      entry.done = true;
+      this.detach(entry);
+      this.statsAgg.failed++;
+      entry.reject(crashed);
+    }
+    for (const slot of this.slots) {
+      slot.worker.terminate();
+    }
+    this.slots = [];
+    this.checkDrained();
+  }
+}
+
+/**
+ * A WorkerPool with first-class task methods baked in from `config.tasks`.
+ * Constructed as `new WorkerPool({ workerUrl, sharedMemory, tasks })` —
+ * `tasks: { queryIncidents: QueryIncidents }` exposes `pool.queryIncidents(q)`
+ * with the contract's declared args/result types.
+ */
+export type WorkerPool<S extends SharedSpec = SharedSpec, T extends TaskMap = TaskMap> =
+  WorkerPoolImpl<S> & PoolTasks<T>;
+
+export const WorkerPool = WorkerPoolImpl as {
+  new <S extends SharedSpec = SharedSpec, T extends TaskMap = TaskMap>(
+    config: WorkerPoolConfig<S, T>
+  ): WorkerPool<S, T>;
+};