diff --git a/.changeset/local-chunk-sharding.md b/.changeset/local-chunk-sharding.md new file mode 100644 index 0000000000..a3d672b672 --- /dev/null +++ b/.changeset/local-chunk-sharding.md @@ -0,0 +1,5 @@ +--- +'@workflow/world-local': patch +--- + +Shard stream chunks into a directory per stream so a tail reader's poll no longer lists every chunk in the world on each tick, and reliably release its emitter listeners and poll timer when the reader is cancelled. Note: stream chunks are now stored at `streams/chunks//`; chunk files written to the old flat layout by an earlier version are not read back (an acceptable tradeoff for local dev data, and stale flat files are left in place rather than cleaned up). diff --git a/.changeset/rich-signals-inline.md b/.changeset/rich-signals-inline.md new file mode 100644 index 0000000000..c98b3552b8 --- /dev/null +++ b/.changeset/rich-signals-inline.md @@ -0,0 +1,5 @@ +--- +'@workflow/core': patch +--- + +Steps that receive an `AbortSignal` argument no longer pay a per-step queue round-trip: the real-time abort-stream reader opened for such a step is now released when the step finishes, so the step can complete inline instead of reporting pending work on every invocation. diff --git a/packages/core/src/runtime/step-executor.ts b/packages/core/src/runtime/step-executor.ts index 130a73fce9..c0a0b4b7fb 100644 --- a/packages/core/src/runtime/step-executor.ts +++ b/packages/core/src/runtime/step-executor.ts @@ -22,6 +22,7 @@ import type { CryptoKey } from '../encryption.js'; import { runtimeLogger, stepLogger } from '../logger.js'; import { getStepFunction } from '../private.js'; import { + cancelAbortReaders, dehydrateStepError, dehydrateStepReturnValue, hydrateStepArguments, @@ -573,41 +574,73 @@ export async function executeStep( : `http://localhost:${(await getPortLazy()) ?? 3000}` ); + // --- User code execution --- + // Wrap only stepFn.apply() (user step code) so cleanup below runs on + // BOTH success and failure. A user-code throw is captured here and + // re-raised after cancelAbortReaders, so it still flows to the outer + // catch (step_failed/step_retrying) — but the abort-stream reader is + // torn down first. Without this, a throwing/retrying signal-bearing + // step would leak a real-time abort reader per attempt. + let userCodeError: unknown; + let userCodeFailed = false; + const executionStartTime = Date.now(); - result = await trace('step.execute', {}, async () => { - return await contextStorage.run( - { - stepMetadata: { - stepName, - stepId, - stepStartedAt: new Date(+stepStartedAt), - attempt, - }, - workflowMetadata: { - workflowName, - workflowRunId, - workflowStartedAt: new Date(+workflowStartedAt), - url: workflowBaseUrl, - features: { encryption: !!encryptionKey }, + try { + result = await trace('step.execute', {}, async () => { + return await contextStorage.run( + { + stepMetadata: { + stepName, + stepId, + stepStartedAt: new Date(+stepStartedAt), + attempt, + }, + workflowMetadata: { + workflowName, + workflowRunId, + workflowStartedAt: new Date(+workflowStartedAt), + url: workflowBaseUrl, + features: { encryption: !!encryptionKey }, + }, + workflowDeploymentId: params.workflowDeploymentId, + ops, + preCompletionOps, + closureVars: hydratedInput.closureVars, + encryptionKey, + // Turbo optimistic start runs this body before `run_started` is + // durable. Expose the barrier so a direct step-body world write + // (e.g. `experimental_setAttributes`) can order itself after the + // run exists. Undefined on the await path (run already durable). + runReadyBarrier: optimisticStart + ? params.runReadyBarrier + : undefined, }, - workflowDeploymentId: params.workflowDeploymentId, - ops, - preCompletionOps, - closureVars: hydratedInput.closureVars, - encryptionKey, - // Turbo optimistic start runs this body before `run_started` is - // durable. Expose the barrier so a direct step-body world write - // (e.g. `experimental_setAttributes`) can order itself after the - // run exists. Undefined on the await path (run already durable). - runReadyBarrier: optimisticStart - ? params.runReadyBarrier - : undefined, - }, - () => stepFn.apply(thisVal, args) - ); - }); + () => stepFn.apply(thisVal, args) + ); + }); + } catch (err) { + userCodeError = err; + userCodeFailed = true; + } const executionTimeMs = Date.now() - executionStartTime; + // Tear down any abort-stream readers opened while hydrating the step's + // arguments (a serialized AbortSignal opens a real-time abort reader for + // the step's duration). Without this the reader's `read()` promise never + // settles, so the `ops` flush below always loses the 500ms race and the + // step reports `hasPendingOps` — forcing the inline loop to queue a + // continuation and paying a full round-trip per signal-bearing step. + // The non-inline `step-handler` path already does this after user code. + // Runs unconditionally (success or failure) so a throwing step doesn't + // leak the reader. + cancelAbortReaders(...args, thisVal, hydratedInput.closureVars); + + // Re-raise a user-code failure now that cleanup has run; the outer + // catch maps it to step_failed/step_retrying. + if (userCodeFailed) { + throw userCodeError; + } + span?.setAttributes({ ...Attribute.QueueExecutionTimeMs(executionTimeMs), }); diff --git a/packages/core/src/runtime/step-handler.test.ts b/packages/core/src/runtime/step-handler.test.ts index 5f126f0280..8e5f95d61c 100644 --- a/packages/core/src/runtime/step-handler.test.ts +++ b/packages/core/src/runtime/step-handler.test.ts @@ -184,7 +184,7 @@ vi.mock('@workflow/utils/get-port', () => ({ })); import { getStepFunction } from '../private.js'; -import { dehydrateStepError } from '../serialization.js'; +import { cancelAbortReaders, dehydrateStepError } from '../serialization.js'; import { getErrorName, getErrorStack, @@ -1192,6 +1192,73 @@ describe('executeStep inline-delta threading', () => { }); }); +describe('executeStep abort-reader cleanup', () => { + const baseParams = { + workflowRunId: 'wrun_test123', + workflowName: 'test-workflow', + workflowStartedAt: Date.now(), + stepId: 'step_abc', + stepName: 'myStep', + }; + + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(getStepFunction).mockReturnValue(mockStepFn); + vi.mocked(normalizeUnknownError).mockImplementation( + async (err: unknown) => ({ + message: err instanceof Error ? err.message : String(err), + name: err instanceof Error ? err.name : 'Error', + stack: err instanceof Error ? err.stack : undefined, + }) + ); + mockStepFn.mockReset().mockResolvedValue('step-result'); + mockStepFn.maxRetries = 3; + mockEventsCreate.mockReset().mockImplementation((_runId, event) => { + if (event.eventType === 'step_started') { + return Promise.resolve({ + step: { + stepId: 'step_abc', + status: 'running', + attempt: 1, + startedAt: new Date(), + input: [], + }, + event: {}, + }); + } + return Promise.resolve({ event: {} }); + }); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + // Regression: a serialized AbortSignal opens a real-time abort-stream reader + // for the step's duration; cleanup must run whether the step succeeds OR + // throws, otherwise a throwing/retrying signal-bearing step leaks that reader + // (a filesystem poll + emitter listeners on world-local) on every attempt. + it('tears down abort readers even when the step function throws', async () => { + mockStepFn.mockReset().mockRejectedValue(new Error('boom')); + + const world = await getWorld(); + const result = await executeStep({ world: world as never, ...baseParams }); + + // The user-code error is still surfaced (retry, since attempt < maxRetries)… + expect(result.type).toBe('retry'); + // …and cleanup ran on the failure path before the error was re-raised. + expect(cancelAbortReaders).toHaveBeenCalledTimes(1); + }); + + it('tears down abort readers on the success path', async () => { + const world = await getWorld(); + const result = await executeStep({ world: world as never, ...baseParams }); + + expect(result.type).toBe('completed'); + expect(cancelAbortReaders).toHaveBeenCalledTimes(1); + }); +}); + describe('executeStep optimistic inline start', () => { const baseParams = { workflowRunId: 'wrun_test123', diff --git a/packages/core/src/serialization.ts b/packages/core/src/serialization.ts index 113c514a93..5cc3c443e0 100644 --- a/packages/core/src/serialization.ts +++ b/packages/core/src/serialization.ts @@ -1712,8 +1712,17 @@ function setupAbortStreamReader( ); }), ]); - reader.releaseLock(); if (result.value && !result.done) { + // An abort packet arrived: propagate it as fast as possible. Release + // the lock (synchronous) rather than cancelling here — on a + // service-backed World `reader.cancel()` can do a network round-trip, + // and awaiting it before `controller.abort()` would delay (or, if it + // hangs, drop) real-time abort delivery to the in-flight step. + try { + reader.releaseLock(); + } catch { + // Reader may already be released; ignore. + } try { // Hydrate via the same machinery the writer used so the reason // round-trips with full type fidelity. Encryption key (if any) @@ -1730,6 +1739,16 @@ function setupAbortStreamReader( } catch { controller.abort(); } + } else { + // The step finished (or the reader was cancelled) without an abort. + // Cancel — not just release — so the underlying World stream is torn + // down: a polling World (e.g. world-local) otherwise leaks a tail + // reader (a 100ms filesystem poll plus emitter listeners) per step + // invocation for the life of the process, since a signal-bearing step + // opens one of these on every revival and usually never aborts. Fire + // and forget: a service-backed World's cancel may hit the network, + // and this path must not block the step's ops-settle window. + void reader.cancel().catch(() => {}); } } catch { // Stream read failed — signal won't propagate in real-time, diff --git a/packages/world-local/src/index.ts b/packages/world-local/src/index.ts index a4012ac95c..a963fcb496 100644 --- a/packages/world-local/src/index.ts +++ b/packages/world-local/src/index.ts @@ -165,17 +165,36 @@ export function createWorld(args?: Partial): LocalWorld { await fs .rm(path.join(basedir, '.locks'), { recursive: true, force: true }) .catch(() => {}); - // Delete tagged stream chunks (.{tag}.bin files) + // Delete tagged stream chunks (.{tag}.bin files). Chunks are sharded + // one directory per stream (streams/chunks//.{tag}.bin), + // so iterate each per-stream directory — the top-level chunks dir now + // holds only subdirectories, so listing it directly would match nothing + // and silently leak tagged chunk files across test sessions. const chunksDir = path.join(basedir, 'streams', 'chunks'); - const taggedBinFiles = await listTaggedFilesByExtension( - chunksDir, - tag, - '.bin' - ); + let streamDirEntries: import('node:fs').Dirent[]; + try { + streamDirEntries = await fs.readdir(chunksDir, { + withFileTypes: true, + }); + } catch { + streamDirEntries = []; + } await Promise.all( - taggedBinFiles.map((f) => - fs.unlink(path.join(chunksDir, f)).catch(() => {}) - ) + streamDirEntries + .filter((entry) => entry.isDirectory()) + .map(async (entry) => { + const streamChunkDir = path.join(chunksDir, entry.name); + const taggedBinFiles = await listTaggedFilesByExtension( + streamChunkDir, + tag, + '.bin' + ); + await Promise.all( + taggedBinFiles.map((f) => + fs.unlink(path.join(streamChunkDir, f)).catch(() => {}) + ) + ); + }) ); // Clear the in-memory write cache so deleted paths are forgotten clearCreatedFilesCache(); diff --git a/packages/world-local/src/storage.test.ts b/packages/world-local/src/storage.test.ts index 4b31d4a2aa..66cfb6fd47 100644 --- a/packages/world-local/src/storage.test.ts +++ b/packages/world-local/src/storage.test.ts @@ -2118,18 +2118,27 @@ describe('Storage', () => { expect(result.hook).toBeUndefined(); }); - it('should return hook_conflict event when the token claim cannot provide a run ID', async () => { - const token = 'legacy-duplicate-test-token'; + it('reaps an unparseable claim of a live hook and recovers the real conflict from the event log', async () => { + // A corrupt claim file whose hook is still live must not be trusted as + // debris and stolen — but it also must not block the token forever. + // The claimant loop reaps the unparseable file after a few observations + // (nothing else deletes it — the releaser can't determine ownership), + // then rebuilds the live hook's claim from the event log, so the + // duplicate still conflicts and now carries the real conflicting run. + const token = 'corrupt-claim-live-hook-token'; await createHook(storage, testRunId, { hookId: 'hook_1', token, }); - await fs.writeFile( - path.join(testDir, 'hooks', 'tokens', `${hashToken(token)}.json`), - '{' + const claimPath = path.join( + testDir, + 'hooks', + 'tokens', + `${hashToken(token)}.json` ); + await fs.writeFile(claimPath, '{'); const result = await storage.events.create(testRunId, { eventType: 'hook_created', @@ -2139,10 +2148,46 @@ describe('Storage', () => { expect(result.event.eventType).toBe('hook_conflict'); expect((result.event as any).eventData.token).toBe(token); - expect( - (result.event as any).eventData.conflictingRunId - ).toBeUndefined(); + // Recovered from the event log rather than the unreadable claim. + expect((result.event as any).eventData.conflictingRunId).toBe( + testRunId + ); expect(result.hook).toBeUndefined(); + + // The corrupt file was replaced by the live hook's real, parseable claim. + const restored = JSON.parse(await fs.readFile(claimPath, 'utf8')); + expect(restored).toMatchObject({ + token, + hookId: 'hook_1', + runId: testRunId, + }); + }); + + it('reaps an unparseable orphan claim so a new hook can reuse the token (#2808)', async () => { + // Regression for the "unparseable claim blocks its token indefinitely" + // gap: a corrupt claim with no live hook behind it (genuine debris — + // e.g. a partial/corrupted write) is never reaped by the releaser + // (ownership is undeterminable). The claimant loop must delete it after + // N observations so the token becomes claimable again. + const token = 'corrupt-orphan-claim-token'; + const claimPath = path.join( + testDir, + 'hooks', + 'tokens', + `${hashToken(token)}.json` + ); + await fs.mkdir(path.dirname(claimPath), { recursive: true }); + await fs.writeFile(claimPath, 'not-json{'); + + const result = await storage.events.create(testRunId, { + eventType: 'hook_created', + correlationId: 'hook_new', + eventData: { token }, + }); + + expect(result.event.eventType).toBe('hook_created'); + expect(result.hook?.token).toBe(token); + expect(result.hook?.hookId).toBe('hook_new'); }); it('should allow multiple hooks with different tokens for the same run', async () => { @@ -2985,50 +3030,12 @@ describe('Storage', () => { ).toHaveLength(0); }); - it('should never journal hook_received after hook_disposed when resume races disposal', async () => { - // Each round races one resume against the hook's disposal. Both - // outcomes are valid — the resume lands before the disposal, or it - // is rejected — but the journal must never contain hook_received - // ordered after hook_disposed for the same hook. - for (let round = 0; round < 20; round++) { - const hookId = `hook_race_2781_${round}`; - await createHook(storage, testRunId, { - hookId, - token: `race-2781-token-${round}`, - }); - - const [resume] = await Promise.allSettled([ - storage.events.create(testRunId, { - eventType: 'hook_received', - correlationId: hookId, - eventData: { payload: new Uint8Array([round]) }, - }), - disposeHook(storage, testRunId, hookId), - ]); - - if (resume.status === 'rejected') { - expect((resume.reason as { name?: string }).name).toBe( - 'HookNotFoundError' - ); - } - - const events = await storage.events.listByCorrelationId({ - correlationId: hookId, - pagination: {}, - }); - const types = events.data.map((e) => e.eventType); - const disposedIndex = types.indexOf('hook_disposed'); - expect(disposedIndex).toBeGreaterThan(-1); - expect(types.lastIndexOf('hook_received')).toBeLessThan( - disposedIndex - ); - // A fulfilled resume must actually be in the log (before the - // disposal), a rejected one must not be journaled at all. - expect(types.filter((t) => t === 'hook_received')).toHaveLength( - resume.status === 'fulfilled' ? 1 : 0 - ); - } - }); + // NOTE: a 20-round "race resume against disposal" soak test previously + // lived here. It was removed as a regression guard: because both + // orderings (resume-before-disposal and rejected-resume) are valid, the + // loop passed even with the ordering fix reverted, so it was a soak, not + // a deterministic guard. The mid-teardown test above (which forces the + // committed-dispose-lock state) is the real regression guard for #2781. }); }); diff --git a/packages/world-local/src/storage/events-storage.ts b/packages/world-local/src/storage/events-storage.ts index 19a058d689..0f38224a0b 100644 --- a/packages/world-local/src/storage/events-storage.ts +++ b/packages/world-local/src/storage/events-storage.ts @@ -56,9 +56,9 @@ import { import { stripEventDataRefs } from './filters.js'; import { getObjectCreatedAt, - hashToken, hookDisposeLockPath, hookRecoveryMarkerPath, + hookTokenClaimPath, isHookDisposalCommitted, monotonicUlid, releaseHookTokenClaimIfOwnedBy, @@ -1629,12 +1629,7 @@ export function createEventsStorage( // Atomically claim the token using an exclusive-create constraint file. // This avoids the TOCTOU race of the previous read-all-then-check approach. - const constraintPath = path.join( - basedir, - 'hooks', - 'tokens', - `${hashToken(hookData.token)}.json` - ); + const constraintPath = hookTokenClaimPath(basedir, hookData.token); // Persist `eventId` in the claim so concurrent / cross- // process retries can converge on a single canonical // `hook_created` event path. See the recovery comment @@ -1668,6 +1663,13 @@ export function createEventsStorage( let tokenClaimed = false; let existingClaim: z.infer | null = null; let releasableObservations = 0; + // A claim file that exists (exclusive-create keeps failing) but never + // parses is debris: `writeExclusive` writes atomically, so a live + // claim is always valid JSON — an unparseable one at the canonical + // path is a corrupt/partial leftover. The releaser leaves it alone + // (ownership is undeterminable), so nothing else reaps it and it would + // block the token forever. Force-delete it after a few observations. + let unparseableObservations = 0; for (let attempt = 0; attempt < 10; attempt++) { // When the claim is absent, the event log is the only durable // source that can distinguish a first hook from a crash-lost @@ -1684,6 +1686,17 @@ export function createEventsStorage( existingClaim = await readHookTokenClaim(constraintPath); if (!existingClaim) { + // The claim either vanished (raced a releaser between the + // exclusive-create attempt and this read → retry resolves it) or + // exists but is unparseable. `readHookTokenClaim` can't tell them + // apart, so count consecutive misses; once we're confident it is + // not a transient race, delete the (presumed corrupt) file so the + // token isn't blocked forever. `deleteJSON` is a no-op if it was + // in fact a vanished-claim race. + unparseableObservations++; + if (unparseableObservations >= 3) { + await deleteJSON(constraintPath); + } continue; } if ( diff --git a/packages/world-local/src/streamer.test.ts b/packages/world-local/src/streamer.test.ts index f44929175e..7c99b433ba 100644 --- a/packages/world-local/src/streamer.test.ts +++ b/packages/world-local/src/streamer.test.ts @@ -96,9 +96,20 @@ describe('streamer', () => { await fs.rm(testDir, { recursive: true, force: true }); } else { const chunksPath = `${testDir}/streams/chunks`; + // Chunks are sharded one directory per stream: + // streams/chunks//.bin let files: string[]; try { - files = await fs.readdir(chunksPath); + const streamDirs = await fs.readdir(chunksPath); + files = ( + await Promise.all( + streamDirs.map(async (streamDir) => + ( + await fs.readdir(`${chunksPath}/${streamDir}`) + ).map((f) => `${streamDir}/${f}`) + ) + ) + ).flat(); } catch { // chunks directory may not exist if the test failed before any writes files = []; @@ -109,9 +120,10 @@ describe('streamer', () => { const chunk = deserializeChunk( await fs.readFile(`${chunksPath}/${file}`) ); - // Extract ULID from filename: "streamName-chnk_ULID.bin" - const chunkIdPart = String(file.split('-').at(-1)).split('.')[0]; // "chnk_ULID" - const ulid = chunkIdPart.replace('chnk_', ''); // Just the ULID + // Filename is ".bin"; chunkId is "chnk_ULID". + const ulid = String(file.split('/').at(-1)) + .split('.')[0] + .replace('chnk_', ''); const time = decodeTime(ulid); const timeDiff = time - lastTime; lastTime = time; @@ -144,12 +156,12 @@ describe('streamer', () => { await streamer.streams.write(TEST_RUN_ID, streamName, 'hello'); await streamer.streams.write(TEST_RUN_ID, streamName, ' world'); - // Verify chunks directory was created - const chunksDir = path.join(testDir, 'streams', 'chunks'); + // Verify the per-stream chunk directory was created + const chunksDir = path.join(testDir, 'streams', 'chunks', streamName); const files = await fs.readdir(chunksDir); expect(files).toHaveLength(2); - expect(files.every((f) => f.startsWith(`${streamName}-`))).toBe(true); + expect(files.every((f) => f.startsWith('chnk_'))).toBe(true); expect(files.every((f) => f.endsWith('.bin'))).toBe(true); }); @@ -162,11 +174,11 @@ describe('streamer', () => { await streamer.streams.write(TEST_RUN_ID, streamName, buffer1); await streamer.streams.write(TEST_RUN_ID, streamName, buffer2); - const chunksDir = path.join(testDir, 'streams', 'chunks'); + const chunksDir = path.join(testDir, 'streams', 'chunks', streamName); const files = await fs.readdir(chunksDir); expect(files).toHaveLength(2); - expect(files.every((f) => f.startsWith(`${streamName}-`))).toBe(true); + expect(files.every((f) => f.startsWith('chnk_'))).toBe(true); }); it('should write Uint8Array chunks to a stream', async () => { @@ -176,11 +188,11 @@ describe('streamer', () => { await streamer.streams.write(TEST_RUN_ID, streamName, uint8Array); - const chunksDir = path.join(testDir, 'streams', 'chunks'); + const chunksDir = path.join(testDir, 'streams', 'chunks', streamName); const files = await fs.readdir(chunksDir); expect(files).toHaveLength(1); - expect(files[0]).toMatch(`${streamName}-`); + expect(files[0]).toMatch('chnk_'); }); it('should handle multiple streams independently', async () => { @@ -190,11 +202,10 @@ describe('streamer', () => { await streamer.streams.write(TEST_RUN_ID, 'stream2', 'data2'); await streamer.streams.write(TEST_RUN_ID, 'stream1', 'data3'); + // Each stream gets its own sharded directory. const chunksDir = path.join(testDir, 'streams', 'chunks'); - const files = await fs.readdir(chunksDir); - - const stream1Files = files.filter((f) => f.startsWith('stream1-')); - const stream2Files = files.filter((f) => f.startsWith('stream2-')); + const stream1Files = await fs.readdir(path.join(chunksDir, 'stream1')); + const stream2Files = await fs.readdir(path.join(chunksDir, 'stream2')); expect(stream1Files).toHaveLength(2); expect(stream2Files).toHaveLength(1); @@ -212,11 +223,11 @@ describe('streamer', () => { 'chunk3', ]); - const chunksDir = path.join(testDir, 'streams', 'chunks'); + const chunksDir = path.join(testDir, 'streams', 'chunks', streamName); const files = await fs.readdir(chunksDir); expect(files).toHaveLength(3); - expect(files.every((f) => f.startsWith(`${streamName}-`))).toBe(true); + expect(files.every((f) => f.startsWith('chnk_'))).toBe(true); }); it('should preserve chunk ordering', async () => { @@ -250,7 +261,7 @@ describe('streamer', () => { await streamer.streams.writeMulti!(TEST_RUN_ID, streamName, []); - const chunksDir = path.join(testDir, 'streams', 'chunks'); + const chunksDir = path.join(testDir, 'streams', 'chunks', streamName); const dirExists = await fs .access(chunksDir) .then(() => true) @@ -259,10 +270,7 @@ describe('streamer', () => { // Directory might not exist if no chunks were written if (dirExists) { const files = await fs.readdir(chunksDir); - const streamFiles = files.filter((f) => - f.startsWith(`${streamName}-`) - ); - expect(streamFiles).toHaveLength(0); + expect(files).toHaveLength(0); } }); @@ -301,11 +309,11 @@ describe('streamer', () => { await streamer.streams.close(TEST_RUN_ID, streamName); - const chunksDir = path.join(testDir, 'streams', 'chunks'); + const chunksDir = path.join(testDir, 'streams', 'chunks', streamName); const files = await fs.readdir(chunksDir); expect(files).toHaveLength(1); - expect(files[0]).toMatch(`${streamName}-`); + expect(files[0]).toMatch('chnk_'); }); it('should close a stream with existing chunks', async () => { @@ -316,7 +324,7 @@ describe('streamer', () => { await streamer.streams.write(TEST_RUN_ID, streamName, 'chunk2'); await streamer.streams.close(TEST_RUN_ID, streamName); - const chunksDir = path.join(testDir, 'streams', 'chunks'); + const chunksDir = path.join(testDir, 'streams', 'chunks', streamName); const files = await fs.readdir(chunksDir); expect(files).toHaveLength(3); // 2 data chunks + 1 EOF chunk @@ -550,7 +558,9 @@ describe('streamer', () => { const streamer = createStreamer(testDir); const streamName = 'poll-test'; - const chunksDir = path.join(testDir, 'streams', 'chunks'); + // Simulate a cross-process writer landing chunks straight on disk in + // the stream's sharded directory. + const chunksDir = path.join(testDir, 'streams', 'chunks', streamName); await fs.mkdir(chunksDir, { recursive: true }); // Start reading — sets up EventEmitter listeners + polling interval @@ -578,10 +588,7 @@ describe('streamer', () => { chunk: Buffer.from('hello'), }); await fs.writeFile( - path.join( - chunksDir, - `${streamName}-chnk_01ARZ3NDEKTSV4RRFFQ69G5FAV.bin` - ), + path.join(chunksDir, `chnk_01ARZ3NDEKTSV4RRFFQ69G5FAV.bin`), chunk1 ); @@ -592,10 +599,7 @@ describe('streamer', () => { chunk: Buffer.from(' world'), }); await fs.writeFile( - path.join( - chunksDir, - `${streamName}-chnk_01ARZ3NDEKTSV4RRFFQ69G5FAW.bin` - ), + path.join(chunksDir, `chnk_01ARZ3NDEKTSV4RRFFQ69G5FAW.bin`), chunk2 ); @@ -607,10 +611,7 @@ describe('streamer', () => { chunk: Buffer.from([]), }); await fs.writeFile( - path.join( - chunksDir, - `${streamName}-chnk_01ARZ3NDEKTSV4RRFFQ69G5FAX.bin` - ), + path.join(chunksDir, `chnk_01ARZ3NDEKTSV4RRFFQ69G5FAX.bin`), eofChunk ); @@ -620,6 +621,75 @@ describe('streamer', () => { }, 10000); }); + describe('reader lifecycle (vercel/workflow#2795, #2797)', () => { + it('tears down emitter listeners and the poll interval on cancel', async () => { + const { streamer } = await setupStreamer(); + const streamName = 'teardown-stream'; + + // Open a reader on a stream with no chunks (the abort-stream shape): + // start() reads an empty directory, then arms the 100ms poll and + // registers chunk/close emitter listeners. + const before = process + .getActiveResourcesInfo() + .filter((r) => r === 'Timeout').length; + + const readable = await streamer.streams.get(TEST_RUN_ID, streamName); + const reader = readable.getReader(); + // Kick off a read so start() runs to completion (arms the poll). + const pending = reader.read(); + await new Promise((resolve) => setTimeout(resolve, 20)); + + // Cancelling must release everything the reader holds open. + await reader.cancel(); + await pending.catch(() => {}); + + const after = process + .getActiveResourcesInfo() + .filter((r) => r === 'Timeout').length; + expect(after).toBeLessThanOrEqual(before); + + // A subsequent write must not reach the cancelled reader's listeners + // (a leaked listener would still try to enqueue on a closed stream). + await expect( + streamer.streams.write(TEST_RUN_ID, streamName, 'after-cancel') + ).resolves.toBeUndefined(); + }); + + it('scopes chunk listing to the stream, not the whole world', async () => { + const { testDir, streamer } = await setupStreamer(); + + // One real chunk on the stream under test. + await streamer.streams.write(TEST_RUN_ID, 'target', 'hi'); + + // Thousands of unrelated chunks in *other* streams. Under the old flat + // layout these all lived in one directory and every tail-reader poll + // re-listed them (O(world chunks)); now each stream is sharded so the + // target's listing is unaffected. + const chunksBase = path.join(testDir, 'streams', 'chunks'); + const otherDir = path.join(chunksBase, 'noise'); + await fs.mkdir(otherDir, { recursive: true }); + await Promise.all( + Array.from({ length: 2000 }, (_, i) => + fs.writeFile(path.join(otherDir, `chnk_seed${i}.bin`), '') + ) + ); + + // The target stream's own directory holds exactly its one chunk. + const targetEntries = await fs.readdir(path.join(chunksBase, 'target')); + expect(targetEntries).toHaveLength(1); + + // And reads return only the target's data, never the noise. + const info = await streamer.streams.getInfo(TEST_RUN_ID, 'target'); + expect(info.tailIndex).toBe(0); + const { data } = await streamer.streams.getChunks( + TEST_RUN_ID, + 'target' + ); + expect(data).toHaveLength(1); + expect(Buffer.from(data[0].data).toString()).toBe('hi'); + }); + }); + describe('integration scenarios', () => { it('should handle complete write-close-read cycle', async () => { const { streamer } = await setupStreamer(); @@ -913,7 +983,9 @@ describe('streamer', () => { const { testDir } = await setupStreamer(); const streamName = 'mixed-format-stream'; const taggedStreamer = createStreamer(testDir, 'vitest-0'); - const chunksDir = path.join(testDir, 'streams', 'chunks'); + // Chunks live in the stream's sharded directory; the filename is just + // the chunk id plus its format/tag suffix (no stream-name prefix). + const chunksDir = path.join(testDir, 'streams', 'chunks', streamName); await fs.mkdir(chunksDir, { recursive: true }); const writeChunk = (fileName: string, text: string, eof = false) => @@ -923,13 +995,13 @@ describe('streamer', () => { ); await Promise.all([ - writeChunk(`${streamName}-chnk_01.json`, 'legacy-shadowed'), - writeChunk(`${streamName}-chnk_01.bin`, 'untagged-shadowed'), - writeChunk(`${streamName}-chnk_01.vitest-0.bin`, 'tagged'), - writeChunk(`${streamName}-chnk_02.json`, 'legacy-shadowed'), - writeChunk(`${streamName}-chnk_02.bin`, 'untagged'), - writeChunk(`${streamName}-chnk_03.json`, 'legacy'), - writeChunk(`${streamName}-chnk_04.vitest-0.bin`, '', true), + writeChunk(`chnk_01.json`, 'legacy-shadowed'), + writeChunk(`chnk_01.bin`, 'untagged-shadowed'), + writeChunk(`chnk_01.vitest-0.bin`, 'tagged'), + writeChunk(`chnk_02.json`, 'legacy-shadowed'), + writeChunk(`chnk_02.bin`, 'untagged'), + writeChunk(`chnk_03.json`, 'legacy'), + writeChunk(`chnk_04.vitest-0.bin`, '', true), ]); const result = await taggedStreamer.streams.getChunks( diff --git a/packages/world-local/src/streamer.ts b/packages/world-local/src/streamer.ts index 02f4e28a92..ccc99e1b3a 100644 --- a/packages/world-local/src/streamer.ts +++ b/packages/world-local/src/streamer.ts @@ -86,19 +86,35 @@ function addChunkFilesByExtension( } } +/** + * Resolve the per-stream chunk directory. Chunks are sharded one directory + * per stream (`streams/chunks//`) so that listing a stream's + * chunks costs O(chunks in that stream) rather than O(chunks in the whole + * world). A tail reader polling for new chunks would otherwise `readdir` the + * entire global chunks directory every 100ms — see vercel/workflow#2797. + */ +function chunkDirForStream(chunksBaseDir: string, name: string): string { + // Name becomes a path segment below; validate it can't escape chunksBaseDir. + assertSafeEntityId('streamName', name); + return path.join(chunksBaseDir, name); +} + /** * List chunk files for a stream, sorted chronologically (ULID order). - * Returns both the sorted file names and a map of file → extension for - * resolving the full path. Handles tagged and legacy (.json) formats. + * Returns the sorted chunk keys (each key is the chunk ULID), a map of + * key → extension for resolving the full path, and the per-stream directory + * the files live in. Handles tagged and legacy (.json) formats. + * + * Files are stored per-stream (`/.bin`), so the + * key returned here is already the chunk id — no stream-name prefix to strip. */ async function listChunkFilesForStream( - chunksDir: string, + chunksBaseDir: string, name: string, tag?: string -): Promise<{ files: string[]; extMap: Map }> { - // Name is used as a filename prefix below; validate it can't escape chunksDir. - assertSafeEntityId('streamName', name); - const entries = await listChunkEntries(chunksDir); +): Promise<{ files: string[]; extMap: Map; dir: string }> { + const dir = chunkDirForStream(chunksBaseDir, name); + const entries = await listChunkEntries(dir); const extMap = new Map(); addChunkFilesByExtension(extMap, entries, '.json'); addChunkFilesByExtension( @@ -114,11 +130,9 @@ async function listChunkFilesForStream( addChunkFilesByExtension(extMap, entries, taggedExtension); } - const files = [...extMap.keys()] - .filter((file) => file.startsWith(`${name}-`)) - .sort(); + const files = [...extMap.keys()].sort(); - return { files, extMap }; + return { files, extMap, dir }; } export function createStreamer(basedir: string, tag?: string): Streamer { @@ -212,10 +226,8 @@ export function createStreamer(basedir: string, tag?: string): Streamer { }); const chunkPath = path.join( - basedir, - 'streams', - 'chunks', - `${name}-${chunkId}${tagSuffix}.bin` + chunkDirForStream(path.join(basedir, 'streams', 'chunks'), name), + `${chunkId}${tagSuffix}.bin` ); await write(chunkPath, serialized); @@ -260,10 +272,8 @@ export function createStreamer(basedir: string, tag?: string): Streamer { }); const chunkPath = path.join( - basedir, - 'streams', - 'chunks', - `${name}-${chunkId}${tagSuffix}.bin` + chunkDirForStream(path.join(basedir, 'streams', 'chunks'), name), + `${chunkId}${tagSuffix}.bin` ); await write(chunkPath, serialized); @@ -298,10 +308,8 @@ export function createStreamer(basedir: string, tag?: string): Streamer { // Register this stream for the run (in case write wasn't called) await registerStreamForRun(runId, name); const chunkPath = path.join( - basedir, - 'streams', - 'chunks', - `${name}-${chunkId}${tagSuffix}.bin` + chunkDirForStream(path.join(basedir, 'streams', 'chunks'), name), + `${chunkId}${tagSuffix}.bin` ); await write( @@ -330,9 +338,12 @@ export function createStreamer(basedir: string, tag?: string): Streamer { options?: GetChunksOptions ): Promise { const limit = options?.limit ?? 100; - const chunksDir = path.join(basedir, 'streams', 'chunks'); - const { files: chunkFiles, extMap: fileExtMap } = - await listChunkFilesForStream(chunksDir, name, tag); + const chunksBaseDir = path.join(basedir, 'streams', 'chunks'); + const { + files: chunkFiles, + extMap: fileExtMap, + dir: chunksDir, + } = await listChunkFilesForStream(chunksBaseDir, name, tag); // Decode cursor let startIndex = 0; @@ -408,9 +419,12 @@ export function createStreamer(basedir: string, tag?: string): Streamer { }, async getInfo(_runId: string, name: string): Promise { - const chunksDir = path.join(basedir, 'streams', 'chunks'); - const { files: chunkFiles, extMap: fileExtMap } = - await listChunkFilesForStream(chunksDir, name, tag); + const chunksBaseDir = path.join(basedir, 'streams', 'chunks'); + const { + files: chunkFiles, + extMap: fileExtMap, + dir: chunksDir, + } = await listChunkFilesForStream(chunksBaseDir, name, tag); // Read only the EOF marker byte because metadata never needs payloads. let streamDone = false; @@ -432,9 +446,19 @@ export function createStreamer(basedir: string, tag?: string): Streamer { }, async get(_runId: string, name: string, startIndex = 0) { - const chunksDir = path.join(basedir, 'streams', 'chunks'); - let removeListeners = () => {}; + const chunksBaseDir = path.join(basedir, 'streams', 'chunks'); + // Tears down everything the reader holds open: both emitter listeners + // and the filesystem poll interval. Assigned once listeners are wired + // up in start(); called on cancel() and on terminal (EOF/close) paths. + // Kept robust (unconditional) so a cancel() while still reading from + // disk can't leak a listener/poll — a signal-bearing step opens one of + // these readers per invocation, so any leak accumulates fast. + let teardown = () => {}; let pollInterval: ReturnType | null = null; + // Set when the controller is closed/cancelled; guards against + // enqueue-after-close in the polling callback when teardown fires + // mid-iteration. + let streamClosed = false; return new ReadableStream({ async start(controller) { @@ -448,9 +472,6 @@ export function createStreamer(basedir: string, tag?: string): Streamer { let isReadingFromDisk = true; // Buffer close event if it arrives during disk reading let pendingClose = false; - // Set when the controller is closed; guards against enqueue-after-close - // in the polling callback when closeListener fires mid-iteration. - let streamClosed = false; const chunkListener = (event: { streamName: string; @@ -489,27 +510,35 @@ export function createStreamer(basedir: string, tag?: string): Streamer { } // Remove listeners before closing streamClosed = true; + teardown(); + try { + controller.close(); + } catch { + // Ignore if controller is already closed (e.g., from cancel() or EOF) + } + }; + // Tear down listeners and the poll unconditionally. Unlike + // closeListener this never defers on isReadingFromDisk, so cancel() + // reliably releases the reader even mid-disk-read. + teardown = () => { streamEmitter.off(`chunk:${name}` as const, chunkListener); streamEmitter.off(`close:${name}` as const, closeListener); if (pollInterval) { clearInterval(pollInterval); pollInterval = null; } - try { - controller.close(); - } catch { - // Ignore if controller is already closed (e.g., from cancel() or EOF) - } }; - removeListeners = closeListener; // Set up listeners FIRST to avoid missing events streamEmitter.on(`chunk:${name}` as const, chunkListener); streamEmitter.on(`close:${name}` as const, closeListener); // Now load existing chunks from disk. - const { files: chunkFiles, extMap: fileExtMap } = - await listChunkFilesForStream(chunksDir, name, tag); + const { + files: chunkFiles, + extMap: fileExtMap, + dir: chunksDir, + } = await listChunkFilesForStream(chunksBaseDir, name, tag); // Resolve negative startIndex relative to the number of data chunks // (excluding the trailing EOF marker chunk, if present). @@ -540,12 +569,9 @@ export function createStreamer(basedir: string, tag?: string): Streamer { let isComplete = false; for (let i = resolvedStartIndex; i < chunkFiles.length; i++) { const file = chunkFiles[i]; - // Extract chunk ID from filename: "streamName-chunkId" or "streamName-chunkId.tag" - const rawChunkId = file.substring(name.length + 1); - // Strip tag suffix (e.g., "chnk_ULID.vitest-0" → "chnk_ULID") - const chunkId = tag - ? rawChunkId.replace(`.${tag}`, '') - : rawChunkId; + // Files are sharded per stream, so the key is already the chunk id + // (no stream-name prefix, tag suffix already stripped). + const chunkId = file; // Skip if already delivered via event if (deliveredChunkIds.has(chunkId)) { @@ -581,7 +607,8 @@ export function createStreamer(basedir: string, tag?: string): Streamer { } if (isComplete) { - removeListeners(); + streamClosed = true; + teardown(); try { controller.close(); } catch { @@ -592,8 +619,8 @@ export function createStreamer(basedir: string, tag?: string): Streamer { // Process any pending close event that arrived during disk reading if (pendingClose) { - streamEmitter.off(`chunk:${name}` as const, chunkListener); - streamEmitter.off(`close:${name}` as const, closeListener); + streamClosed = true; + teardown(); try { controller.close(); } catch { @@ -608,12 +635,17 @@ export function createStreamer(basedir: string, tag?: string): Streamer { i < resolvedStartIndex && i < chunkFiles.length; i++ ) { - const file = chunkFiles[i]; - const rawChunkId = file.substring(name.length + 1); - const chunkId = tag - ? rawChunkId.replace(`.${tag}`, '') - : rawChunkId; - deliveredChunkIds.add(chunkId); + // Files are sharded per stream: the key is already the chunk id. + deliveredChunkIds.add(chunkFiles[i]); + } + + // If the reader was already cancelled/closed while we were reading + // from disk above (start() yields at every await), don't arm the + // poll — cancel()'s teardown ran before this point and would leave + // the freshly-created interval orphaned. + if (streamClosed) { + teardown(); + return; } // Start filesystem polling for cross-process streaming support. @@ -626,13 +658,11 @@ export function createStreamer(basedir: string, tag?: string): Streamer { isPolling = true; try { const { files: currentFiles, extMap: currentExtMap } = - await listChunkFilesForStream(chunksDir, name, tag); + await listChunkFilesForStream(chunksBaseDir, name, tag); for (const file of currentFiles) { - const rawChunkId = file.substring(name.length + 1); - const chunkId = tag - ? rawChunkId.replace(`.${tag}`, '') - : rawChunkId; + // Files are sharded per stream: the key is already the chunk id. + const chunkId = file; if (deliveredChunkIds.has(chunkId)) continue; deliveredChunkIds.add(chunkId); @@ -644,12 +674,7 @@ export function createStreamer(basedir: string, tag?: string): Streamer { if (chunk?.eof === true) { streamClosed = true; - if (pollInterval) { - clearInterval(pollInterval); - pollInterval = null; - } - streamEmitter.off(`chunk:${name}` as const, chunkListener); - streamEmitter.off(`close:${name}` as const, closeListener); + teardown(); try { controller.close(); } catch { @@ -679,11 +704,8 @@ export function createStreamer(basedir: string, tag?: string): Streamer { }, cancel() { - removeListeners(); - if (pollInterval) { - clearInterval(pollInterval); - pollInterval = null; - } + streamClosed = true; + teardown(); }, }); }, diff --git a/packages/world-local/src/tag.test.ts b/packages/world-local/src/tag.test.ts index c95f5be548..35b240c1d8 100644 --- a/packages/world-local/src/tag.test.ts +++ b/packages/world-local/src/tag.test.ts @@ -351,6 +351,44 @@ describe('File tagging', () => { await world.close?.(); }); + + it('should clear tagged stream chunks in the sharded layout', async () => { + const { createLocalWorld } = await import('./index.js'); + + const world = createLocalWorld({ dataDir: testDir, tag: 'vitest-0' }); + await world.start?.(); + + const run = await createRun(world, { + deploymentId: 'dep-1', + workflowName: 'stream-wf', + input: new Uint8Array(), + }); + await updateRun(world, run.runId, 'run_started'); + + // Chunks land under a per-stream subdirectory as `.vitest-0.bin`. + await world.streams.write(run.runId, 'strm_a', 'hello'); + await world.streams.write(run.runId, 'strm_b', 'world'); + + const chunksDir = path.join(testDir, 'streams', 'chunks'); + const streamADir = path.join(chunksDir, 'strm_a'); + const taggedBin = (files: string[]) => + files.filter((f) => f.endsWith('.vitest-0.bin')); + const readOrEmpty = (d: string) => + fs.readdir(d).catch(() => [] as string[]); + expect(taggedBin(await fs.readdir(streamADir))).toHaveLength(1); + + await world.clear(); + + // Regression: the tag-scoped clear previously listed only the top-level + // chunks dir, which after sharding holds only subdirectories, so tagged + // chunk files leaked across test sessions. + expect(taggedBin(await readOrEmpty(streamADir))).toHaveLength(0); + expect( + taggedBin(await readOrEmpty(path.join(chunksDir, 'strm_b'))) + ).toHaveLength(0); + + await world.close?.(); + }); }); describe('untagged clear()', () => { diff --git a/packages/world-testing/src/inline-execution.mts b/packages/world-testing/src/inline-execution.mts index 34deb70c50..50f72c910e 100644 --- a/packages/world-testing/src/inline-execution.mts +++ b/packages/world-testing/src/inline-execution.mts @@ -150,6 +150,45 @@ export function inlineExecution(world: string) { } ); + test( + 'sequential steps carrying an AbortSignal still run in a single flow invocation', + { timeout: 30_000 }, + async () => { + // Regression for vercel/workflow#2795: a serialized AbortSignal in step + // arguments opens a real-time abort-stream reader for the step's + // duration. If that reader isn't torn down before the step's ops-settle + // check, its never-resolving read() makes every signal-bearing step + // report `hasPendingOps`, forcing a queue continuation per step (N+1 + // invocations) instead of running inline. + const server = await startServer({ world }).then(createFetcher); + const count = 5; + const result = await server.invoke( + 'workflows/signal-cost.ts', + 'signalCostWorkflow', + [{ count, withSignal: true }] + ); + + const run = await vi.waitFor( + async () => { + const run = await server.getRun(result.runId); + expect(run.status).toBe('completed'); + return run; + }, + { interval: 200, timeout: 29_000 } + ); + + const output = await hydrateWorkflowReturnValue( + run.output!, + run.runId, + undefined + ); + expect(output).toBe('done'); + + const invocations = await server.getFlowInvocationCount(result.runId); + expect(invocations).toBe(1); + } + ); + // Hook invocation counting is tested by the existing hooks test suite. // The hook pattern requires external resume, which involves complex // timing. The invocation count for hooks is 2: create + resume. diff --git a/packages/world-testing/workflows/signal-cost.ts b/packages/world-testing/workflows/signal-cost.ts new file mode 100644 index 0000000000..a033ee09e6 --- /dev/null +++ b/packages/world-testing/workflows/signal-cost.ts @@ -0,0 +1,27 @@ +// Repro for vercel/workflow#2795 / #2797: threading an AbortSignal into a step +// adds fixed per-step latency on world-local (and degrades further as the +// world accumulates stream chunks). This workflow runs N trivial steps, with +// the signal optionally omitted from the step input. + +async function trivialStep(input: { + index: number; + abortSignal?: AbortSignal; +}): Promise { + 'use step'; + return input.index; +} + +export async function signalCostWorkflow(input: { + count: number; + withSignal: boolean; +}): Promise { + 'use workflow'; + const controller = new AbortController(); + for (let index = 0; index < input.count; index++) { + await trivialStep({ + abortSignal: input.withSignal ? controller.signal : undefined, + index, + }); + } + return 'done'; +} diff --git a/workbench/vitest/test/hook-token-reuse.test.ts b/workbench/vitest/test/hook-token-reuse.test.ts index 966f4fea9b..ffb18f1501 100644 --- a/workbench/vitest/test/hook-token-reuse.test.ts +++ b/workbench/vitest/test/hook-token-reuse.test.ts @@ -1,6 +1,5 @@ import { waitForHook } from '@workflow/vitest'; import { describe, expect, it } from 'vitest'; -import type { Run } from 'workflow/api'; import { resumeHook, start } from 'workflow/api'; import { claimTokenOnceWorkflow, @@ -27,29 +26,37 @@ describe('hook token reuse after dispose', () => { await expect(run.returnValue).resolves.toBe('ok'); }, 60_000); - // Issue #2778: when run A disposes its hook and completes, run B claiming - // the same token immediately afterwards must not conflict against run A. + // Issue #2778: once a run disposes its hook AND COMPLETES, the token claim + // is released, so the next run reusing the same token must claim it cleanly + // instead of conflicting against the finished run. Each round drives its run + // to completion before the next starts — matching the guarantee's own + // precondition ("disposes AND completes"). An earlier version resumed and + // immediately started the next claimant without awaiting completion; because + // `resumeHook` only enqueues the target run's continuation (it does not wait + // for it to process the resume + dispose), the next run could legitimately + // observe the previous run still holding the token — a correct conflict, but + // one that made the assertion ~10% flaky. it('next run can claim the token right after the previous run disposed it', async () => { const token = `handoff-reuse-${Math.random().toString(36).slice(2)}`; const rounds = 5; - const runs: Run[] = []; for (let round = 0; round < rounds; round++) { const run = await start(claimTokenOnceWorkflow, [token]); - runs.push(run); + // The token is free (the previous round's run has completed), so this + // run registers a hook and suspends rather than returning a conflict. + // A regression (spurious conflict against a released claim) would resolve + // `returnValue` first with a `conflict:*` value and fail here. const settled = await Promise.race([ waitForHook(run, { token }).then(() => 'hook' as const), run.returnValue.then((value) => ({ value })), ]); expect(settled, `round ${round} should register a hook`).toBe('hook'); - // Resume the hook and immediately start the next claimant without - // waiting for this run to settle (the "fast handoff" timing). await resumeHook(token, { n: round }); + // Wait for the run to dispose the hook and complete — releasing the + // token claim — before the next round reuses the token. + await expect(run.returnValue).resolves.toBe('ok'); } - - const results = await Promise.all(runs.map((run) => run.returnValue)); - expect(results).toEqual(Array(rounds).fill('ok')); }, 60_000); });