Something went wrong. Try again.
Offload functions to worker threads with shared memory primitives for Node.js.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197import { MessageChannel } from 'node:worker_threads';import type { MessagePort, Transferable } from 'node:worker_threads';import { AsyncIterableTask } from './stream-task.ts';import { runStreamOnDedicated } from './dedicated-runner.ts';import { newPipeFlags, CANCEL, DEFAULT_HIGH_WATER } from './pipe.ts';import type { PipeFlags, StreamHandle } from './pipe.ts';import { Int32Atomic } from './shared/int32-atomic.ts';import { serializeArg } from './shared/reconstruct.ts';import { collectTransferables } from './transfer.ts';
const CHANNEL = Symbol.for('moroutine.channel');
/** Options for configuring a streaming dispatch or channel. */export interface ChannelOptions { /** * Maximum number of in-flight items on a stream before the producer parks. * For a streaming task (`run(streamTask)`), this is how many items the * worker may have emitted ahead of the parent's consumption. For a * `channel()`, it's the per-consumer cap — the distributor skips any * consumer at cap and gives items to another ready consumer. * Defaults to 16. */ highWaterMark?: number;}
interface Consumer { port: MessagePort; flags: PipeFlags;}
/** * One-producer, many-consumer source with atomics-based backpressure. * * A single `tryPull` loop iterates the shared source and dispatches each * item to the first below-cap consumer starting from `cursor` (skip-based * round-robin). When every consumer is at cap, the loop parks on a shared * `readySignal` atomic until any consumer pulls from its port (which * decrements its inflight and bumps readySignal, waking the loop). */export class Channel<T> { private readonly iter: AsyncIterator<T>; private readonly consumers: Consumer[] = []; private readonly readySignal: Int32Atomic = new Int32Atomic(); private readonly highWater: number; private cursor = 0; private pulling = false; private done = false; private error: Error | null = null; readonly [CHANNEL] = true;
constructor(source: AsyncIterable<T>, opts?: ChannelOptions) { this.highWater = opts?.highWaterMark ?? DEFAULT_HIGH_WATER; const src = source instanceof AsyncIterableTask ? (runStreamOnDedicated((source as any).id, (source as any).args) as AsyncIterable<T>) : source; this.iter = src[Symbol.asyncIterator](); }
/** Per-consumer stream handle — complete description of this consumer's * producer-facing side of the channel. */ addConsumer(): StreamHandle { const { port1, port2 } = new MessageChannel(); port1.unref(); const flags = newPipeFlags();
const handle: StreamHandle = { port: port2, flags, readySignal: this.readySignal, highWater: this.highWater, };
if (this.done) { if (this.error) { port1.postMessage({ done: true, error: this.error }); } else { port1.postMessage({ done: true }); } try { port1.close(); } catch {} return handle; }
const consumer: Consumer = { port: port1, flags }; this.consumers.push(consumer);
port1.on('close', () => { consumer.flags.fields.state.store(CANCEL); const idx = this.consumers.indexOf(consumer); if (idx >= 0) this.consumers.splice(idx, 1); this.readySignal.add(1); this.readySignal.notify(); });
// Ensure the pull loop is running; also bump readySignal in case the // existing loop is currently parked (the new consumer is ready). this.readySignal.add(1); this.readySignal.notify(); void this.tryPull();
return handle; }
// Skip-based round-robin: scan from `cursor`, pick the first non-cancelled // consumer below cap. Cursor advances past the picked consumer so ties // rotate across the below-cap set. Consumers at cap are skipped — this // lets faster consumers keep pulling when a peer is saturated, trading // strict fairness for resilience and throughput. private findReadyConsumer(): Consumer | null { const hw = this.highWater; const n = this.consumers.length; if (n === 0) return null; for (let k = 0; k < n; k++) { const i = (this.cursor + k) % n; const c = this.consumers[i]; if (c.flags.fields.state.load() === CANCEL) continue; if (c.flags.fields.inflight.load() < hw) { this.cursor = (i + 1) % n; return c; } } return null; }
private async tryPull(): Promise<void> { if (this.pulling || this.done) return; this.pulling = true;
try { while (!this.done) { let consumer = this.findReadyConsumer(); if (!consumer) { const signal = this.readySignal.load(); consumer = this.findReadyConsumer(); if (!consumer) { if (this.consumers.length === 0) return; await this.readySignal.waitAsync(signal); continue; } }
const result = await this.iter.next(); if (result.done) { this.done = true; for (const c of this.consumers) { try { c.port.postMessage({ done: true }); } catch {} try { c.port.close(); } catch {} } return; }
const serialized = serializeArg(result.value); const transferList: Transferable[] = []; collectTransferables(result.value, transferList); try { consumer.port.postMessage({ value: serialized, done: false }, transferList); consumer.flags.fields.inflight.add(1); } catch { // Consumer may have closed between scan and post; move on. } } } catch (err) { this.error = err instanceof Error ? err : new Error(String(err)); this.done = true; for (const c of this.consumers) { try { c.port.postMessage({ done: true, error: this.error }); } catch {} try { c.port.close(); } catch {} } } finally { this.pulling = false; } }}
/** * Creates a channel for streaming values to workers. Supports fan-out when * the same channel is passed to multiple tasks. * @param iterable - The async iterable or StreamTask to distribute. * @param opts - Optional. `highWaterMark` bounds per-consumer in-flight items. * @returns A Channel, typed as AsyncIterable<T> for transparent moroutine arg use. */export function channel<T>(iterable: AsyncIterable<T>, opts?: ChannelOptions): AsyncIterable<T> { return new Channel(iterable, opts) as unknown as AsyncIterable<T>;}
export { CHANNEL };