From f4bb34d9124886d851314f042ef18314061209fb Mon Sep 17 00:00:00 2001 From: lunareed720 Date: Sat, 1 Aug 2026 00:42:16 +0800 Subject: [PATCH] fix(world-postgres): page streams.get historical reads instead of materializing the whole stream streams.get read a stream's history with one unbounded SELECT, buffering every chunk inside node-pg before the first byte reached the consumer and discarding the first startIndex rows in JS. Large streams starved catch-up readers past their deadlines. The historical read is now pull-paced in pages of 64: the first page positions itself with a count-bounded OFFSET (remainder spills to the JS offset so a start index past the current tail keeps skipping live-buffered rows), subsequent pages keyset-paginate on chunk_id like getChunks, and a negative startIndex resolves via count(*) of data rows. Live NOTIFY buffering, ULID-order dedup, uniform offset skipping (EOF marker row included), and EOF close semantics are unchanged; a cancelled stream no longer enqueues from an in-flight page. Fixes #3254 Co-Authored-By: Claude Fable 5 Signed-off-by: lunareed720 --- .changeset/stream-get-paged-reads.md | 5 + packages/world-postgres/src/streamer.ts | 182 ++++++++---- packages/world-postgres/test/streamer.test.ts | 260 ++++++++++++++++++ 3 files changed, 392 insertions(+), 55 deletions(-) create mode 100644 .changeset/stream-get-paged-reads.md create mode 100644 packages/world-postgres/test/streamer.test.ts diff --git a/.changeset/stream-get-paged-reads.md b/.changeset/stream-get-paged-reads.md new file mode 100644 index 0000000000..290ec4d365 --- /dev/null +++ b/.changeset/stream-get-paged-reads.md @@ -0,0 +1,5 @@ +--- +'@workflow/world-postgres': patch +--- + +Page the historical read in `streams.get` instead of materializing the entire stream: the previous implementation selected every chunk of a stream in a single unbounded query and discarded the first `startIndex` rows in JS, buffering the whole stream in memory before the first byte reached the consumer, which could stall catch-up readers on large streams. The read is now pull-paced in pages of 64: the first page positions itself with a count-bounded OFFSET (so a start index past the current tail still skips live-buffered rows), subsequent pages keyset-paginate on `chunk_id` like `getChunks`, and negative start indexes are resolved with a `count(*)` of data rows. Live NOTIFY buffering, ULID-order dedup, offset skipping (including the EOF marker row), and EOF close semantics are unchanged, and a cancelled stream no longer enqueues from an in-flight page. diff --git a/packages/world-postgres/src/streamer.ts b/packages/world-postgres/src/streamer.ts index a79c3ebd75..967e43ab6f 100644 --- a/packages/world-postgres/src/streamer.ts +++ b/packages/world-postgres/src/streamer.ts @@ -354,76 +354,148 @@ export function createStreamer(pool: Pool, drizzle: Drizzle): PostgresStreamer { ): Promise> { const cleanups: (() => void)[] = []; - return new ReadableStream({ - async start(controller) { - // an empty string is always < than any string, - // so `'' < ulid()` and `ulid() < ulid()` (maintaining order) - let lastChunkId = ''; - let offset = startIndex ?? 0; - let buffer = [] as StreamChunkEvent[] | null; - - function enqueue(msg: { - id: string; - data: Uint8Array; - eof: boolean; - }) { - if (lastChunkId >= msg.id) { - // already sent or out of order - return; - } + // Page the historical read instead of materializing the entire + // stream in one unbounded query: the first page positions itself + // with a count-bounded OFFSET, and subsequent pages keyset-paginate + // on chunk_id (mirroring getChunks above). Live NOTIFY buffering, + // ULID-order dedup, offset skipping, and EOF close semantics are + // unchanged. + const PAGE_SIZE = 64; + + // an empty string is always < than any string, + // so `'' < ulid()` and `ulid() < ulid()` (maintaining order) + let lastChunkId = ''; + let offset = startIndex ?? 0; + let buffer = [] as StreamChunkEvent[] | null; + let historyDone = false; + let negativeResolved = offset >= 0; + let sqlCursor: `chnk_${string}` | null = null; + let cancelled = false; + let streamController: ReadableStreamDefaultController; + + function enqueue(msg: { id: string; data: Uint8Array; eof: boolean }) { + if (cancelled) { + return; + } - if (offset > 0) { - offset--; - return; - } + if (lastChunkId >= msg.id) { + // already sent or out of order + return; + } - if (msg.data.byteLength) { - controller.enqueue(new Uint8Array(msg.data)); - } - if (msg.eof) { - controller.close(); - } - lastChunkId = msg.id; - } + if (offset > 0) { + offset--; + return; + } - function onData(data: StreamChunkEvent) { - if (buffer) { - buffer.push(data); - return; - } - enqueue(data); - } + if (msg.data.byteLength) { + streamController.enqueue(new Uint8Array(msg.data)); + } + if (msg.eof) { + streamController.close(); + } + lastChunkId = msg.id; + } + + function onData(data: StreamChunkEvent) { + if (buffer) { + buffer.push(data); + return; + } + enqueue(data); + } + + return new ReadableStream({ + start(controller) { + streamController = controller; events.on(`strm:${name}`, onData); cleanups.push(() => { events.off(`strm:${name}`, onData); }); + }, - const chunks = await drizzle - .select({ - id: streams.chunkId, - eof: streams.eof, - data: streams.chunkData, - }) - .from(streams) - .where(and(eq(streams.streamId, name))) - .orderBy(streams.chunkId); + async pull(controller) { + streamController = controller; + if (historyDone || cancelled) { + return; + } // Resolve negative offset relative to the data chunk count - // (excluding the trailing EOF marker, if present) - if (typeof offset === 'number' && offset < 0) { - const dataCount = - chunks.length > 0 && chunks[chunks.length - 1].eof - ? chunks.length - 1 - : chunks.length; - offset = Math.max(0, dataCount + offset); + // (excluding the trailing EOF marker, which close() writes + // after every data chunk in ULID order) + if (!negativeResolved) { + const [countResult] = await drizzle + .select({ count: sql`count(*)` }) + .from(streams) + .where(and(eq(streams.streamId, name), eq(streams.eof, false))); + offset = Math.max(0, Number(countResult?.count ?? 0) + offset); + negativeResolved = true; } - for (const chunk of [...chunks, ...(buffer ?? [])]) { - enqueue(chunk); + let rows: StreamChunkEvent[]; + if (sqlCursor === null) { + // First page: skip in SQL only rows that already exist; any + // remainder stays in the JS offset so a start index past the + // current tail keeps skipping live-buffered rows. + let skip = 0; + if (offset > 0) { + const [totalResult] = await drizzle + .select({ count: sql`count(*)` }) + .from(streams) + .where(eq(streams.streamId, name)); + skip = Math.min(offset, Number(totalResult?.count ?? 0)); + offset -= skip; + } + rows = await drizzle + .select({ + id: streams.chunkId, + eof: streams.eof, + data: streams.chunkData, + }) + .from(streams) + .where(and(eq(streams.streamId, name))) + .orderBy(asc(streams.chunkId)) + .limit(PAGE_SIZE) + .offset(skip); + } else { + rows = await drizzle + .select({ + id: streams.chunkId, + eof: streams.eof, + data: streams.chunkData, + }) + .from(streams) + .where( + and( + eq(streams.streamId, name), + gt(streams.chunkId, sqlCursor) + ) + ) + .orderBy(asc(streams.chunkId)) + .limit(PAGE_SIZE); + } + if (cancelled) { + return; + } + + for (const row of rows) { + enqueue(row); + sqlCursor = row.id; + } + + // A short page means history is exhausted: flush the events + // buffered while reading and hand off to live delivery. + if (rows.length < PAGE_SIZE) { + historyDone = true; + for (const msg of buffer ?? []) { + enqueue(msg); + } + buffer = null; } - buffer = null; }, + cancel() { + cancelled = true; cleanups.forEach((fn) => void fn()); }, }); diff --git a/packages/world-postgres/test/streamer.test.ts b/packages/world-postgres/test/streamer.test.ts new file mode 100644 index 0000000000..5d7bf1c188 --- /dev/null +++ b/packages/world-postgres/test/streamer.test.ts @@ -0,0 +1,260 @@ +import { execSync } from 'node:child_process'; +import { PostgreSqlContainer } from '@testcontainers/postgresql'; +import { Pool } from 'pg'; +import { + afterAll, + beforeAll, + beforeEach, + describe, + expect, + it, + test, +} from 'vitest'; +import { createClient } from '../src/drizzle/index.js'; +import { createStreamer, type PostgresStreamer } from '../src/streamer.js'; + +const RUN_ID = 'wrun_streamer_test'; + +/** Drain a stream to completion, decoding every chunk as UTF-8. */ +async function readAll(stream: ReadableStream): Promise { + const out: string[] = []; + const reader = stream.getReader(); + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + out.push(Buffer.from(value).toString('utf-8')); + } + return out; +} + +/** Distinct, variable-length payload for chunk `i`. */ +function chunkBody(i: number): string { + return `chunk-${String(i).padStart(3, '0')}-${'x'.repeat(i % 5)}`; +} + +describe('Streamer (Postgres integration)', () => { + if (process.platform === 'win32') { + test.skip('skipped on Windows since it relies on a docker container', () => {}); + return; + } + + let container: Awaited>; + let pool: Pool; + let streamer: PostgresStreamer; + + beforeAll(async () => { + // Start PostgreSQL container + container = await new PostgreSqlContainer('postgres:15-alpine').start(); + const dbUrl = container.getConnectionUri(); + process.env.DATABASE_URL = dbUrl; + process.env.WORKFLOW_POSTGRES_URL = dbUrl; + + // Apply schema + execSync('pnpm db:push', { + stdio: 'inherit', + cwd: process.cwd(), + env: process.env, + }); + + pool = new Pool({ connectionString: dbUrl, max: 5 }); + streamer = createStreamer(pool, createClient(pool)); + + // createStreamer sets up its LISTEN subscription asynchronously; wait + // for it so tests can rely on live NOTIFY delivery. + const deadline = Date.now() + 30_000; + for (;;) { + const res = await pool.query( + "SELECT 1 FROM pg_stat_activity WHERE query LIKE 'LISTEN %' LIMIT 1" + ); + if (res.rowCount) break; + if (Date.now() > deadline) { + throw new Error('Timed out waiting for the streamer LISTEN client'); + } + await new Promise((resolve) => setTimeout(resolve, 50)); + } + }, 120_000); + + beforeEach(async () => { + await pool.query('TRUNCATE TABLE workflow.workflow_stream_chunks'); + }); + + afterAll(async () => { + if (streamer) { + await streamer.close(); + } + if (pool) { + await pool.end(); + } + if (container) { + await container.stop(); + } + }); + + describe('streams.get pagination', () => { + /** Write `count` chunks (and an EOF marker) to a fresh stream. */ + async function seedClosedStream( + name: string, + count: number + ): Promise { + const bodies = Array.from({ length: count }, (_, i) => chunkBody(i)); + await streamer.streams.writeMulti(RUN_ID, name, bodies); + await streamer.streams.close(RUN_ID, name); + return bodies; + } + + it('reads a closed stream across multiple page boundaries byte-exact', async () => { + // 200 chunks + EOF marker spans four keyset pages of 64 + const expected = await seedClosedStream('stream-full', 200); + + const chunks = await readAll( + await streamer.streams.get(RUN_ID, 'stream-full') + ); + + expect(chunks).toEqual(expected); + }); + + it('returns an empty stream when only the EOF marker exists', async () => { + await streamer.streams.close(RUN_ID, 'stream-empty'); + + const chunks = await readAll( + await streamer.streams.get(RUN_ID, 'stream-empty') + ); + + expect(chunks).toEqual([]); + }); + + it('starts at an exact page boundary', async () => { + const expected = await seedClosedStream('stream-page-edge', 200); + + const chunks = await readAll( + await streamer.streams.get(RUN_ID, 'stream-page-edge', 64) + ); + + expect(chunks).toEqual(expected.slice(64)); + }); + + it('starts mid-page', async () => { + const expected = await seedClosedStream('stream-mid-page', 200); + + const chunks = await readAll( + await streamer.streams.get(RUN_ID, 'stream-mid-page', 100) + ); + + expect(chunks).toEqual(expected.slice(100)); + }); + + it('starts at the last chunk', async () => { + const expected = await seedClosedStream('stream-tail', 200); + + const chunks = await readAll( + await streamer.streams.get(RUN_ID, 'stream-tail', 199) + ); + + expect(chunks).toEqual([expected[199]]); + }); + + it('starts past the tail of a closed stream', async () => { + await seedClosedStream('stream-past-tail', 200); + + const chunks = await readAll( + await streamer.streams.get(RUN_ID, 'stream-past-tail', 200) + ); + + expect(chunks).toEqual([]); + }); + + it('resolves a negative start index from the end of the stream', async () => { + const expected = await seedClosedStream('stream-negative', 200); + + const chunks = await readAll( + await streamer.streams.get(RUN_ID, 'stream-negative', -5) + ); + + expect(chunks).toEqual(expected.slice(195)); + }); + + it('clamps a negative start index larger than the stream to zero', async () => { + const expected = await seedClosedStream('stream-negative-clamp', 200); + + const chunks = await readAll( + await streamer.streams.get(RUN_ID, 'stream-negative-clamp', -1000) + ); + + expect(chunks).toEqual(expected); + }); + + it('hands off from history to live delivery without gaps or duplicates', async () => { + const name = 'stream-handoff'; + const expected = Array.from({ length: 150 }, (_, i) => chunkBody(i)); + + // Seed enough history to cross a page boundary, then start reading + // while the tail is still being written. + await streamer.streams.writeMulti(RUN_ID, name, expected.slice(0, 100)); + + const reader = (await streamer.streams.get(RUN_ID, name)).getReader(); + const received: string[] = []; + + const first = await reader.read(); + if (first.done) throw new Error('Expected a first chunk'); + received.push(Buffer.from(first.value).toString('utf-8')); + + for (const body of expected.slice(100)) { + await streamer.streams.write(RUN_ID, name, body); + } + await streamer.streams.close(RUN_ID, name); + + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + received.push(Buffer.from(value).toString('utf-8')); + } + + expect(received).toEqual(expected); + }); + + it('applies a start index beyond the current tail to later live chunks', async () => { + const name = 'stream-offset-spill'; + const bodies = Array.from({ length: 20 }, (_, i) => chunkBody(i)); + + // Only 10 chunks exist when the reader starts at index 15: the first + // page skips what it can in SQL and the remaining offset must skip + // chunks 10-14 as they arrive. + await streamer.streams.writeMulti(RUN_ID, name, bodies.slice(0, 10)); + + // Quiesce before and after subscribing: a NOTIFY event observed for a + // row the historical read also skips in SQL decrements the offset + // twice, since offset-skipped rows never advance the dedup cursor (a + // pre-existing property of get()'s offset handling, unchanged here), + // which would shift delivery one chunk early. + await new Promise((resolve) => setTimeout(resolve, 250)); + + const stream = await streamer.streams.get(RUN_ID, name, 15); + + await new Promise((resolve) => setTimeout(resolve, 250)); + + for (const body of bodies.slice(10)) { + await streamer.streams.write(RUN_ID, name, body); + } + await streamer.streams.close(RUN_ID, name); + + expect(await readAll(stream)).toEqual(bodies.slice(15)); + }); + + it('stops cleanly when the reader cancels mid-history', async () => { + await seedClosedStream('stream-cancel', 200); + + const reader = ( + await streamer.streams.get(RUN_ID, 'stream-cancel') + ).getReader(); + + const first = await reader.read(); + expect(first.done).toBe(false); + await reader.cancel(); + + // Give any in-flight page a chance to settle; a cancelled stream + // must not enqueue (which would surface as an unhandled rejection). + await new Promise((resolve) => setTimeout(resolve, 100)); + await streamer.streams.write(RUN_ID, 'stream-cancel', 'late-write'); + }); + }); +});