diff --git a/.changeset/world-postgres-read-after-eof.md b/.changeset/world-postgres-read-after-eof.md new file mode 100644 index 0000000000..9a140114d6 --- /dev/null +++ b/.changeset/world-postgres-read-after-eof.md @@ -0,0 +1,5 @@ +--- +'@workflow/world-postgres': patch +--- + +Fix `readFromStream` erroring the whole stream — and dropping every chunk still queued — when rows were written after the stream's first EOF marker (e.g. a producer that retried its terminal write). Rows past the first EOF are now ignored. diff --git a/packages/world-postgres/src/streamer.test.ts b/packages/world-postgres/src/streamer.test.ts new file mode 100644 index 0000000000..948afacc03 --- /dev/null +++ b/packages/world-postgres/src/streamer.test.ts @@ -0,0 +1,141 @@ +import type { Pool } from 'pg'; +import { describe, expect, it, vi } from 'vitest'; +import type { Drizzle } from './drizzle/index.js'; +import { createStreamer } from './streamer.js'; + +// `listenChannel` opens a dedicated LISTEN client; the live-tail leg is not +// under test here, so stub it out rather than open a socket. +vi.mock('pg', () => ({ + Client: class { + async connect() {} + async query() {} + on() {} + removeListener() {} + async end() {} + }, +})); + +type Row = { id: string; eof: boolean; data: Buffer }; + +/** The one query the catch-up loop runs: select(...).from(...).where(...).orderBy(...). */ +function fakeDrizzle(rows: Row[]): Drizzle { + const chain = { + from: () => chain, + where: () => chain, + orderBy: async () => rows, + limit: async () => rows, + }; + return { select: () => chain } as unknown as Drizzle; +} + +async function drain(stream: ReadableStream): Promise { + const reader = stream.getReader(); + const decoder = new TextDecoder(); + let out = ''; + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + out += decoder.decode(value, { stream: true }); + } + return out; +} + +describe('streams.get()', () => { + it('ignores rows written after the first EOF instead of erroring the stream', async () => { + // Several data rows before the EOF, on purpose: the catch-up loop is + // synchronous, so all but the first sit in the controller's queue when + // the EOF closes it. A post-EOF `enqueue` then threw out of `start()`, + // which errors the stream and drops that queue — with a single queued + // chunk the consumer happens to drain it first and the bug hides. + const rows: Row[] = ['a', 'b', 'c', 'd', 'e'].map((text, i) => ({ + id: `chnk_0${i}`, + eof: false, + data: Buffer.from(`${text}\n`), + })); + rows.push( + { id: 'chnk_10', eof: true, data: Buffer.alloc(0) }, + // A retried terminal write: the frame appended again, the stream + // closed again. + { id: 'chnk_11', eof: false, data: Buffer.from('e\n') }, + { id: 'chnk_12', eof: true, data: Buffer.alloc(0) } + ); + const streamer = createStreamer( + { options: {} } as unknown as Pool, + fakeDrizzle(rows) + ); + try { + const stream = await streamer.streams.get('run_1', 'stream-1', 0); + await expect(drain(stream)).resolves.toBe('a\nb\nc\nd\ne\n'); + } finally { + await streamer.close(); + } + }); + + it.each([ + 5, 6, + ])('a start index at or past the data count (%i) closes at the first EOF instead of consuming it', async (startIndex) => { + // Offsets count data chunks (`getInfo` reports tailIndex = dataCount - 1), + // so the EOF marker must never be skipped as if it were data: before this + // guard the marker was swallowed, the post-EOF duplicate came back as + // live data, and with no later EOF the stream never closed. + const rows: Row[] = ['a', 'b', 'c', 'd', 'e'].map((text, i) => ({ + id: `chnk_0${i}`, + eof: false, + data: Buffer.from(`${text}\n`), + })); + rows.push( + { id: 'chnk_10', eof: true, data: Buffer.alloc(0) }, + { id: 'chnk_11', eof: false, data: Buffer.from('e\n') } + ); + const streamer = createStreamer( + { options: {} } as unknown as Pool, + fakeDrizzle(rows) + ); + try { + const stream = await streamer.streams.get( + 'run_1', + 'stream-1', + startIndex + ); + await expect(drain(stream)).resolves.toBe(''); + } finally { + await streamer.close(); + } + }, 5_000); + + it.each([ + [-1, 'e\n'], + [-2, 'd\ne\n'], + ])( + 'a negative start index (%i) counts data chunks up to the first EOF only', + async (startIndex, expected) => { + // The negative offset used to resolve against every row, including the + // ones after the first EOF, so `-1` on A…E + EOF + duplicate-E skipped + // the last valid chunk and returned nothing. + const rows: Row[] = ['a', 'b', 'c', 'd', 'e'].map((text, i) => ({ + id: `chnk_0${i}`, + eof: false, + data: Buffer.from(`${text}\n`), + })); + rows.push( + { id: 'chnk_10', eof: true, data: Buffer.alloc(0) }, + { id: 'chnk_11', eof: false, data: Buffer.from('e\n') } + ); + const streamer = createStreamer( + { options: {} } as unknown as Pool, + fakeDrizzle(rows) + ); + try { + const stream = await streamer.streams.get( + 'run_1', + 'stream-1', + startIndex + ); + await expect(drain(stream)).resolves.toBe(expected); + } finally { + await streamer.close(); + } + }, + 5_000 + ); +}); diff --git a/packages/world-postgres/src/streamer.ts b/packages/world-postgres/src/streamer.ts index a79c3ebd75..1a3c36fd68 100644 --- a/packages/world-postgres/src/streamer.ts +++ b/packages/world-postgres/src/streamer.ts @@ -361,18 +361,34 @@ export function createStreamer(pool: Pool, drizzle: Drizzle): PostgresStreamer { let lastChunkId = ''; let offset = startIndex ?? 0; let buffer = [] as StreamChunkEvent[] | null; + // Set by the first EOF marker. A producer that retries a + // terminal write (lost ACK, overlapping attempts) can append + // data and EOF rows after it; `enqueue`/`close` on the already + // closed controller would throw out of `start()`, erroring the + // stream and discarding every chunk still queued, so rows past + // the first EOF are ignored instead. + let closed = false; function enqueue(msg: { id: string; data: Uint8Array; eof: boolean; }) { + if (closed) { + return; + } + if (lastChunkId >= msg.id) { // already sent or out of order return; } - if (offset > 0) { + // The EOF marker is not a data chunk (`getInfo`'s tailIndex + // excludes it), so it never counts toward `offset`: a start + // index at or past the data count must still close the + // stream rather than consume the marker and then hang, or + // surface rows written after it. + if (offset > 0 && !msg.eof) { offset--; return; } @@ -381,6 +397,7 @@ export function createStreamer(pool: Pool, drizzle: Drizzle): PostgresStreamer { controller.enqueue(new Uint8Array(msg.data)); } if (msg.eof) { + closed = true; controller.close(); } lastChunkId = msg.id; @@ -408,13 +425,13 @@ export function createStreamer(pool: Pool, drizzle: Drizzle): PostgresStreamer { .where(and(eq(streams.streamId, name))) .orderBy(streams.chunkId); - // Resolve negative offset relative to the data chunk count - // (excluding the trailing EOF marker, if present) + // Resolve negative offset relative to the data chunk count: the + // rows before the first EOF marker. Rows after it (a retried + // terminal write) are ignored by `enqueue`, so they must not + // count here either. if (typeof offset === 'number' && offset < 0) { - const dataCount = - chunks.length > 0 && chunks[chunks.length - 1].eof - ? chunks.length - 1 - : chunks.length; + const firstEof = chunks.findIndex((chunk) => chunk.eof); + const dataCount = firstEof === -1 ? chunks.length : firstEof; offset = Math.max(0, dataCount + offset); }