From fa004f3639ffa2291474a9374e4e1e3a6535d7cb Mon Sep 17 00:00:00 2001 From: Peter Wielander Date: Fri, 17 Jul 2026 18:16:52 -0700 Subject: [PATCH 1/2] fix(core): batch stream writes via writeMulti instead of one PUT per chunk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit flushablePipe awaited each writer.write() and the WritableStream sink serializes chunks one at a time, so WorkflowServerWritableStream's buffer never held more than one chunk and its writeMulti batching path (gated on chunksToFlush.length > 1) never engaged — every chunk became its own server round trip. WorkflowServerWritableStream now exposes a durable batch write under STREAM_WRITE_BATCH_SYMBOL. flushablePipe feature-detects it and runs a producer/consumer that coalesces chunks arriving while a previous batch is still in flight into a single writeMulti, preserving the existing durability guarantees (pendingOps counts read-but-not-durable chunks; source-done awaits every batch before closing). Sinks without the batch entry point keep the unchanged per-chunk path. Co-Authored-By: Claude Opus 4.8 (1M context) --- .changeset/stream-writemulti-batching.md | 5 + packages/core/src/flushable-stream.test.ts | 196 ++++++++++++++++- packages/core/src/flushable-stream.ts | 235 ++++++++++++++++++++- packages/core/src/serialization.ts | 22 ++ packages/core/src/symbols.ts | 16 ++ packages/core/src/writable-stream.test.ts | 98 +++++++++ 6 files changed, 570 insertions(+), 2 deletions(-) create mode 100644 .changeset/stream-writemulti-batching.md diff --git a/.changeset/stream-writemulti-batching.md b/.changeset/stream-writemulti-batching.md new file mode 100644 index 0000000000..f291c7c83e --- /dev/null +++ b/.changeset/stream-writemulti-batching.md @@ -0,0 +1,5 @@ +--- +'@workflow/core': patch +--- + +Fix stream writes never batching: `flushablePipe` awaited each `writer.write()` and the WritableStream sink serialized chunks one at a time, so the server writable's buffer never held more than one chunk and its `writeMulti` batching path never engaged — every chunk became its own server round trip. The server writable now exposes a durable batch write that `flushablePipe` uses to coalesce chunks arriving while a previous write is in flight into a single `writeMulti`, while preserving the existing durability guarantees. diff --git a/packages/core/src/flushable-stream.test.ts b/packages/core/src/flushable-stream.test.ts index 423e2f08bd..1da6a51742 100644 --- a/packages/core/src/flushable-stream.test.ts +++ b/packages/core/src/flushable-stream.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from 'vitest'; +import { afterEach, describe, expect, it } from 'vitest'; import { createFlushableState, flushablePipe, @@ -6,6 +6,46 @@ import { pollReadableLock, pollWritableLock, } from './flushable-stream.js'; +import { STREAM_WRITE_BATCH_SYMBOL } from './symbols.js'; + +const tick = (ms = 0) => new Promise((r) => setTimeout(r, ms)); + +/** + * A batch-capable mock sink, mirroring the durable batch entry point that + * `WorkflowServerWritableStream` exposes under `STREAM_WRITE_BATCH_SYMBOL`. + * `gate` (when provided) lets a test hold a batch write "in flight" so it can + * enqueue more chunks and observe coalescing. + */ +function makeBatchSink( + gate?: (call: number, chunks: Uint8Array[]) => Promise | void +) { + const batches: Uint8Array[][] = []; + let closed = false; + let call = 0; + const sink = new WritableStream({ + close() { + closed = true; + }, + }) as WritableStream & { + [STREAM_WRITE_BATCH_SYMBOL]: (chunks: Uint8Array[]) => Promise; + }; + sink[STREAM_WRITE_BATCH_SYMBOL] = async (chunks: Uint8Array[]) => { + const n = call++; + batches.push(chunks.slice()); + await gate?.(n, chunks); + }; + return { sink, batches, isClosed: () => closed }; +} + +function makeControlledSource() { + let controller!: ReadableStreamDefaultController; + const source = new ReadableStream({ + start(c) { + controller = c; + }, + }); + return { source, controller: () => controller }; +} describe('flushable stream behavior', () => { it('does not emit an unhandled rejection before the runtime awaits a failed operation', async () => { @@ -400,3 +440,157 @@ describe('flushable stream behavior', () => { expect(state.streamEnded).toBe(true); }); }); + +describe('flushablePipe batching (STREAM_WRITE_BATCH_SYMBOL sinks)', () => { + afterEach(() => { + delete process.env.WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS; + }); + + it('coalesces chunks that arrive during an in-flight batch into one write', async () => { + // Hold the first batch write "in flight" so the chunks that arrive while + // it is pending accumulate and go out together as the next batch. + let releaseFirst!: () => void; + const firstInFlight = new Promise((r) => { + releaseFirst = r; + }); + const { sink, batches } = makeBatchSink((call) => + call === 0 ? firstInFlight : undefined + ); + const { source, controller } = makeControlledSource(); + const state = createFlushableState(); + + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + // First chunk: consumer picks it up alone and blocks on the gate. + controller().enqueue(new Uint8Array([1])); + await tick(); + expect(batches).toHaveLength(1); + expect(batches[0]).toHaveLength(1); + + // While the first write is in flight, three more chunks arrive. + controller().enqueue(new Uint8Array([2])); + controller().enqueue(new Uint8Array([3])); + controller().enqueue(new Uint8Array([4])); + await tick(); + // Still only the first batch has been dispatched — the rest are queued. + expect(batches).toHaveLength(1); + // They are counted as pending (read but not yet durable): 1 in-flight + 3. + expect(state.pendingOps).toBe(4); + + // Release the first write; the queued chunks flush as a single batch. + releaseFirst(); + await tick(); + expect(batches).toHaveLength(2); + expect(Array.from(batches[1].map((c) => c[0]))).toEqual([2, 3, 4]); + + controller().close(); + await pipe; + await expect(state.promise).resolves.toBeUndefined(); + expect(state.pendingOps).toBe(0); + }); + + it('delivers every chunk in order and closes the sink on completion', async () => { + const { sink, batches, isClosed } = makeBatchSink(); + const { source, controller } = makeControlledSource(); + const state = createFlushableState(); + + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + for (let i = 0; i < 25; i++) controller().enqueue(new Uint8Array([i])); + controller().close(); + await pipe; + + const delivered = batches.flat().map((c) => c[0]); + expect(delivered).toEqual(Array.from({ length: 25 }, (_, i) => i)); + expect(isClosed()).toBe(true); + await expect(state.promise).resolves.toBeUndefined(); + expect(state.pendingOps).toBe(0); + }); + + it('keeps pendingOps > 0 until batches are durable (lock-release durability)', async () => { + let releaseWrite!: () => void; + const inFlight = new Promise((r) => { + releaseWrite = r; + }); + const { sink } = makeBatchSink(() => inFlight); + const { source, controller } = makeControlledSource(); + const state = createFlushableState(); + + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + controller().enqueue(new Uint8Array([1])); + controller().close(); + await tick(); + + // Source is done, but the write has not landed yet: a lock-release poll + // must NOT resolve while the chunk is still un-durable. + pollReadableLock(source, state); + await tick(LOCK_POLL_INTERVAL_MS + 20); + expect(state.doneResolved).toBe(false); + expect(state.pendingOps).toBe(1); + + releaseWrite(); + await pipe; + expect(state.pendingOps).toBe(0); + await expect(state.promise).resolves.toBeUndefined(); + }); + + it('propagates a batch-write failure through state.promise', async () => { + const { sink } = makeBatchSink((call) => { + if (call === 0) throw new Error('batch write failed'); + }); + const { source, controller } = makeControlledSource(); + const state = createFlushableState(); + + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + controller().enqueue(new Uint8Array([1])); + await tick(); + + await pipe; + await expect(state.promise).rejects.toThrow('batch write failed'); + expect(state.streamEnded).toBe(true); + }); + + it('applies backpressure once too many chunks are outstanding', async () => { + process.env.WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS = '2'; + + let releaseFirst!: () => void; + const firstInFlight = new Promise((r) => { + releaseFirst = r; + }); + const { sink, batches } = makeBatchSink((call) => + call === 0 ? firstInFlight : undefined + ); + + // A source that records how many chunks the pipe has pulled, so we can + // assert the producer stops reading under backpressure. + let pulled = 0; + const total = 6; + const source = new ReadableStream({ + pull(controller) { + if (pulled < total) { + controller.enqueue(new Uint8Array([pulled])); + pulled++; + } else { + controller.close(); + } + }, + }); + const state = createFlushableState(); + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + // First chunk is in flight; the producer may read ahead only up to the cap + // (2 outstanding) and then must wait, so it cannot drain all 6. + await tick(20); + expect(batches).toHaveLength(1); + expect(pulled).toBeLessThan(total); + expect(state.pendingOps).toBeLessThanOrEqual(2); + + // Releasing the in-flight write lets the pipe drain the rest. + releaseFirst(); + await pipe; + expect(batches.flat()).toHaveLength(total); + expect(state.pendingOps).toBe(0); + }); +}); diff --git a/packages/core/src/flushable-stream.ts b/packages/core/src/flushable-stream.ts index dfddacf670..dac37c0b79 100644 --- a/packages/core/src/flushable-stream.ts +++ b/packages/core/src/flushable-stream.ts @@ -1,6 +1,29 @@ import { WorkflowRuntimeError } from '@workflow/errors'; import { type PromiseWithResolvers, withResolvers } from '@workflow/utils'; import { envNumber } from '@workflow/world'; +import { STREAM_WRITE_BATCH_SYMBOL } from './symbols.js'; + +/** + * A batched, durable write entry point a sink may expose under + * {@link STREAM_WRITE_BATCH_SYMBOL}. Resolves once every chunk in the batch + * has reached the server. See `WorkflowServerWritableStream`. + */ +type BatchWrite = (chunks: Uint8Array[]) => Promise; + +/** + * Upper bound on chunks read-but-not-yet-durably-written while coalescing. + * Once this many chunks are outstanding the producer stops reading until the + * consumer drains a batch, so a fast producer paired with a slow server can't + * grow the in-memory queue without bound. Override: + * `WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS`. + */ +export const MAX_INFLIGHT_CHUNKS = 1000; + +const getMaxInflightChunks = (): number => + envNumber('WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS', MAX_INFLIGHT_CHUNKS, { + integer: true, + min: 1, + }); /** * Polling interval (in ms) for lock release detection. @@ -217,7 +240,32 @@ export function pollReadableLock( * @param state - The flushable state tracker * @returns Promise that resolves when stream ends (not when done promise resolves) */ -export async function flushablePipe( +export function flushablePipe( + source: ReadableStream, + sink: WritableStream, + state: FlushableStreamState +): Promise { + // When the sink can accept a batch of chunks in one durable write (server + // writables do), coalesce chunks that arrive while a previous batch is still + // in flight into a single server write. Without this, the WHATWG + // `WritableStream` contract serializes the sink one chunk per write() and the + // per-chunk `await writer.write()` in the fallback means the sink's buffer + // never holds more than one chunk — so its `writeMulti` batching path never + // engages and every chunk becomes its own server round trip. + const batchWrite = (sink as { [STREAM_WRITE_BATCH_SYMBOL]?: BatchWrite })[ + STREAM_WRITE_BATCH_SYMBOL + ]; + return typeof batchWrite === 'function' + ? flushablePipeCoalescing(source, sink, state, batchWrite) + : flushablePipePerChunk(source, sink, state); +} + +/** + * Per-chunk variant of {@link flushablePipe}: awaits each `writer.write()` + * before reading the next chunk. Used for sinks without a durable batch write + * (plain `WritableStream`s, `TransformStream` writables on the read path). + */ +async function flushablePipePerChunk( source: ReadableStream, sink: WritableStream, state: FlushableStreamState @@ -290,3 +338,188 @@ export async function flushablePipe( writer.releaseLock(); } } + +/** + * Shared, mutable coordination state between the coalescing pipe's producer + * (reads from source into `queue`) and consumer (drains `queue` in batches). + */ +interface CoalesceContext { + state: FlushableStreamState; + batchWrite: BatchWrite; + /** Chunks read but not yet handed to a batch write. */ + queue: Uint8Array[]; + /** Set once the source has reported `done`. */ + sourceDone: boolean; + /** Resolver that wakes the consumer when the queue grows or the source ends. */ + wakeConsumer: (() => void) | null; + /** Resolver that wakes the producer when the consumer relieves backpressure. */ + wakeProducer: (() => void) | null; +} + +function wake(ctx: CoalesceContext, which: 'wakeConsumer' | 'wakeProducer') { + const resolve = ctx[which]; + if (resolve) { + ctx[which] = null; + resolve(); + } +} + +/** + * Consumer loop: drains the whole queue into one `batchWrite` at a time, + * awaiting each batch so it is durable on the server before its chunks are + * removed from `pendingOps`. Because it takes *everything* queued each round, + * chunks that arrived while the previous batch was in flight coalesce into the + * next server write. + */ +async function drainBatches(ctx: CoalesceContext): Promise { + while (true) { + if (ctx.queue.length === 0) { + if (ctx.sourceDone) return; + await new Promise((resolve) => { + ctx.wakeConsumer = resolve; + }); + continue; + } + const batch = ctx.queue; + ctx.queue = []; + try { + await ctx.batchWrite(batch); + } finally { + ctx.state.pendingOps -= batch.length; + // Relieve producer backpressure now that this batch is durable. + wake(ctx, 'wakeProducer'); + } + } +} + +/** + * Await the next chunk while a batch write is in flight, racing three outcomes: + * a read result, the sink closing under us, or a server-write failure from the + * consumer (surfaced promptly even while blocked on the read). On normal + * completion the consumer only settles after the producer sets `sourceDone`, so + * its branch throws only on real failure. + */ +function readNextChunk( + reader: ReadableStreamDefaultReader, + writer: WritableStreamDefaultWriter, + consumer: Promise +): Promise>> { + return Promise.race([ + reader.read(), + writer.closed.then(() => { + throw new WorkflowRuntimeError('Writable stream closed prematurely'); + }), + consumer.then( + () => { + throw new WorkflowRuntimeError('Stream consumer ended prematurely'); + }, + (err) => { + throw err; + } + ), + ]); +} + +/** + * Backpressure: once too many chunks are outstanding, stop reading until the + * consumer drains a batch. Racing the consumer avoids a deadlock if it ends or + * fails while we wait. + */ +async function awaitBackpressureRelief( + ctx: CoalesceContext, + consumer: Promise, + maxInflight: number +): Promise { + if (ctx.state.pendingOps < maxInflight) return; + await Promise.race([ + new Promise((resolve) => { + ctx.wakeProducer = resolve; + }), + consumer.then( + () => undefined, + () => undefined + ), + ]); +} + +/** + * Batching variant of {@link flushablePipe} for sinks that expose a durable + * batch write ({@link STREAM_WRITE_BATCH_SYMBOL}). A producer reads from + * `source` into a queue; {@link drainBatches} coalesces queued chunks into as + * few server writes as possible — turning a fast burst of chunks into a handful + * of round trips instead of one per chunk. + * + * Durability is preserved exactly as in the per-chunk path: `state.pendingOps` + * counts chunks that have been read but not yet durably written, so the + * lock-release completion (`pollWritableLock` / `pollReadableLock`) never fires + * while data is still queued or in flight, and the source-done path awaits every + * batch before closing. + */ +async function flushablePipeCoalescing( + source: ReadableStream, + sink: WritableStream, + state: FlushableStreamState, + batchWrite: BatchWrite +): Promise { + const reader = source.getReader(); + const writer = sink.getWriter(); + const maxInflight = getMaxInflightChunks(); + let cancelReason: unknown; + + const ctx: CoalesceContext = { + state, + batchWrite, + queue: [], + sourceDone: false, + wakeConsumer: null, + wakeProducer: null, + }; + const consumer = drainBatches(ctx); + + try { + while (!state.streamEnded) { + const readResult = await readNextChunk(reader, writer, consumer); + if (state.streamEnded) return; + + if (readResult.done) { + // Signal end-of-source and wait for every queued/in-flight batch to + // reach the server before closing, so no chunk is dropped. + ctx.sourceDone = true; + wake(ctx, 'wakeConsumer'); + await consumer; + state.streamEnded = true; + await writer.close(); + if (!state.doneResolved) { + state.doneResolved = true; + state.resolve(); + } + return; + } + + // Count the chunk as pending the moment it is read — it stays counted + // until its batch is durably written. + ctx.queue.push(readResult.value as Uint8Array); + state.pendingOps++; + wake(ctx, 'wakeConsumer'); + + await awaitBackpressureRelief(ctx, consumer, maxInflight); + } + } catch (err) { + state.streamEnded = true; + cancelReason = err; + // Let the consumer settle so its rejection (if any) is observed here rather + // than surfacing as an unhandled rejection. + ctx.sourceDone = true; + wake(ctx, 'wakeConsumer'); + await consumer.catch(() => {}); + if (!state.doneResolved) { + state.doneResolved = true; + state.reject(err); + } + throw err; + } finally { + reader.cancel(cancelReason).catch(() => {}); + reader.releaseLock(); + writer.releaseLock(); + } +} diff --git a/packages/core/src/serialization.ts b/packages/core/src/serialization.ts index 466263f5c8..1e3a48da60 100644 --- a/packages/core/src/serialization.ts +++ b/packages/core/src/serialization.ts @@ -81,6 +81,7 @@ import { STREAM_SERVER_DEPLOYMENT_ID_SYMBOL, STREAM_SERVER_RUN_ID_SYMBOL, STREAM_TYPE_SYMBOL, + STREAM_WRITE_BATCH_SYMBOL, WEBHOOK_RESPONSE_WRITABLE, } from './symbols.js'; import * as Attr from './telemetry/semantic-conventions.js'; @@ -1231,6 +1232,27 @@ export class WorkflowServerWritableStream extends WritableStream { for (const w of waiters) w.reject(abortError); }, }); + + // Batched, durable write entry point used by `flushablePipe` to coalesce + // chunks that arrive while a previous batch is still in flight into a + // single server write. It buffers every chunk and awaits one `flush()`, + // so the whole batch goes out as one `writeMulti` (or, when the world has + // no `writeMulti`, sequential `write`s) and resolves only once the batch + // has reached the server. It shares the buffer/flush machinery with the + // per-chunk sink `write()`, but the two are never used concurrently on the + // same stream: `flushablePipe` uses either this path or the writer, never + // both. On failure `flush()` retains the batch in the buffer and rethrows, + // so the caller's durability tracking stays accurate. + Object.defineProperty(this, STREAM_WRITE_BATCH_SYMBOL, { + value: async (chunks: Uint8Array[]): Promise => { + if (chunks.length === 0) return; + if (batchStartAt === undefined) batchStartAt = Date.now(); + for (const chunk of chunks) buffer.push(chunk); + await flush(); + }, + enumerable: false, + writable: false, + }); } } diff --git a/packages/core/src/symbols.ts b/packages/core/src/symbols.ts index a22712fbb3..ed95dd418d 100644 --- a/packages/core/src/symbols.ts +++ b/packages/core/src/symbols.ts @@ -34,6 +34,22 @@ export const STREAM_SERVER_RUN_ID_SYMBOL = Symbol.for( export const STREAM_SERVER_DEPLOYMENT_ID_SYMBOL = Symbol.for( 'WORKFLOW_STREAM_SERVER_DEPLOYMENT_ID' ); +/** + * Stamped on a `WorkflowServerWritableStream` instance to expose a batched, + * durable write entry point: `(chunks) => Promise` that flushes every + * chunk in one `writeMulti` (falling back to sequential `write`s when the + * world lacks `writeMulti`) and resolves only once the whole batch has + * reached the server. + * + * `flushablePipe` feature-detects this to coalesce chunks that arrive while a + * previous batch is still in flight into a single server write. Sinks without + * it (plain `WritableStream`s, `TransformStream` writables on the read path) + * fall back to the per-chunk write loop. See `flushablePipe`. + */ +export const STREAM_WRITE_BATCH_SYMBOL = Symbol.for( + 'WORKFLOW_STREAM_WRITE_BATCH' +); + export const BODY_INIT_SYMBOL = Symbol.for('BODY_INIT'); export const WEBHOOK_RESPONSE_WRITABLE = Symbol.for( 'WEBHOOK_RESPONSE_WRITABLE' diff --git a/packages/core/src/writable-stream.test.ts b/packages/core/src/writable-stream.test.ts index 858f750c02..6b2ebf48b8 100644 --- a/packages/core/src/writable-stream.test.ts +++ b/packages/core/src/writable-stream.test.ts @@ -1,5 +1,6 @@ import { SPEC_VERSION_CURRENT } from '@workflow/world'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { createFlushableState, flushablePipe } from './flushable-stream.js'; import { setWorld } from './runtime/world.js'; import { dehydrateStepReturnValue, @@ -459,4 +460,101 @@ describe('WorkflowServerWritableStream', () => { expect(order).toEqual(['barrier', 'close']); }); }); + + describe('writeMulti batching via flushablePipe', () => { + it('coalesces chunks that arrive during an in-flight write into one writeMulti', async () => { + // Hold the first server write in flight so the chunks produced while it + // is pending accumulate and flush together via writeMulti. Without the + // coalescing pipe, each chunk would be its own single write() and + // writeMulti would never be called with more than one chunk. + let releaseFirst!: () => void; + const firstInFlight = new Promise((r) => { + releaseFirst = r; + }); + mockStreams.write.mockImplementationOnce(async () => { + await firstInFlight; + }); + + const serverWritable = new WorkflowServerWritableStream( + 'run-123', + 'test-stream' + ); + const state = createFlushableState(); + + let controller!: ReadableStreamDefaultController; + const source = new ReadableStream({ + start(c) { + controller = c; + }, + }); + + const pipe = flushablePipe(source, serverWritable, state).catch(() => {}); + + // First chunk goes out alone (a single write) and blocks in flight. + controller.enqueue(new Uint8Array([1])); + await new Promise((r) => setTimeout(r, 10)); + expect(mockStreams.write).toHaveBeenCalledTimes(1); + expect(mockStreams.writeMulti).not.toHaveBeenCalled(); + + // Three more chunks arrive while the first write is in flight. + controller.enqueue(new Uint8Array([2])); + controller.enqueue(new Uint8Array([3])); + controller.enqueue(new Uint8Array([4])); + await new Promise((r) => setTimeout(r, 10)); + + // Release the first write; the queued chunks flush as ONE writeMulti. + releaseFirst(); + controller.close(); + await pipe; + + expect(mockStreams.writeMulti).toHaveBeenCalledTimes(1); + const [, , batched] = mockStreams.writeMulti.mock.calls[0]; + expect(batched.map((c: Uint8Array) => c[0])).toEqual([2, 3, 4]); + expect(mockStreams.close).toHaveBeenCalledTimes(1); + await expect(state.promise).resolves.toBeUndefined(); + }); + + it('falls back to sequential writes when the world lacks writeMulti', async () => { + delete (mockStreams as { writeMulti?: unknown }).writeMulti; + + let releaseFirst!: () => void; + const firstInFlight = new Promise((r) => { + releaseFirst = r; + }); + mockStreams.write.mockImplementationOnce(async () => { + await firstInFlight; + }); + + const serverWritable = new WorkflowServerWritableStream( + 'run-123', + 'test-stream' + ); + const state = createFlushableState(); + + let controller!: ReadableStreamDefaultController; + const source = new ReadableStream({ + start(c) { + controller = c; + }, + }); + const pipe = flushablePipe(source, serverWritable, state).catch(() => {}); + + controller.enqueue(new Uint8Array([1])); + await new Promise((r) => setTimeout(r, 10)); + controller.enqueue(new Uint8Array([2])); + controller.enqueue(new Uint8Array([3])); + await new Promise((r) => setTimeout(r, 10)); + releaseFirst(); + controller.close(); + await pipe; + + // All three chunks are delivered via sequential write() calls. + expect(mockStreams.write).toHaveBeenCalledTimes(3); + const delivered = mockStreams.write.mock.calls.map( + (call: unknown[]) => (call[2] as Uint8Array)[0] + ); + expect(delivered).toEqual([1, 2, 3]); + await expect(state.promise).resolves.toBeUndefined(); + }); + }); }); From 3cb6bf9ca852d2957abc41ed4a47721d41fdd922 Mon Sep 17 00:00:00 2001 From: Peter Wielander Date: Mon, 20 Jul 2026 08:46:04 -0700 Subject: [PATCH 2/2] fix(core): cap coalesced batches at chunk/byte wire limits; keep pendingOps exact on failure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address PR review: - Split each coalesced writeMulti at a chunk-count cap (WORKFLOW_STREAM_MAX_CHUNKS_PER_BATCH, default 1000) and a byte cap (WORKFLOW_STREAM_MAX_BYTES_PER_BATCH, default 1 MiB), independent of the read-ahead backpressure knob — so raising backpressure can't push one request past the server's per-multi-write chunk limit or a platform request-body limit. A single chunk larger than the byte cap still goes out alone. - Decrement pendingOps only after batchWrite succeeds (was in a finally): on failure flush() retains the batch in the sink buffer, so those chunks are still read-but-not-durable and must keep counting. - Document the no-writeMulti fallback stall bound and the three env vars. - Tests: chunk-count split, byte split, oversized-single-chunk, failed batch keeps pendingOps, and failed batch does not re-send or close. Co-Authored-By: Claude Opus 4.8 (1M context) --- .changeset/stream-writemulti-batching.md | 2 +- .../docs/v5/configuration/runtime-tuning.mdx | 15 +++ packages/core/src/flushable-stream.test.ts | 107 ++++++++++++++++++ packages/core/src/flushable-stream.ts | 96 +++++++++++++--- packages/core/src/serialization.ts | 20 ++-- packages/core/src/writable-stream.test.ts | 46 ++++++++ 6 files changed, 260 insertions(+), 26 deletions(-) diff --git a/.changeset/stream-writemulti-batching.md b/.changeset/stream-writemulti-batching.md index f291c7c83e..15eaed2a1d 100644 --- a/.changeset/stream-writemulti-batching.md +++ b/.changeset/stream-writemulti-batching.md @@ -2,4 +2,4 @@ '@workflow/core': patch --- -Fix stream writes never batching: `flushablePipe` awaited each `writer.write()` and the WritableStream sink serialized chunks one at a time, so the server writable's buffer never held more than one chunk and its `writeMulti` batching path never engaged — every chunk became its own server round trip. The server writable now exposes a durable batch write that `flushablePipe` uses to coalesce chunks arriving while a previous write is in flight into a single `writeMulti`, while preserving the existing durability guarantees. +Fix stream writes never batching: `flushablePipe` awaited each `writer.write()` and the WritableStream sink serialized chunks one at a time, so the server writable's buffer never held more than one chunk and its `writeMulti` batching path never engaged — every chunk became its own server round trip. The server writable now exposes a durable batch write that `flushablePipe` uses to coalesce chunks arriving while a previous write is in flight into a single `writeMulti`, while preserving the existing durability guarantees. Coalesced batches are split at chunk-count and byte wire limits (`WORKFLOW_STREAM_MAX_CHUNKS_PER_BATCH`, `WORKFLOW_STREAM_MAX_BYTES_PER_BATCH`), independent of the read-ahead backpressure bound (`WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS`). diff --git a/docs/content/docs/v5/configuration/runtime-tuning.mdx b/docs/content/docs/v5/configuration/runtime-tuning.mdx index 34603add1e..c8dd85b454 100644 --- a/docs/content/docs/v5/configuration/runtime-tuning.mdx +++ b/docs/content/docs/v5/configuration/runtime-tuning.mdx @@ -133,6 +133,21 @@ These variables are primarily for tests, debugging, or unusual deployments. - Stream write buffering interval. - Also available as `streamFlushIntervalMs` on Worlds that expose it. +### `WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS` + +- Default: `1000` +- Flow-control bound: how many stream chunks may be read-but-not-yet-durably-written while writes are coalesced. Once this many are outstanding, the producer pauses reading until a batch is written. Bounds in-memory buffering; distinct from the per-request batch caps below. + +### `WORKFLOW_STREAM_MAX_CHUNKS_PER_BATCH` + +- Default: `1000` +- Wire limit: maximum number of chunks sent in a single coalesced multi-write. Larger bursts are split across requests so one request can't exceed the server's per-multi-write chunk cap. + +### `WORKFLOW_STREAM_MAX_BYTES_PER_BATCH` + +- Default: `1048576` (1 MiB) +- Wire limit: maximum cumulative bytes in a single coalesced multi-write, so large chunks don't produce a request body that platform limits reject. A single chunk larger than this is still sent on its own. + ### `WORKFLOW_FRAMED_STREAM_MAX_RECONNECTS` - Default: `50` diff --git a/packages/core/src/flushable-stream.test.ts b/packages/core/src/flushable-stream.test.ts index 1da6a51742..21503acfcf 100644 --- a/packages/core/src/flushable-stream.test.ts +++ b/packages/core/src/flushable-stream.test.ts @@ -444,6 +444,8 @@ describe('flushable stream behavior', () => { describe('flushablePipe batching (STREAM_WRITE_BATCH_SYMBOL sinks)', () => { afterEach(() => { delete process.env.WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS; + delete process.env.WORKFLOW_STREAM_MAX_CHUNKS_PER_BATCH; + delete process.env.WORKFLOW_STREAM_MAX_BYTES_PER_BATCH; }); it('coalesces chunks that arrive during an in-flight batch into one write', async () => { @@ -552,6 +554,111 @@ describe('flushablePipe batching (STREAM_WRITE_BATCH_SYMBOL sinks)', () => { expect(state.streamEnded).toBe(true); }); + it('does NOT decrement pendingOps for a failed (retained, un-durable) batch', async () => { + // On failure the sink retains the batch in its buffer, so those chunks are + // still "read but not durable" — pendingOps must keep counting them. + const { sink } = makeBatchSink((call) => { + if (call === 0) throw new Error('batch write failed'); + }); + const { source, controller } = makeControlledSource(); + const state = createFlushableState(); + + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + controller().enqueue(new Uint8Array([1])); + controller().enqueue(new Uint8Array([2])); + await tick(); + + await pipe; + await expect(state.promise).rejects.toThrow('batch write failed'); + // Both chunks were read; the batch write failed and retained them, so the + // count stays at 2 rather than being silently zeroed in a `finally`. + expect(state.pendingOps).toBe(2); + }); + + it('splits a coalesced batch at the chunk-count wire limit', async () => { + process.env.WORKFLOW_STREAM_MAX_CHUNKS_PER_BATCH = '2'; + + let releaseFirst!: () => void; + const firstInFlight = new Promise((r) => { + releaseFirst = r; + }); + const { sink, batches } = makeBatchSink((call) => + call === 0 ? firstInFlight : undefined + ); + const { source, controller } = makeControlledSource(); + const state = createFlushableState(); + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + // First chunk goes out alone and blocks; five more queue behind it. + controller().enqueue(new Uint8Array([1])); + await tick(); + for (const n of [2, 3, 4, 5, 6]) controller().enqueue(new Uint8Array([n])); + await tick(); + + releaseFirst(); + controller().close(); + await pipe; + + // No batch exceeds the 2-chunk cap, and every chunk is delivered in order. + expect(batches.every((b) => b.length <= 2)).toBe(true); + expect(batches.flat().map((c) => c[0])).toEqual([1, 2, 3, 4, 5, 6]); + await expect(state.promise).resolves.toBeUndefined(); + }); + + it('splits a coalesced batch at the byte wire limit', async () => { + // Cap at 10 bytes; each chunk is 4 bytes, so at most 2 fit per batch. + process.env.WORKFLOW_STREAM_MAX_BYTES_PER_BATCH = '10'; + + let releaseFirst!: () => void; + const firstInFlight = new Promise((r) => { + releaseFirst = r; + }); + const { sink, batches } = makeBatchSink((call) => + call === 0 ? firstInFlight : undefined + ); + const { source, controller } = makeControlledSource(); + const state = createFlushableState(); + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + controller().enqueue(new Uint8Array([0, 0, 0, 1])); + await tick(); + for (let i = 2; i <= 5; i++) { + controller().enqueue(new Uint8Array([0, 0, 0, i])); + } + await tick(); + + releaseFirst(); + controller().close(); + await pipe; + + for (const b of batches) { + const bytes = b.reduce((sum, c) => sum + c.byteLength, 0); + expect(bytes).toBeLessThanOrEqual(10); + } + expect(batches.flat().map((c) => c[3])).toEqual([1, 2, 3, 4, 5]); + await expect(state.promise).resolves.toBeUndefined(); + }); + + it('sends an oversized single chunk alone rather than stalling', async () => { + process.env.WORKFLOW_STREAM_MAX_BYTES_PER_BATCH = '4'; + + const { sink, batches } = makeBatchSink(); + const { source, controller } = makeControlledSource(); + const state = createFlushableState(); + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + // A single chunk larger than the byte cap must still be delivered (alone). + controller().enqueue(new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8])); + controller().close(); + await pipe; + + expect(batches).toHaveLength(1); + expect(batches[0]).toHaveLength(1); + expect(batches[0][0].byteLength).toBe(8); + await expect(state.promise).resolves.toBeUndefined(); + }); + it('applies backpressure once too many chunks are outstanding', async () => { process.env.WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS = '2'; diff --git a/packages/core/src/flushable-stream.ts b/packages/core/src/flushable-stream.ts index dac37c0b79..725f9dcb37 100644 --- a/packages/core/src/flushable-stream.ts +++ b/packages/core/src/flushable-stream.ts @@ -11,11 +11,16 @@ import { STREAM_WRITE_BATCH_SYMBOL } from './symbols.js'; type BatchWrite = (chunks: Uint8Array[]) => Promise; /** - * Upper bound on chunks read-but-not-yet-durably-written while coalescing. - * Once this many chunks are outstanding the producer stops reading until the - * consumer drains a batch, so a fast producer paired with a slow server can't - * grow the in-memory queue without bound. Override: + * Flow-control knob: upper bound on chunks read-but-not-yet-durably-written + * while coalescing. Once this many chunks are outstanding the producer stops + * reading until the consumer drains a batch, so a fast producer paired with a + * slow server can't grow the in-memory queue without bound. Override: * `WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS`. + * + * This is deliberately distinct from the per-request batch caps below: this + * bounds how much is *buffered*, those bound how much goes out in one + * `writeMulti`. Raising this must never let a single request exceed a wire + * limit — batch sizing enforces that independently. */ export const MAX_INFLIGHT_CHUNKS = 1000; @@ -25,6 +30,36 @@ const getMaxInflightChunks = (): number => min: 1, }); +/** + * Wire limit: maximum number of chunks in a single coalesced `writeMulti`. + * The server enforces a per-multi-write chunk cap (1,000 today); a batch is + * split at this bound so it can never be rejected wholesale, independently of + * the backpressure knob above. Override: `WORKFLOW_STREAM_MAX_CHUNKS_PER_BATCH`. + */ +export const MAX_CHUNKS_PER_BATCH = 1000; + +const getMaxChunksPerBatch = (): number => + envNumber('WORKFLOW_STREAM_MAX_CHUNKS_PER_BATCH', MAX_CHUNKS_PER_BATCH, { + integer: true, + min: 1, + }); + +/** + * Wire limit: maximum cumulative bytes in a single coalesced `writeMulti`. + * Chunk *count* alone is not enough — 1,000 small chunks are ~100KB but 1,000 + * file-sized chunks can be hundreds of MB, which platform request-body limits + * reject long before the count cap matters. A batch is split once adding the + * next chunk would exceed this (a single chunk larger than the cap still goes + * out alone). Default 1 MiB. Override: `WORKFLOW_STREAM_MAX_BYTES_PER_BATCH`. + */ +export const MAX_BYTES_PER_BATCH = 1024 * 1024; + +const getMaxBytesPerBatch = (): number => + envNumber('WORKFLOW_STREAM_MAX_BYTES_PER_BATCH', MAX_BYTES_PER_BATCH, { + integer: true, + min: 1, + }); + /** * Polling interval (in ms) for lock release detection. * @@ -354,6 +389,10 @@ interface CoalesceContext { wakeConsumer: (() => void) | null; /** Resolver that wakes the producer when the consumer relieves backpressure. */ wakeProducer: (() => void) | null; + /** Wire limit: max chunks per coalesced `writeMulti`. */ + maxChunksPerBatch: number; + /** Wire limit: max cumulative bytes per coalesced `writeMulti`. */ + maxBytesPerBatch: number; } function wake(ctx: CoalesceContext, which: 'wakeConsumer' | 'wakeProducer') { @@ -365,11 +404,32 @@ function wake(ctx: CoalesceContext, which: 'wakeConsumer' | 'wakeProducer') { } /** - * Consumer loop: drains the whole queue into one `batchWrite` at a time, - * awaiting each batch so it is durable on the server before its chunks are - * removed from `pendingOps`. Because it takes *everything* queued each round, - * chunks that arrived while the previous batch was in flight coalesce into the - * next server write. + * How many leading queued chunks fit in one `writeMulti` without crossing the + * chunk-count or byte wire limits. Always returns at least 1 so a single chunk + * larger than the byte cap still makes progress (alone). + */ +function nextBatchSize(ctx: CoalesceContext): number { + let count = 0; + let bytes = 0; + for (const chunk of ctx.queue) { + if (count >= ctx.maxChunksPerBatch) break; + if (count > 0 && bytes + chunk.byteLength > ctx.maxBytesPerBatch) break; + count++; + bytes += chunk.byteLength; + } + return count; +} + +/** + * Consumer loop: drains the queue one batch at a time, awaiting each batch so it + * is durable on the server before its chunks leave `pendingOps`. It takes as + * much of the queue as the wire limits allow each round, so chunks that arrived + * while the previous batch was in flight coalesce into the next server write. + * + * `pendingOps` is decremented only after `batchWrite` *succeeds*: on failure the + * chunks stay retained in the sink's buffer (see `WorkflowServerWritableStream`) + * and are therefore still "read but not durable", so they must keep counting. + * A throw propagates out of the loop; the pipe's producer then tears down. */ async function drainBatches(ctx: CoalesceContext): Promise { while (true) { @@ -380,15 +440,13 @@ async function drainBatches(ctx: CoalesceContext): Promise { }); continue; } - const batch = ctx.queue; - ctx.queue = []; - try { - await ctx.batchWrite(batch); - } finally { - ctx.state.pendingOps -= batch.length; - // Relieve producer backpressure now that this batch is durable. - wake(ctx, 'wakeProducer'); - } + const count = nextBatchSize(ctx); + const batch = ctx.queue.slice(0, count); + ctx.queue = ctx.queue.slice(count); + await ctx.batchWrite(batch); + ctx.state.pendingOps -= batch.length; + // Relieve producer backpressure now that this batch is durable. + wake(ctx, 'wakeProducer'); } } @@ -473,6 +531,8 @@ async function flushablePipeCoalescing( sourceDone: false, wakeConsumer: null, wakeProducer: null, + maxChunksPerBatch: getMaxChunksPerBatch(), + maxBytesPerBatch: getMaxBytesPerBatch(), }; const consumer = drainBatches(ctx); diff --git a/packages/core/src/serialization.ts b/packages/core/src/serialization.ts index 1e3a48da60..ffd5815d01 100644 --- a/packages/core/src/serialization.ts +++ b/packages/core/src/serialization.ts @@ -1236,13 +1236,19 @@ export class WorkflowServerWritableStream extends WritableStream { // Batched, durable write entry point used by `flushablePipe` to coalesce // chunks that arrive while a previous batch is still in flight into a // single server write. It buffers every chunk and awaits one `flush()`, - // so the whole batch goes out as one `writeMulti` (or, when the world has - // no `writeMulti`, sequential `write`s) and resolves only once the batch - // has reached the server. It shares the buffer/flush machinery with the - // per-chunk sink `write()`, but the two are never used concurrently on the - // same stream: `flushablePipe` uses either this path or the writer, never - // both. On failure `flush()` retains the batch in the buffer and rethrows, - // so the caller's durability tracking stays accurate. + // so the whole batch goes out as one `writeMulti` and resolves only once + // the batch has reached the server. It shares the buffer/flush machinery + // with the per-chunk sink `write()`, but the two are never used + // concurrently on the same stream: `flushablePipe` uses either this path + // or the writer, never both. On failure `flush()` retains the batch in the + // buffer and rethrows, so the caller's durability tracking stays accurate. + // + // No-`writeMulti` fallback: when the world lacks `writeMulti`, `flush()` + // degrades to sequential `write`s for the batch's chunks — one round trip + // each, within this single call, while backpressure holds the producer. + // `flushablePipe` bounds that stall by capping each coalesced batch (see + // `MAX_CHUNKS_PER_BATCH` / `MAX_BYTES_PER_BATCH`), so the fallback can't + // turn one drain into an unbounded sequential run. Object.defineProperty(this, STREAM_WRITE_BATCH_SYMBOL, { value: async (chunks: Uint8Array[]): Promise => { if (chunks.length === 0) return; diff --git a/packages/core/src/writable-stream.test.ts b/packages/core/src/writable-stream.test.ts index 6b2ebf48b8..8627f7abcd 100644 --- a/packages/core/src/writable-stream.test.ts +++ b/packages/core/src/writable-stream.test.ts @@ -556,5 +556,51 @@ describe('WorkflowServerWritableStream', () => { expect(delivered).toEqual([1, 2, 3]); await expect(state.promise).resolves.toBeUndefined(); }); + + it('errors the stream on a failed batch without re-sending or closing', async () => { + // Let the first (single) write land, then make the coalesced writeMulti + // fail. flush() retains the batch in the buffer and rethrows; the pipe + // errors. Nothing should re-send the buffered chunks or close the stream. + let releaseFirst!: () => void; + const firstInFlight = new Promise((r) => { + releaseFirst = r; + }); + mockStreams.write.mockImplementationOnce(async () => { + await firstInFlight; + }); + mockStreams.writeMulti.mockRejectedValueOnce(new Error('batch failed')); + + const serverWritable = new WorkflowServerWritableStream( + 'run-123', + 'test-stream' + ); + const state = createFlushableState(); + + let controller!: ReadableStreamDefaultController; + const source = new ReadableStream({ + start(c) { + controller = c; + }, + }); + const pipe = flushablePipe(source, serverWritable, state).catch(() => {}); + + controller.enqueue(new Uint8Array([1])); + await new Promise((r) => setTimeout(r, 10)); + controller.enqueue(new Uint8Array([2])); + controller.enqueue(new Uint8Array([3])); + await new Promise((r) => setTimeout(r, 10)); + releaseFirst(); + await pipe; + + await expect(state.promise).rejects.toThrow('batch failed'); + expect(mockStreams.writeMulti).toHaveBeenCalledTimes(1); + expect(mockStreams.close).not.toHaveBeenCalled(); + + // Wait past any flush interval: the retained chunks are not re-sent by a + // leaked timer or later path, and the stream stays closed-off. + await new Promise((r) => setTimeout(r, 30)); + expect(mockStreams.writeMulti).toHaveBeenCalledTimes(1); + expect(mockStreams.close).not.toHaveBeenCalled(); + }); }); });