diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 7dec8fee610..c925854d981 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -132,6 +132,14 @@ interface ArrowStreamContext { const readClients = new WeakMap(); +// gRPC status codes safe to resume from a row offset on ReadRows. +const GRPC_STATUS_INTERNAL = 13; +const GRPC_STATUS_UNAVAILABLE = 14; + +// How many times in a row a failed ReadRows call is reopened before the +// error is passed to the caller. Any rows read reset the count. +const MAX_STREAM_RETRIES = 3; + /** * Returns the first project ID from `candidates`, falling back to the * `BigQuery` client's project and ignoring the uninitialized `'{{projectId}}'` @@ -431,14 +439,46 @@ async function readDestinationSchema( return session?.arrowSchema?.serializedSchema || undefined; } +/** + * Returns `true` if a `ReadRows` gRPC stream error is transient and safe to + * resume from the next unread row offset. + * + * @private + */ +function isRetryableStreamError(err: Error & {code?: number}): boolean { + // DEADLINE_EXCEEDED isn't retried: it means the caller's deadline ran out. + if (err.code === GRPC_STATUS_UNAVAILABLE) { + return true; + } + // Long-lived HTTP/2 streams closed by the server with RST_STREAM surface as + // INTERNAL (code 13) in gRPC and are safe to reconnect at the next offset. + return ( + err.code === GRPC_STATUS_INTERNAL && + Boolean(err.message?.includes('RST_STREAM')) + ); +} + +/** + * Progress through a Storage Read API stream, kept across reconnects. + * + * @private + */ +interface ReadRowsProgress { + // The next unread row. + offset: number; + emittedBatch: boolean; + // The last Arrow schema seen, needed to decode later record batches. + serializedSchema?: Uint8Array; +} + /** * Streams Arrow `RecordBatch`es from a Storage Read API stream starting at - * `offset`, pausing `ReadRows` while the consumer is behind. + * `offset`. If `ReadRows` fails with a transient error, reopens it at the next + * unread row. * - * Resolves when `ReadRows` ends or the consumer destroys the stream, with - * whether any batch was emitted and the last Arrow schema seen, and rejects - * on a `ReadRows` or decode error. In every case, the listeners added here - * are removed, and `ReadRows` is cancelled unless it already finished. + * Resolves when the stream ends or the consumer destroys it, with whether any + * batch was emitted and the last Arrow schema seen. Rejects with the error + * once it isn't retryable or the retries run out. * * @private */ @@ -450,20 +490,60 @@ async function streamReadRows( initialSchema?: Uint8Array | string; }, ): Promise<{emittedBatch: boolean; serializedSchema?: Uint8Array}> { - const {bq, arrow, storage, stream} = ctx; const {readStream, offset, initialSchema} = options; + const progress: ReadRowsProgress = { + offset, + emittedBatch: false, + serializedSchema: initialSchema ? toUint8Array(initialSchema) : undefined, + }; + + let retries = 0; + while (!ctx.stream.destroyed) { + const attemptOffset = progress.offset; + try { + await readRowsOnce(ctx, readStream, progress); + break; + } catch (err) { + if (progress.offset > attemptOffset) { + retries = 0; + } + if ( + !isRetryableStreamError(err as Error & {code?: number}) || + retries >= MAX_STREAM_RETRIES + ) { + throw err; + } + retries++; + } + } + const {emittedBatch, serializedSchema} = progress; + return {emittedBatch, serializedSchema}; +} + +/** + * Opens one `ReadRows` call at `progress.offset` and pushes its Arrow + * `RecordBatch`es, pausing `ReadRows` while the consumer is behind. Updates + * `progress` for each batch, so a retry resumes at the next unread row. + * + * Resolves when `ReadRows` ends or the consumer destroys the stream, and + * rejects on a `ReadRows` or decode error. In every case, the listeners added + * here are removed, and `ReadRows` is cancelled unless it already finished. + * + * @private + */ +async function readRowsOnce( + ctx: ArrowStreamContext, + readStream: string, + progress: ReadRowsProgress, +): Promise { + const {bq, arrow, storage, stream} = ctx; return new Promise((resolve, reject) => { const grpcStream = getStorageReadClient(bq, storage).readRows({ readStream, - offset, + offset: progress.offset, }); - let serializedSchema: Uint8Array | undefined = initialSchema - ? toUint8Array(initialSchema) - : undefined; - let emittedBatch = false; - let settled = false; // Detaches every listener added here and settles the promise. Cancels // only this ReadRows call, keeping the shared BigQueryReadClient open. @@ -489,7 +569,7 @@ async function streamReadRows( if (opts.err) { reject(opts.err); } else { - resolve({emittedBatch, serializedSchema}); + resolve(); } }; @@ -502,7 +582,9 @@ async function streamReadRows( // and may also emit responses without record batch bytes (e.g., // heartbeats). if (resp.arrowSchema?.serializedSchema) { - serializedSchema = toUint8Array(resp.arrowSchema.serializedSchema); + progress.serializedSchema = toUint8Array( + resp.arrowSchema.serializedSchema, + ); } const rawBatch = resp.arrowRecordBatch?.serializedRecordBatch; const hasBatchBytes = @@ -514,13 +596,18 @@ async function streamReadRows( return; } try { - const batches = decodeArrowPayload(arrow, serializedSchema, rawBatch); + const batches = decodeArrowPayload( + arrow, + progress.serializedSchema, + rawBatch, + ); for (const batch of batches) { if (stream.destroyed) { settle({cancel: true}); return; } - emittedBatch = true; + progress.offset += batch.numRows; + progress.emittedBatch = true; if (!stream.push(batch)) { grpcStream.pause(); } @@ -579,7 +666,7 @@ async function streamDestinationTable( initialSchema: sessionSchema, }); // A stream that ends without rows still yields a 0-row RecordBatch. - if (!emittedBatch) { + if (!emittedBatch && !stream.destroyed) { pushBatches(stream, decodeArrowPayload(arrow, serializedSchema)); } } diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 0069967a658..81afb8f7f75 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -182,20 +182,47 @@ function stubJobGetMetadata( } /** - * Returns a fake ReadRows stream that emits `responses`, then ends, on the - * next tick. + * Returns a fake ReadRows stream that emits `responses` on the next tick, + * then ends, or fails with `err` once the consumer has read them. */ -function fakeReadRowsStream(responses: object[] = []): EventEmitter { +function fakeReadRowsStream( + responses: object[] = [], + err?: Error, +): EventEmitter { const grpcStream = new EventEmitter(); process.nextTick(() => { for (const resp of responses) { grpcStream.emit('data', resp); } - grpcStream.emit('end'); + if (err) { + setImmediate(() => grpcStream.emit('error', err)); + } else { + grpcStream.emit('end'); + } }); return grpcStream; } +/** Returns an error with a gRPC status `code`, like google-gax emits. */ +function grpcError(message: string, code: number): Error { + return Object.assign(new Error(message), {code}); +} + +/** Returns a ReadRowsResponse carrying a serialized Arrow table as bytes. */ +function toReadRowsResponse(payload: ReturnType) { + return { + arrowSchema: { + serializedSchema: Buffer.from(payload.serializedSchema, 'base64'), + }, + arrowRecordBatch: { + serializedRecordBatch: Buffer.from( + payload.serializedRecordBatch, + 'base64', + ), + }, + }; +} + /** A ReadRowsResponse carrying `SAMPLE_PAYLOAD` as bytes. */ const SAMPLE_READ_ROWS_RESPONSE = { arrowSchema: { @@ -1979,4 +2006,186 @@ describe('Arrow Support', () => { assert.strictEqual(readSessions, 0); }); }); + + describe('BigQuery#createQueryArrowStream (stream resumption)', () => { + const jobReference = { + projectId: 'test-project', + jobId: 'job-resume-1', + location: 'US', + }; + const readStream = + 'projects/test-project/locations/US/jobs/job-resume-1/streams/_default'; + let bq: BigQuery; + // The ReadRows streams opened, in order. + let grpcStreams: EventEmitter[]; + + beforeEach(() => { + bq = new BigQuery({projectId: 'test-project'}); + grpcStreams = []; + }); + + const track = (grpcStream: EventEmitter) => { + grpcStreams.push(grpcStream); + return grpcStream; + }; + + // Stubs jobs.query to return `inline` as its first page out of + // `totalRows`, so the remaining rows are read with ReadRows. + const stubFirstPage = ( + totalRows: number, + inline?: ReturnType, + inlineRows = 0, + ) => + stubJobsQuery(bq, { + res: { + jobComplete: true, + jobReference, + pageRowCount: String(inlineRows), + totalRows: String(totalRows), + arrowSchema: {serializedSchema: inline?.serializedSchema}, + arrowRecordBatch: { + serializedRecordBatch: inline?.serializedRecordBatch, + }, + }, + }); + + it('should resume ReadRows at offset + rowsReceived on UNAVAILABLE and RST_STREAM errors without gaps or duplicates', async () => { + stubFirstPage(4, serializeArrowTable(makeSampleTable([1], ['alice'])), 1); + const page2 = toReadRowsResponse( + serializeArrowTable(makeSampleTable([2, 3], ['bob', 'carol'])), + ); + const page3 = toReadRowsResponse( + serializeArrowTable(makeSampleTable([4], ['dave'])), + ); + // Each ReadRows call's responses and the error it fails with, if any. + const attempts: Array<[object[], Error?]> = [ + [[page2], grpcError('Service unavailable', 14)], + [[], grpcError('Connection reset', 14)], + [[], grpcError('Received RST_STREAM with code 2', 13)], + // Without a schema, so the batch decodes with the one sent earlier. + [[{arrowRecordBatch: page3.arrowRecordBatch}]], + ]; + const requests: ReadRowsRequest[] = []; + stubStorageReadClient(bq, { + readRows: req => { + const [responses, err] = attempts[requests.push(req) - 1]; + return track(fakeReadRowsStream(responses, err)); + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT id, name FROM tbl') + .toArray(); + + // Verify ReadRows started after the inline row, and every reopen + // resumed after the 2 rows read before the first error. + assert.deepStrictEqual(requests, [ + {readStream, offset: 1}, + {readStream, offset: 3}, + {readStream, offset: 3}, + {readStream, offset: 3}, + ]); + // Verify all rows arrived in order without gaps or duplicates. + const table = new arrow.Table(batches); + assert.deepStrictEqual( + Array.from(table.getChild('id')!.toArray()), + [1, 2, 3, 4], + ); + assert.deepStrictEqual(table.getChild('name')!.toArray(), [ + 'alice', + 'bob', + 'carol', + 'dave', + ]); + // Verify each attempt, failed or not, kept none of our listeners. + grpcStreams.forEach(assertReadRowsListenersRemoved); + }); + + it('should not retry DEADLINE_EXCEEDED, which means the caller deadline ran out', async () => { + stubFirstPage(2); + stubStorageReadClient(bq, { + readRows: () => + track(fakeReadRowsStream([], grpcError('Deadline exceeded', 4))), + }); + + await assert.rejects( + bq.createQueryArrowStream('SELECT id, name FROM tbl').toArray(), + /Deadline exceeded/, + ); + // Verify ReadRows wasn't reopened. + assert.strictEqual(grpcStreams.length, 1); + assertReadRowsListenersRemoved(grpcStreams[0]); + }); + + it('should surface non-retryable ReadRows errors after yielding earlier batches', async () => { + stubFirstPage(3, serializeArrowTable(makeSampleTable([1], ['alice'])), 1); + stubStorageReadClient(bq, { + readRows: () => + track( + fakeReadRowsStream( + [ + toReadRowsResponse( + serializeArrowTable(makeSampleTable([2], ['bob'])), + ), + ], + // INTERNAL without RST_STREAM isn't safe to resume. + grpcError('Corrupted stream state', 13), + ), + ), + }); + + const batches: arrow.RecordBatch[] = []; + await assert.rejects(async () => { + for await (const batch of bq.createQueryArrowStream( + 'SELECT id, name FROM tbl', + )) { + batches.push(batch); + } + }, /Corrupted stream state/); + + // Verify the inline and streamed batches arrived, and ReadRows wasn't + // reopened. + assert.strictEqual(batches.length, 2); + assert.strictEqual(grpcStreams.length, 1); + assertReadRowsListenersRemoved(grpcStreams[0]); + }); + + it('should stop retrying and surface the error when consecutive retryable errors exceed MAX_STREAM_RETRIES', async () => { + stubFirstPage(2); + stubStorageReadClient(bq, { + readRows: () => + track( + fakeReadRowsStream([], grpcError('Persistent UNAVAILABLE', 14)), + ), + }); + + await assert.rejects( + bq.createQueryArrowStream('SELECT id, name FROM tbl').toArray(), + /Persistent UNAVAILABLE/, + ); + + // Verify the first call and 3 retries each kept none of our listeners. + assert.strictEqual(grpcStreams.length, 4); + grpcStreams.forEach(assertReadRowsListenersRemoved); + }); + + it('should surface a synchronous readRows error when reopening ReadRows', async () => { + stubFirstPage(1); + let readRowsCalls = 0; + stubStorageReadClient(bq, { + readRows: () => { + if (++readRowsCalls > 1) { + throw new Error('readRows failed'); + } + return fakeReadRowsStream([], grpcError('Service unavailable', 14)); + }, + }); + + await assert.rejects( + bq.createQueryArrowStream('SELECT id, name FROM tbl').toArray(), + /readRows failed/, + ); + assert.strictEqual(readRowsCalls, 2); + }); + }); });