Something went wrong. Try again.
Offload functions to worker threads with shared memory primitives for Node.js.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256import { parentPort, MessagePort } from 'node:worker_threads';import type { Transferable } from 'node:worker_threads';import { registry } from './registry.ts';import { deserializeArg, serializeArg } from './shared/reconstruct.ts';import { collectTransferables } from './transfer.ts';import { pipeIterable, CANCEL, DEFAULT_HIGH_WATER, deserializeStreamHandle, isSerializedStreamHandle } from './pipe.ts';import type { StreamHandle, SerializedStreamHandle } from './pipe.ts';
const imported = new Set<string>();const taskCache = new Map<number, unknown>();
type Fn = (...args: unknown[]) => unknown;const fnCache = new Map<string, Fn>();
function isTaskArg(arg: unknown): arg is { __task__: number; id: string; args: unknown[] } { return typeof arg === 'object' && arg !== null && '__task__' in arg;}
function needsAsyncResolve(arg: unknown): boolean { return arg instanceof MessagePort || isSerializedStreamHandle(arg) || isTaskArg(arg);}
// Returns the function synchronously when cached — callers must not await// unconditionally, or they pay a microtask hop for every dispatch.function resolveFn(id: string): Fn | Promise<Fn> { const cached = fnCache.get(id); if (cached) return cached; return resolveFnSlow(id);}
async function resolveFnSlow(id: string): Promise<Fn> { const url = id.slice(0, id.lastIndexOf('#')); if (!imported.has(url)) { await import(url); imported.add(url); } const fn = registry.get(id) as Fn | undefined; if (!fn) throw new Error(`Moroutine not found: ${id}`); fnCache.set(id, fn); return fn;}
function portToAsyncIterable<T>(handle: StreamHandle): AsyncIterable<T> { const { port, highWater, readySignal } = handle; const inflight = handle.flags.fields.inflight; const state = handle.flags.fields.state;
const queue: T[] = []; let done = false; let error: Error | null = null; let waiting: (() => void) | null = null;
port.on('message', (msg: { value?: unknown; done?: boolean; error?: Error }) => { if (msg.error) { error = msg.error; done = true; if (waiting) { waiting(); waiting = null; } return; } if (msg.done) { done = true; if (waiting) { waiting(); waiting = null; } return; } queue.push(deserializeArg(msg.value) as T); if (waiting) { waiting(); waiting = null; } });
return { [Symbol.asyncIterator]() { return { async next(): Promise<IteratorResult<T>> { while (true) { if (queue.length > 0) { const value = queue.shift()!; // Atomics-based backpressure: decrement inflight. In stream // mode the producer parks on inflight, so we notify it; in // channel mode the producer parks on readySignal, and we // only wake it on the cap→below-cap transition. const prev = inflight.sub(1); if (readySignal) { if (prev === highWater) { readySignal.add(1); readySignal.notify(); } } else { inflight.notify(); } return { done: false, value }; } if (error) throw error; if (done) return { done: true, value: undefined }; await new Promise<void>((resolve) => { waiting = resolve; }); } }, async return(): Promise<IteratorResult<T>> { state.store(CANCEL); inflight.notify(); if (readySignal) { readySignal.add(1); readySignal.notify(); } port.close(); return { done: true, value: undefined }; }, }; }, };}
async function resolveArg(arg: unknown): Promise<unknown> { if (isSerializedStreamHandle(arg)) { return portToAsyncIterable(deserializeStreamHandle(arg)); } if (isTaskArg(arg)) { if (taskCache.has(arg.__task__)) { return taskCache.get(arg.__task__); } const inputArgs = arg.args; const resolvedArgs = new Array(inputArgs.length); for (let i = 0; i < inputArgs.length; i++) { const a = inputArgs[i]; resolvedArgs[i] = needsAsyncResolve(a) ? await resolveArg(a) : deserializeArg(a); } const fnM = resolveFn(arg.id); const fn = fnM instanceof Promise ? await fnM : fnM; const raw = fn(...resolvedArgs); const value = raw instanceof Promise ? await raw : raw; taskCache.set(arg.__task__, value); return value; } return deserializeArg(arg);}
type TaskMsg = { callId?: number; id: string; args: unknown[]; /** Present on streaming dispatches; bundles port + flags + highWater. */ stream?: SerializedStreamHandle;};
function postResult(callId: number, value: unknown): void { const returnValue = serializeArg(value); const transferList: Transferable[] = []; collectTransferables(value, transferList); parentPort!.postMessage({ callId, value: returnValue }, transferList);}
function postError(callId: number, err: unknown): void { parentPort!.postMessage({ callId, error: err instanceof Error ? err : new Error(String(err)), });}
// Value-task path, structured to pay zero microtask hops when the whole// chain (fn resolution, arg resolution, function invocation) is sync.function handleValueTask(msg: TaskMsg): void { const callId = msg.callId!; try { const fnM = resolveFn(msg.id); if (fnM instanceof Promise) { fnM.then( (fn) => invokeWithArgs(callId, fn, msg.args), (err) => postError(callId, err), ); return; } invokeWithArgs(callId, fnM, msg.args); } catch (err) { postError(callId, err); }}
function invokeWithArgs(callId: number, fn: Fn, args: unknown[]): void { try { // Sync prologue: walk args in-place until we hit one that needs async // resolution. If we finish the loop synchronously, go straight to invoke. const resolved = new Array(args.length); let i = 0; for (; i < args.length; i++) { const arg = args[i]; if (needsAsyncResolve(arg)) break; resolved[i] = deserializeArg(arg); } if (i === args.length) { invokeAndRespond(callId, fn, resolved); return; } // Async tail: same loop body, awaiting where needed. (async () => { for (; i < args.length; i++) { const arg = args[i]; resolved[i] = needsAsyncResolve(arg) ? await resolveArg(arg) : deserializeArg(arg); } invokeAndRespond(callId, fn, resolved); })().catch((err) => postError(callId, err)); } catch (err) { postError(callId, err); }}
function invokeAndRespond(callId: number, fn: Fn, args: unknown[]): void { try { const raw = fn(...args); if (raw instanceof Promise) { raw.then( (v) => postResult(callId, v), (err) => postError(callId, err), ); return; } postResult(callId, raw); } catch (err) { postError(callId, err); }}
async function handleStreamTask(msg: TaskMsg): Promise<void> { const { id, args } = msg; const callId = msg.callId; try { const fnM = resolveFn(id); const fn = fnM instanceof Promise ? await fnM : fnM; const resolvedArgs = new Array(args.length); for (let i = 0; i < args.length; i++) { const a = args[i]; resolvedArgs[i] = needsAsyncResolve(a) ? await resolveArg(a) : deserializeArg(a); } const gen = fn(...resolvedArgs) as AsyncGenerator; await pipeIterable(gen, deserializeStreamHandle(msg.stream!)); } catch (err) { if (callId != null) postError(callId, err); }}
parentPort!.on('message', (msg: TaskMsg) => { if (msg.stream) { void handleStreamTask(msg); } else { handleValueTask(msg); }});