diff --git a/src/execute.ts b/src/execute.ts index 197b741..150c4d7 100644 --- a/src/execute.ts +++ b/src/execute.ts @@ -3,6 +3,7 @@ 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 { getDefineId } from './define.ts'; import { serializeArg, deserializeArg } from './shared/reconstruct.ts'; import { extractTransferables, collectTransferables } from './transfer.ts'; import { PromiseLikeTask } from './task.ts'; @@ -50,6 +51,14 @@ function pipeArgToWorker(iterable: AsyncIterable): SerializedStreamHand } function prepareArg(arg: unknown): unknown { + // Defined values cross by reference; the worker resolves via module import. + // This check must win over structural detection (e.g. async-generator + // sniffing) since a defined value is user-branded, not structurally typed. + const defineId = getDefineId(arg); + if (defineId !== undefined) { + freezeModule(defineId.slice(0, defineId.lastIndexOf('#'))); + return { __define__: defineId }; + } // Auto-detect AbortSignal args — mark transferable and include in transfer list if (arg instanceof AbortSignal) { const signal = transferableAbortSignal(arg); diff --git a/src/worker-entry.ts b/src/worker-entry.ts index df1ea95..0627b8d 100644 --- a/src/worker-entry.ts +++ b/src/worker-entry.ts @@ -1,6 +1,7 @@ import { parentPort, MessagePort } from 'node:worker_threads'; import type { Transferable } from 'node:worker_threads'; import { registry } from './registry.ts'; +import { defineRegistry } from './define.ts'; import { deserializeArg, serializeArg } from './shared/reconstruct.ts'; import { collectTransferables } from './transfer.ts'; import { pipeIterable, CANCEL, DEFAULT_HIGH_WATER, deserializeStreamHandle, isSerializedStreamHandle } from './pipe.ts'; @@ -16,10 +17,20 @@ function isTaskArg(arg: unknown): arg is { __task__: number; id: string; args: u return typeof arg === 'object' && arg !== null && '__task__' in arg; } +function isDefinedRef(arg: unknown): arg is { __define__: string } { + return typeof arg === 'object' && arg !== null && '__define__' in arg; +} + function needsAsyncResolve(arg: unknown): boolean { + if (isDefinedRef(arg)) return !defineRegistry.has(arg.__define__); return arg instanceof MessagePort || isSerializedStreamHandle(arg) || isTaskArg(arg); } +function deserializeLocal(arg: unknown): unknown { + if (isDefinedRef(arg)) return defineRegistry.get(arg.__define__); + return deserializeArg(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 { @@ -120,6 +131,15 @@ function portToAsyncIterable(handle: StreamHandle): AsyncIterable { } async function resolveArg(arg: unknown): Promise { + if (isDefinedRef(arg)) { + const id = arg.__define__; + if (!defineRegistry.has(id)) { + const url = id.slice(0, id.lastIndexOf('#')); + await import(url); + if (!defineRegistry.has(id)) throw new Error(`Defined value not found: ${id}`); + } + return defineRegistry.get(id); + } if (isSerializedStreamHandle(arg)) { return portToAsyncIterable(deserializeStreamHandle(arg)); } @@ -131,7 +151,7 @@ async function resolveArg(arg: unknown): Promise { 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); + resolvedArgs[i] = needsAsyncResolve(a) ? await resolveArg(a) : deserializeLocal(a); } const fnM = resolveFn(arg.id); const fn = fnM instanceof Promise ? await fnM : fnM; @@ -193,7 +213,7 @@ function invokeWithArgs(callId: number, fn: Fn, args: unknown[]): void { for (; i < args.length; i++) { const arg = args[i]; if (needsAsyncResolve(arg)) break; - resolved[i] = deserializeArg(arg); + resolved[i] = deserializeLocal(arg); } if (i === args.length) { invokeAndRespond(callId, fn, resolved); @@ -203,7 +223,7 @@ function invokeWithArgs(callId: number, fn: Fn, args: unknown[]): void { (async () => { for (; i < args.length; i++) { const arg = args[i]; - resolved[i] = needsAsyncResolve(arg) ? await resolveArg(arg) : deserializeArg(arg); + resolved[i] = needsAsyncResolve(arg) ? await resolveArg(arg) : deserializeLocal(arg); } invokeAndRespond(callId, fn, resolved); })().catch((err) => postError(callId, err)); @@ -237,7 +257,7 @@ async function handleStreamTask(msg: TaskMsg): Promise { 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); + resolvedArgs[i] = needsAsyncResolve(a) ? await resolveArg(a) : deserializeLocal(a); } const gen = fn(...resolvedArgs) as AsyncGenerator; await pipeIterable(gen, deserializeStreamHandle(msg.stream!)); diff --git a/test/define.test.ts b/test/define.test.ts index 925b72a..e0ce8bf 100644 --- a/test/define.test.ts +++ b/test/define.test.ts @@ -1,8 +1,8 @@ import { describe, it } from 'node:test'; import assert from 'node:assert/strict'; -import { define } from 'moroutine'; +import { define, workers } from 'moroutine'; import { getDefineId } from '../src/define.ts'; -import { limits } from './fixtures/define.ts'; +import { limits, helpers, readLimits, useHelpers, lateDefine } from './fixtures/define.ts'; describe('define()', () => { it('returns the same object reference', () => { @@ -61,3 +61,38 @@ describe('getDefineId()', () => { assert.equal(getDefineId(derived), undefined); }); }); + +describe('define() as task args', () => { + it('resolves a defined value on a pool worker', async () => { + using run = workers(1); + const result = await run(readLimits(limits)); + assert.equal(result, 67); + }); + + it('resolves by module evaluation, not serialization (functions cross)', async () => { + using run = workers(1); + const result = await run(useHelpers(helpers)); + assert.equal(result, 42); + }); + + it('resolves on a dedicated worker via bare await', async () => { + const result = await useHelpers(helpers); + assert.equal(result, 42); + }); + + it('resolves repeatedly on the same worker (cached path)', async () => { + using run = workers(1); + const a = await run(readLimits(limits)); + const b = await run(readLimits(limits)); + assert.equal(a, 67); + assert.equal(b, 67); + }); + + it('freezes the module once a defined value has been dispatched', async () => { + using run = workers(1); + await run(readLimits(limits)); + assert.throws(lateDefine, { + message: /Cannot call define\(\).*after a task or defined value from this module has been dispatched/, + }); + }); +});