Something went wrong. Try again.
Offload functions to worker threads with shared memory primitives for Node.js.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197import type { Worker } from 'node:worker_threads';import { MessageChannel } from 'node:worker_threads';import type { MessagePort, Transferable } from 'node:worker_threads';import { transferableAbortSignal } from 'node:util';import { freezeModule } from './registry.ts';import { serializeArg, deserializeArg } from './shared/reconstruct.ts';import { extractTransferables, collectTransferables } from './transfer.ts';import { PromiseLikeTask } from './task.ts';import { AsyncIterableTask } from './stream-task.ts';import { runStreamOnDedicated } from './dedicated-runner.ts';import { Channel } from './channel.ts';import { pipeIterable, newPipeFlags, CANCEL, DEFAULT_HIGH_WATER, serializeStreamHandle } from './pipe.ts';import type { StreamHandle, SerializedStreamHandle } from './pipe.ts';import type { ChannelOptions } from './channel.ts';
let nextCallId = 0;const pending = new Map<number, { resolve: (value: any) => void; reject: (reason: any) => void }>();const streamPortStack: MessagePort[][] = [];
export function setupWorker(worker: Worker): void { worker.on('message', (msg: { callId: number; value?: unknown; error?: Error }) => { const call = pending.get(msg.callId); if (!call) return; pending.delete(msg.callId); if (msg.error !== undefined) { call.reject(msg.error); } else { call.resolve(deserializeArg(msg.value)); } });}
function isAsyncGenerator(arg: unknown): boolean { return typeof arg === 'object' && arg !== null && (arg as any)[Symbol.toStringTag] === 'AsyncGenerator';}
/** * Pipes an AsyncIterable from the parent thread to a worker-facing MessagePort. * Creates a producer-side pipeIterable loop with atomics backpressure and * returns a serialized consumer-side handle for the worker. */function pipeArgToWorker(iterable: AsyncIterable<unknown>): SerializedStreamHandle { const { port1, port2 } = new MessageChannel(); port1.unref(); const flags = newPipeFlags(); const producerHandle: StreamHandle = { port: port1, flags, highWater: DEFAULT_HIGH_WATER }; void pipeIterable(iterable, producerHandle, { extractTransfers: true }); streamPortStack[streamPortStack.length - 1].push(port2); return serializeStreamHandle({ port: port2, flags, highWater: DEFAULT_HIGH_WATER });}
function prepareArg(arg: unknown): unknown { // Auto-detect AbortSignal args — mark transferable and include in transfer list if (arg instanceof AbortSignal) { const signal = transferableAbortSignal(arg); streamPortStack[streamPortStack.length - 1].push(signal as unknown as MessagePort); return signal; } // Auto-detect AsyncGenerator args — pipe via MessageChannel + atomics backpressure if (isAsyncGenerator(arg)) { return pipeArgToWorker(arg as AsyncIterable<unknown>); } // Auto-detect StreamTask args — dispatch to dedicated worker, pipe output if (arg instanceof AsyncIterableTask) { return pipeArgToWorker(runStreamOnDedicated(arg.id, arg.args) as AsyncIterable<unknown>); } // Channel wrapper — single distributor loop + per-consumer atomics backpressure. if (arg instanceof Channel) { const handle = arg.addConsumer(); streamPortStack[streamPortStack.length - 1].push(handle.port); return serializeStreamHandle(handle); } if (arg instanceof PromiseLikeTask) { return { __task__: arg.uid, id: arg.id, args: arg.args.map(prepareArg) }; } return serializeArg(arg);}
export function execute<T>(worker: Worker, id: string, args: unknown[]): Promise<T> { const url = id.slice(0, id.lastIndexOf('#')); freezeModule(url); const callId = nextCallId++; return new Promise<T>((resolve, reject) => { pending.set(callId, { resolve, reject }); const extracted = extractTransferables(args); streamPortStack.push([]); const preparedArgs = extracted.args.map(prepareArg); const ports = streamPortStack.pop()!; const msg = { callId, id, args: preparedArgs }; worker.postMessage(msg, [...extracted.transfer, ...ports] as any[]); });}
export interface StreamDispatch<T> { iterable: AsyncIterable<T>; done: Promise<void>;}
export function dispatchStream<T>( worker: Worker, id: string, args: unknown[], opts?: ChannelOptions,): StreamDispatch<T> { const url = id.slice(0, id.lastIndexOf('#')); freezeModule(url); const highWater = opts?.highWaterMark ?? DEFAULT_HIGH_WATER;
const { port1, port2 } = new MessageChannel(); // Atomics-based backpressure: shared inflight count + cancel flag. // Worker parks on flags.inflight when it hits highWater; parent decrements // + notifies on pull. Replaces pause/resume messages and the // adaptive-yield setImmediate dance. const flags = newPipeFlags(); const handle: StreamHandle = { port: port2, flags, highWater };
const extracted = extractTransferables(args); streamPortStack.push([]); const preparedArgs = extracted.args.map(prepareArg); const ports = streamPortStack.pop()!; const msg = { id, args: preparedArgs, stream: serializeStreamHandle(handle) }; worker.postMessage(msg, [...extracted.transfer, ...ports, port2] as any[]);
let resolveDone: () => void; const donePromise = new Promise<void>((r) => { resolveDone = r; });
const queue: T[] = []; let done = false; let error: Error | null = null; let waiting: (() => void) | null = null;
port1.on('message', (msg: { value?: unknown; done?: boolean; error?: Error }) => { if (msg.error) { error = msg.error; done = true; port1.close(); resolveDone!(); if (waiting) { waiting(); waiting = null; } return; } if (msg.done) { done = true; port1.close(); resolveDone!(); if (waiting) { waiting(); waiting = null; } return; } queue.push(deserializeArg(msg.value) as T); if (waiting) { waiting(); waiting = null; } }); port1.unref();
return { iterable: { [Symbol.asyncIterator]() { return { async next(): Promise<IteratorResult<T>> { while (true) { if (queue.length > 0) { const value = queue.shift()!; // Signal consumption to producer flags.fields.inflight.sub(1); flags.fields.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>> { flags.fields.state.store(CANCEL); flags.fields.inflight.notify(); port1.close(); resolveDone!(); return { done: true, value: undefined }; }, }; }, }, done: donePromise, };}