diff --git a/js/src/utils.ts b/js/src/utils.ts index 060458e4..a7cd0df0 100644 --- a/js/src/utils.ts +++ b/js/src/utils.ts @@ -51,32 +51,54 @@ export function isConnectionClosedError(error: unknown): boolean { export async function* readLines(stream: ReadableStream) { const reader = stream.getReader() - let buffer = '' + const decoder = new TextDecoder() + // Text accumulated since the last newline, kept as an array of fragments: + // appending is O(1), so total work stays linear even when a single line + // spans many chunks (growing a string with `buffer += chunk` re-copies the + // whole buffer on every chunk, which is O(n²) and can OOM on large stdout). + const pending: string[] = [] try { while (true) { const { done, value } = await reader.read() if (value !== undefined) { - buffer += new TextDecoder().decode(value) + // { stream: true } keeps multi-byte UTF-8 sequences split across + // chunk boundaries intact instead of decoding them as U+FFFD. + const chunk = decoder.decode(value, { stream: true }) + let newlineIdx = chunk.indexOf('\n') + + if (newlineIdx === -1) { + if (chunk.length > 0) { + pending.push(chunk) + } + } else { + pending.push(chunk.slice(0, newlineIdx)) + yield pending.join('') + pending.length = 0 + + let start = newlineIdx + 1 + while ((newlineIdx = chunk.indexOf('\n', start)) !== -1) { + yield chunk.slice(start, newlineIdx) + start = newlineIdx + 1 + } + + if (start < chunk.length) { + pending.push(chunk.slice(start)) + } + } } if (done) { - if (buffer.length > 0) { - yield buffer + const trailing = decoder.decode() + if (trailing) { + pending.push(trailing) + } + if (pending.length > 0) { + yield pending.join('') } break } - - let newlineIdx = -1 - - do { - newlineIdx = buffer.indexOf('\n') - if (newlineIdx !== -1) { - yield buffer.slice(0, newlineIdx) - buffer = buffer.slice(newlineIdx + 1) - } - } while (newlineIdx !== -1) } } finally { reader.releaseLock() diff --git a/js/tests/utils.test.ts b/js/tests/utils.test.ts new file mode 100644 index 00000000..7c2e3c25 --- /dev/null +++ b/js/tests/utils.test.ts @@ -0,0 +1,96 @@ +import { expect, test } from 'vitest' + +import { readLines } from '../src/utils' + +function streamFromChunks(chunks: Uint8Array[]): ReadableStream { + return new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(chunk) + } + controller.close() + }, + }) +} + +function encodeChunks(chunks: string[]): Uint8Array[] { + const encoder = new TextEncoder() + return chunks.map((chunk) => encoder.encode(chunk)) +} + +async function collect(stream: ReadableStream): Promise { + const lines: string[] = [] + for await (const line of readLines(stream)) { + lines.push(line) + } + return lines +} + +test('yields lines split across a single chunk', async () => { + const lines = await collect(streamFromChunks(encodeChunks(['a\nb\nc\n']))) + expect(lines).toEqual(['a', 'b', 'c']) +}) + +test('yields trailing text without a final newline', async () => { + const lines = await collect(streamFromChunks(encodeChunks(['a\nb']))) + expect(lines).toEqual(['a', 'b']) +}) + +test('handles a line spanning multiple chunks', async () => { + const lines = await collect( + streamFromChunks(encodeChunks(['first pa', 'rt', ' of line\nsecond\n'])) + ) + expect(lines).toEqual(['first part of line', 'second']) +}) + +test('handles empty lines', async () => { + const lines = await collect(streamFromChunks(encodeChunks(['a\n\nb\n']))) + expect(lines).toEqual(['a', '', 'b']) +}) + +test('handles chunk boundaries around newlines', async () => { + const lines = await collect( + streamFromChunks(encodeChunks(['a', '\n', 'b\n', 'c'])) + ) + expect(lines).toEqual(['a', 'b', 'c']) +}) + +test('yields nothing for an empty stream', async () => { + const lines = await collect(streamFromChunks([])) + expect(lines).toEqual([]) +}) + +test('decodes multi-byte UTF-8 characters split across chunks', async () => { + // '🚀' is 4 bytes in UTF-8, so split it in the middle of the sequence + const bytes = new TextEncoder().encode('before 🚀 after\n') + const splitAt = bytes.indexOf(0xf0) + 2 + const lines = await collect( + streamFromChunks([bytes.slice(0, splitAt), bytes.slice(splitAt)]) + ) + expect(lines).toEqual(['before 🚀 after']) +}) + +test('decodes an incomplete multi-byte sequence at stream end', async () => { + // Stream ends mid-🚀: the decoder should flush a replacement character + // instead of silently dropping the bytes + const bytes = new TextEncoder().encode('abc🚀') + const lines = await collect( + streamFromChunks([bytes.slice(0, bytes.length - 2)]) + ) + expect(lines).toHaveLength(1) + expect(lines[0].startsWith('abc')).toBe(true) +}) + +test('handles large single-line output split into many chunks', async () => { + // A ~8 MB line delivered in 64 KiB chunks must complete without quadratic + // buffer rebuilds (this test hangs for minutes with `buffer += chunk`) + const chunkSize = 64 * 1024 + const chunkCount = 128 + const chunk = 'x'.repeat(chunkSize) + const chunks = encodeChunks([ + ...Array.from({ length: chunkCount }, () => chunk), + '\ntail\n', + ]) + const lines = await collect(streamFromChunks(chunks)) + expect(lines).toEqual(['x'.repeat(chunkSize * chunkCount), 'tail']) +})