From f452d6f6bcdbe4a3e2bf3493940659aba8d5e838 Mon Sep 17 00:00:00 2001 From: Alka Trivedi Date: Thu, 27 Aug 2026 19:48:44 +0530 Subject: [PATCH 1/7] perf: optimize stream pipeline by eliminating events-intercept --- handwritten/spanner/package.json | 3 --- .../spanner/src/partial-result-stream.ts | 27 +++++++++---------- 2 files changed, 12 insertions(+), 18 deletions(-) diff --git a/handwritten/spanner/package.json b/handwritten/spanner/package.json index 717f82b15383..2f3b01b2b4e6 100644 --- a/handwritten/spanner/package.json +++ b/handwritten/spanner/package.json @@ -73,13 +73,11 @@ "big.js": "^7.0.0", "checkpoint-stream": "^0.1.2", "duplexify": "^4.1.3", - "events-intercept": "^2.0.0", "extend": "^3.0.2", "google-auth-library": "^10.0.0-rc.1", "google-gax": "5.0.6", "grpc-gcp": "^1.0.1", "lodash.snakecase": "^4.1.1", - "merge-stream": "^2.0.0", "p-queue": "^6.0.2", "protobufjs": "^7.4.0", "retry-request": "^8.0.0", @@ -98,7 +96,6 @@ "@types/extend": "^3.0.4", "@types/is": "^0.0.25", "@types/lodash.snakecase": "^4.1.9", - "@types/merge-stream": "^2.0.0", "@types/mocha": "^10.0.10", "@types/mv": "^2.1.4", "@types/ncp": "^2.0.8", diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index 2a1533c5419a..a75066e4f1d9 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -16,8 +16,6 @@ import {GrpcService} from './common-grpc/service'; import * as checkpointStream from 'checkpoint-stream'; -import * as eventsIntercept from 'events-intercept'; -import mergeStream = require('merge-stream'); import {common as p} from 'protobufjs'; import {Readable, Transform} from 'stream'; import * as streamEvents from 'stream-events'; @@ -578,17 +576,18 @@ export function partialResultStream( const maxQueued = 10; let lastResumeToken: ResumeToken; let lastRequestStream: Readable; + let errorListener: (err: grpc.ServiceError) => void; const startTime = Date.now(); const timeout = options?.gaxOptions?.timeout ?? Infinity; - // mergeStream allows multiple streams to be connected into one. This is good; + // requestsStream allows multiple streams to be connected into one. This is good; // if we need to retry a request and pipe more data to the user's stream. // We also add an additional stream that can be used to flush any remaining // items in the checkpoint stream that have been received, and that did not // contain a resume token. - const requestsStream = mergeStream(); + const requestsStream = new stream.PassThrough({objectMode: true}); const flushStream = new stream.PassThrough({objectMode: true}); - requestsStream.add(flushStream); + flushStream.pipe(requestsStream, {end: false}); const partialRSStream = new PartialResultStream(options); const userStream = streamEvents(partialRSStream); // We keep track of the number of PartialResultSets that did not include a @@ -626,7 +625,11 @@ export function partialResultStream( } lastRequestStream = requestFn(lastResumeToken); lastRequestStream.on('end', endListener); - requestsStream.add(lastRequestStream); + errorListener = (err: grpc.ServiceError) => { + setImmediate(() => retry(err)); + }; + lastRequestStream.on('error', errorListener); + lastRequestStream.pipe(requestsStream, {end: false}); }; const retry = (err: grpc.ServiceError): void => { @@ -659,6 +662,9 @@ export function partialResultStream( if (lastRequestStream) { lastRequestStream.removeListener('end', endListener); + if (errorListener) { + lastRequestStream.removeListener('error', errorListener); + } lastRequestStream.destroy(); } // Delay the retry until all the values that are already in the stream @@ -674,15 +680,6 @@ export function partialResultStream( }; userStream.once('reading', makeRequest); - eventsIntercept.patch(requestsStream); - - // need types for events-intercept - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (requestsStream as any).intercept('error', err => - // Retry __after__ all pending data has been processed to ensure that the - // checkpoint stream is reset at the correct position. - setImmediate(() => retry(err)), - ); return ( requestsStream From 700e088c9bb09701278ab01a925cfecfe2c665be Mon Sep 17 00:00:00 2001 From: Alka Trivedi Date: Fri, 28 Aug 2026 11:11:00 +0530 Subject: [PATCH 2/7] gemini review comments --- handwritten/spanner/src/partial-result-stream.ts | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index a75066e4f1d9..923c65ad3389 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -17,7 +17,7 @@ import {GrpcService} from './common-grpc/service'; import * as checkpointStream from 'checkpoint-stream'; import {common as p} from 'protobufjs'; -import {Readable, Transform} from 'stream'; +import { PassThrough, Readable, Transform } from 'stream'; import * as streamEvents from 'stream-events'; import {grpc, CallOptions} from 'google-gax'; import {DeadlineError, isRetryableInternalError} from './transaction-runner'; @@ -25,7 +25,6 @@ import {DeadlineError, isRetryableInternalError} from './transaction-runner'; import {codec, JSONOptions, Json, Field, Value} from './codec'; import {protos} from '@google-cloud/spanner-api'; import google = protos.google; -import * as stream from 'stream'; import {isDefined, isEmpty, isString} from './helper'; const originalDecode = codec.decode; @@ -585,9 +584,9 @@ export function partialResultStream( // We also add an additional stream that can be used to flush any remaining // items in the checkpoint stream that have been received, and that did not // contain a resume token. - const requestsStream = new stream.PassThrough({objectMode: true}); - const flushStream = new stream.PassThrough({objectMode: true}); - flushStream.pipe(requestsStream, {end: false}); + const requestsStream = new PassThrough({ objectMode: true }); + const flushStream = new PassThrough({ objectMode: true }); + flushStream.pipe(requestsStream); const partialRSStream = new PartialResultStream(options); const userStream = streamEvents(partialRSStream); // We keep track of the number of PartialResultSets that did not include a @@ -616,7 +615,6 @@ export function partialResultStream( // then push `null` to end the stream. flushStream.push({resumeToken: '_'}); flushStream.push(null); - requestsStream.end(); }); }; const makeRequest = (): void => { From a706f0e921729e2cbeb6ba327f1e1ff0fbea7fec Mon Sep 17 00:00:00 2001 From: Alka Trivedi Date: Fri, 28 Aug 2026 12:14:57 +0530 Subject: [PATCH 3/7] refactor --- handwritten/spanner/src/partial-result-stream.ts | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index 923c65ad3389..b78b09ff5efd 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -17,7 +17,7 @@ import {GrpcService} from './common-grpc/service'; import * as checkpointStream from 'checkpoint-stream'; import {common as p} from 'protobufjs'; -import { PassThrough, Readable, Transform } from 'stream'; +import {PassThrough, Readable, Transform} from 'stream'; import * as streamEvents from 'stream-events'; import {grpc, CallOptions} from 'google-gax'; import {DeadlineError, isRetryableInternalError} from './transaction-runner'; @@ -584,8 +584,8 @@ export function partialResultStream( // We also add an additional stream that can be used to flush any remaining // items in the checkpoint stream that have been received, and that did not // contain a resume token. - const requestsStream = new PassThrough({ objectMode: true }); - const flushStream = new PassThrough({ objectMode: true }); + const requestsStream = new PassThrough({objectMode: true}); + const flushStream = new PassThrough({objectMode: true}); flushStream.pipe(requestsStream); const partialRSStream = new PartialResultStream(options); const userStream = streamEvents(partialRSStream); @@ -660,9 +660,8 @@ export function partialResultStream( if (lastRequestStream) { lastRequestStream.removeListener('end', endListener); - if (errorListener) { - lastRequestStream.removeListener('error', errorListener); - } + lastRequestStream.removeAllListeners('error'); + lastRequestStream.on('error', () => {}); // Prevent unhandled exception crash lastRequestStream.destroy(); } // Delay the retry until all the values that are already in the stream From 21e9811322589badbef1d890fa6d4d504393bd36 Mon Sep 17 00:00:00 2001 From: Alka Trivedi Date: Thu, 3 Sep 2026 18:30:34 +0530 Subject: [PATCH 4/7] review comments --- .../spanner/src/partial-result-stream.ts | 102 ++++++++++++++++-- 1 file changed, 92 insertions(+), 10 deletions(-) diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index b78b09ff5efd..fb1b107788c5 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -15,7 +15,6 @@ */ import {GrpcService} from './common-grpc/service'; -import * as checkpointStream from 'checkpoint-stream'; import {common as p} from 'protobufjs'; import {PassThrough, Readable, Transform} from 'stream'; import * as streamEvents from 'stream-events'; @@ -567,6 +566,76 @@ export class PartialResultStream extends Transform implements ResultEvents { * @param {RowOptions} [options] Options for formatting rows. * @returns {PartialResultStream} */ +class CheckpointStream extends Transform { + private queue: google.spanner.v1.PartialResultSet[] = []; + private maxQueued: number; + private isCheckpointFn: ( + chunk: google.spanner.v1.PartialResultSet, + ) => boolean; + + constructor(options: { + maxQueued?: number; + isCheckpointFn: (chunk: google.spanner.v1.PartialResultSet) => boolean; + }) { + super({objectMode: true}); + this.maxQueued = options.maxQueued ?? 10; + this.isCheckpointFn = options.isCheckpointFn; + } + + _transform( + chunk: google.spanner.v1.PartialResultSet, + enc: string, + callback: () => void, + ): void { + this.queue.push(chunk); + const isCheckpoint = this.isCheckpointFn(chunk); + let shouldFlush = false; + if (isCheckpoint) { + this.emit('checkpoint', chunk); + shouldFlush = true; + } else if (this.queue.length > this.maxQueued) { + shouldFlush = true; + } + + if (!shouldFlush) { + return callback(); + } + + this._flushQueue(callback); + } + + private _flushQueue(callback: () => void): void { + const loopyloop = () => { + if (this.destroyed) { + return callback(); + } + + if (this.queue.length > 0) { + this.push(this.queue.shift()); + setImmediate(loopyloop); + } else { + callback(); + } + }; + + loopyloop(); + } + + flushAndDestroy(err: Error): void { + this._flushQueue(() => { + this.destroy(err); + }); + } + + reset(): void { + this.queue = []; + } + + _flush(callback: () => void): void { + this._flushQueue(callback); + } +} + export function partialResultStream( requestFn: RequestFunction, options?: RowOptions, @@ -593,7 +662,7 @@ export function partialResultStream( // resume token, as that is an indication whether it is safe to retry the // stream halfway. let withoutCheckpointCount = 0; - const batchAndSplitOnTokenStream = checkpointStream.obj({ + const batchAndSplitOnTokenStream = new CheckpointStream({ maxQueued, isCheckpointFn: (chunk: google.spanner.v1.PartialResultSet): boolean => { const withCheckpoint = _hasResumeToken(chunk); @@ -617,6 +686,17 @@ export function partialResultStream( flushStream.push(null); }); }; + + const destroyRequestStream = (): void => { + if (lastRequestStream) { + lastRequestStream.removeListener('end', endListener); + lastRequestStream.removeAllListeners('error'); + lastRequestStream.on('error', () => {}); + lastRequestStream.unpipe(requestsStream); + lastRequestStream.destroy(); + } + }; + const makeRequest = (): void => { if (isDefined(lastResumeToken) && lastResumeToken.length > 0) { partialRSStream._resetPendingValues(); @@ -624,6 +704,7 @@ export function partialResultStream( lastRequestStream = requestFn(lastResumeToken); lastRequestStream.on('end', endListener); errorListener = (err: grpc.ServiceError) => { + destroyRequestStream(); setImmediate(() => retry(err)); }; lastRequestStream.on('error', errorListener); @@ -631,13 +712,14 @@ export function partialResultStream( }; const retry = (err: grpc.ServiceError): void => { + destroyRequestStream(); const elapsed = Date.now() - startTime; if (elapsed >= timeout) { // The timeout has reached so this will flush any rows the // checkpoint stream has queued. After that, we will destroy the // user's stream with the Deadline exceeded error. setImmediate(() => - batchAndSplitOnTokenStream.destroy(new DeadlineError(err)), + batchAndSplitOnTokenStream.flushAndDestroy(new DeadlineError(err)), ); return; } @@ -654,16 +736,10 @@ export function partialResultStream( // This is not a retryable error so this will flush any rows the // checkpoint stream has queued. After that, we will destroy the // user's stream with the same error. - setImmediate(() => batchAndSplitOnTokenStream.destroy(err)); + setImmediate(() => batchAndSplitOnTokenStream.flushAndDestroy(err)); return; } - if (lastRequestStream) { - lastRequestStream.removeListener('end', endListener); - lastRequestStream.removeAllListeners('error'); - lastRequestStream.on('error', () => {}); // Prevent unhandled exception crash - lastRequestStream.destroy(); - } // Delay the retry until all the values that are already in the stream // pipeline have been handled. This ensures that the checkpoint stream is // reset to the correct point. Calling .reset() directly here could cause @@ -677,6 +753,12 @@ export function partialResultStream( }; userStream.once('reading', makeRequest); + userStream.once('close', () => { + destroyRequestStream(); + requestsStream.destroy(); + flushStream.destroy(); + batchAndSplitOnTokenStream.destroy(); + }); return ( requestsStream From 152d1abb90b337e7ecf7243bf68988b7a294b79b Mon Sep 17 00:00:00 2001 From: Alka Trivedi Date: Thu, 3 Sep 2026 23:00:23 +0530 Subject: [PATCH 5/7] test: PartialResultStream --- handwritten/spanner/package.json | 1 - .../spanner/src/partial-result-stream.ts | 8 +- .../spanner/test/partial-result-stream.ts | 198 +++++++++++++++--- 3 files changed, 173 insertions(+), 34 deletions(-) diff --git a/handwritten/spanner/package.json b/handwritten/spanner/package.json index 2f3b01b2b4e6..33c01cdf35dd 100644 --- a/handwritten/spanner/package.json +++ b/handwritten/spanner/package.json @@ -71,7 +71,6 @@ "@types/big.js": "^6.2.2", "@types/stack-trace": "^0.0.33", "big.js": "^7.0.0", - "checkpoint-stream": "^0.1.2", "duplexify": "^4.1.3", "extend": "^3.0.2", "google-auth-library": "^10.0.0-rc.1", diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index fb1b107788c5..28f340f8a824 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -605,20 +605,20 @@ class CheckpointStream extends Transform { } private _flushQueue(callback: () => void): void { - const loopyloop = () => { + const loop = () => { if (this.destroyed) { return callback(); } if (this.queue.length > 0) { this.push(this.queue.shift()); - setImmediate(loopyloop); + setImmediate(loop); } else { callback(); } }; - loopyloop(); + loop(); } flushAndDestroy(err: Error): void { @@ -691,7 +691,7 @@ export function partialResultStream( if (lastRequestStream) { lastRequestStream.removeListener('end', endListener); lastRequestStream.removeAllListeners('error'); - lastRequestStream.on('error', () => {}); + lastRequestStream.on('error', () => { }); lastRequestStream.unpipe(requestsStream); lastRequestStream.destroy(); } diff --git a/handwritten/spanner/test/partial-result-stream.ts b/handwritten/spanner/test/partial-result-stream.ts index 38700a57d245..603f399be256 100644 --- a/handwritten/spanner/test/partial-result-stream.ts +++ b/handwritten/spanner/test/partial-result-stream.ts @@ -17,7 +17,6 @@ import * as assert from 'assert'; import {before, beforeEach, afterEach, describe, it} from 'mocha'; // eslint-disable-next-line @typescript-eslint/no-var-requires -const checkpointStream = require('checkpoint-stream'); // eslint-disable-next-line @typescript-eslint/no-var-requires const concat = require('concat-stream'); import * as proxyquire from 'proxyquire'; @@ -102,7 +101,6 @@ describe('PartialResultStream', () => { before(() => { const prsExports = proxyquire('../src/partial-result-stream.js', { - 'checkpoint-stream': checkpointStream, stream: {Transform}, './codec': {codec}, }); @@ -636,12 +634,7 @@ describe('PartialResultStream', () => { // This test will emit two rows total: // - UNAVAILABLE error (should retry) // - Two rows - // - Confirm all rows were received. - const fakeCheckpointStream = through.obj(); // eslint-disable-next-line @typescript-eslint/no-explicit-any - (fakeCheckpointStream as any).reset = () => {}; - - sandbox.stub(checkpointStream, 'obj').returns(fakeCheckpointStream); const firstFakeRequestStream = through.obj(); const secondFakeRequestStream = through.obj(); @@ -668,9 +661,7 @@ describe('PartialResultStream', () => { setTimeout(() => { secondFakeRequestStream.push(RESULT_WITH_TOKEN); - fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN); secondFakeRequestStream.push(RESULT_WITH_TOKEN); - fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN); secondFakeRequestStream.end(); }, 500); @@ -689,12 +680,6 @@ describe('PartialResultStream', () => { }); it('should get Deadline exceeded error if timeout has reached', done => { - const fakeCheckpointStream = through.obj(); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (fakeCheckpointStream as any).reset = () => {}; - - sandbox.stub(checkpointStream, 'obj').returns(fakeCheckpointStream); - const firstFakeRequestStream = through.obj(); const requestFnStub = sandbox.stub(); @@ -726,11 +711,6 @@ describe('PartialResultStream', () => { // - Error event (should retry) // - Two rows // - Confirm all rows were received. - const fakeCheckpointStream = through.obj(); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (fakeCheckpointStream as any).reset = () => {}; - sandbox.stub(checkpointStream, 'obj').returns(fakeCheckpointStream); - const firstFakeRequestStream = through.obj(); const secondFakeRequestStream = through.obj(); @@ -739,9 +719,7 @@ describe('PartialResultStream', () => { requestFnStub.onCall(0).callsFake(() => { setTimeout(() => { firstFakeRequestStream.push(RESULT_WITH_TOKEN); - fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN); firstFakeRequestStream.push(RESULT_WITH_TOKEN); - fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN); setTimeout(() => { // This causes a new request stream to be created. @@ -760,9 +738,7 @@ describe('PartialResultStream', () => { setTimeout(() => { secondFakeRequestStream.push(RESULT_WITH_TOKEN); - fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN); secondFakeRequestStream.push(RESULT_WITH_TOKEN); - fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN); secondFakeRequestStream.end(); }, 500); @@ -782,9 +758,6 @@ describe('PartialResultStream', () => { it('should emit non-retryable error', done => { // This test will emit two rows and then an error. - const fakeCheckpointStream = through.obj(); - sandbox.stub(checkpointStream, 'obj').returns(fakeCheckpointStream); - const fakeRequestStream = through.obj(); const requestFnStub = sandbox.stub(); @@ -792,9 +765,7 @@ describe('PartialResultStream', () => { requestFnStub.onCall(0).callsFake(() => { setTimeout(() => { fakeRequestStream.push(RESULT_WITH_TOKEN); - fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN); fakeRequestStream.push(RESULT_WITH_TOKEN); - fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN); setTimeout(() => { fakeRequestStream.emit('error', { @@ -839,6 +810,175 @@ describe('PartialResultStream', () => { fakeRequestStream.push(RESULT); fakeRequestStream.destroy(error); }); + + it('should successfully retry when the failed stream emits an error followed by end', done => { + const firstStream = through.obj(); + const secondStream = through.obj(); + const requestFnStub = sandbox.stub(); + + // First request fails with UNAVAILABLE and immediately ends + requestFnStub.onCall(0).callsFake(() => { + setImmediate(() => { + firstStream.emit('error', { + code: grpc.status.UNAVAILABLE, + message: 'Unavailable', + } as grpc.ServiceError); + firstStream.end(); + }); + return firstStream; + }); + + // Retried request succeeds and delivers data + requestFnStub.onCall(1).callsFake(() => { + setImmediate(() => { + secondStream.push(RESULT_WITH_TOKEN); + secondStream.end(); + }); + return secondStream; + }); + + const receivedRows: Row[] = []; + partialResultStream(requestFnStub) + .on('data', row => receivedRows.push(row)) + .on('error', done) + .on('end', () => { + try { + assert.strictEqual( + requestFnStub.callCount, + 2, + 'Should have retried once', + ); + assert.strictEqual( + receivedRows.length, + 1, + 'Should receive data from retried stream', + ); + done(); + } catch (e) { + done(e); + } + }); + }); + + it('should only spawn a single retry when multiple errors are emitted in rapid succession', done => { + const firstStream = through.obj(); + const secondStream = through.obj(); + const requestFnStub = sandbox.stub(); + + firstStream.on('error', () => {}); // Prevent unhandled exception in test runner + + // First request emits two error events synchronously + requestFnStub.onCall(0).callsFake(() => { + setImmediate(() => { + const err = { + code: grpc.status.UNAVAILABLE, + message: 'Unavailable', + } as grpc.ServiceError; + firstStream.emit('error', err); + firstStream.emit('error', err); + }); + return firstStream; + }); + + // Second request succeeds + requestFnStub.onCall(1).callsFake(() => { + setImmediate(() => { + secondStream.push(RESULT_WITH_TOKEN); + secondStream.end(); + }); + return secondStream; + }); + + partialResultStream(requestFnStub) + .on('error', done) + .pipe( + concat(rows => { + try { + assert.strictEqual( + requestFnStub.callCount, + 2, + 'Should only trigger one retry request', + ); + assert.strictEqual(rows.length, 1); + done(); + } catch (e) { + done(e); + } + }), + ); + }); + + it('should destroy the request stream and detach listeners on non-retryable errors', done => { + const fakeStream = through.obj(); + const destroySpy = sandbox.spy(fakeStream, 'destroy'); + + const requestFnStub = sandbox.stub().callsFake(() => { + setImmediate(() => { + fakeStream.emit('error', { + code: grpc.status.INVALID_ARGUMENT, + message: 'Invalid query argument.', + } as grpc.ServiceError); + }); + return fakeStream; + }); + + partialResultStream(requestFnStub) + .on('data', () => {}) + .on('error', err => { + try { + assert.strictEqual(err.code, grpc.status.INVALID_ARGUMENT); + assert.strictEqual( + destroySpy.called, + true, + 'Request stream should be destroyed on non-retryable error', + ); + assert.strictEqual( + fakeStream.listenerCount('end'), + 0, + 'endListener should be removed', + ); + assert.strictEqual( + fakeStream.listenerCount('error'), + 1, + 'Should have 1 dummy listener to swallow late errors', + ); + done(); + } catch (e) { + done(e); + } + }); + }); + + it('should destroy the underlying request stream when the user destroys the returned stream', done => { + const fakeStream = through.obj(); + const destroySpy = sandbox.spy(fakeStream, 'destroy'); + + const requestFnStub = sandbox.stub().returns(fakeStream); + + const stream = partialResultStream(requestFnStub); + + // Read first row and immediately destroy stream + stream.on('data', () => { + stream.destroy(); + }); + + stream.on('close', () => { + setImmediate(() => { + try { + assert.strictEqual( + destroySpy.called, + true, + 'Underlying request stream must be destroyed when user cancels the stream', + ); + done(); + } catch (e) { + done(e); + } + }); + }); + + fakeStream.push(RESULT_WITH_TOKEN); + }); }); }); From 6289364bdf387d6d67dc07b17d6d8f18eebcf86e Mon Sep 17 00:00:00 2001 From: Alka Trivedi Date: Fri, 4 Sep 2026 09:08:24 +0530 Subject: [PATCH 6/7] doc: class CheckpointStream --- .../spanner/src/partial-result-stream.ts | 65 +++++++++++++++---- 1 file changed, 52 insertions(+), 13 deletions(-) diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index 28f340f8a824..bacbef6f775b 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -551,20 +551,17 @@ export class PartialResultStream extends Transform implements ResultEvents { } /** - * Rows returned from queries may be chunked, requiring them to be stitched - * together. This function returns a stream that will properly assemble these - * rows, as well as retry after an error. Rows are only emitted if they hit a - * "checkpoint", which is when a `resumeToken` is returned from the API. Without - * that token, it's unsafe for the query to be retried, as we wouldn't want to - * emit the same data multiple times. + * A custom Transform stream that buffers PartialResultSet chunks and flushes them + * asynchronously to prevent blocking the event loop. * - * @private + * It holds chunks in a queue until a "checkpoint" is reached (as determined by + * `isCheckpointFn`) or until the queue exceeds `maxQueued` items. * - * @param {RequestFunction} requestFn The function that makes an API request. It - * will receive one argument, `resumeToken`, which should be used however is - * necessary to send to the API for additional requests. - * @param {RowOptions} [options] Options for formatting rows. - * @returns {PartialResultStream} + * This matches the legacy behavior of buffering chunks and flushing them + * asynchronously using `setImmediate` to yield control back to the event loop, + * preventing long-running synchronous loops from blocking other processing. + * + * @private */ class CheckpointStream extends Transform { private queue: google.spanner.v1.PartialResultSet[] = []; @@ -582,6 +579,13 @@ class CheckpointStream extends Transform { this.isCheckpointFn = options.isCheckpointFn; } + /** + * Buffers chunks and flushes queue asynchronously on checkpoints or when max limit is reached. + * + * @param {google.spanner.v1.PartialResultSet} chunk The chunk to transform. + * @param {string} enc Encoding (unused). + * @param {Function} callback Callback to signal completion of transformation. + */ _transform( chunk: google.spanner.v1.PartialResultSet, enc: string, @@ -604,6 +608,12 @@ class CheckpointStream extends Transform { this._flushQueue(callback); } + /** + * Flushes queued chunks asynchronously using `setImmediate` to avoid blocking the event loop. + * + * @private + * @param {Function} callback Callback to call when all chunks are flushed. + */ private _flushQueue(callback: () => void): void { const loop = () => { if (this.destroyed) { @@ -621,21 +631,50 @@ class CheckpointStream extends Transform { loop(); } + /** + * Flushes remaining queued chunks before destroying the stream with the provided error. + * + * @param {Error} err The error to destroy the stream with. + */ flushAndDestroy(err: Error): void { this._flushQueue(() => { this.destroy(err); }); } + /** + * Clears the queue without flushing, useful when retrying and discarding partial data. + */ reset(): void { this.queue = []; } + /** + * Flushes all remaining queued chunks when the stream ends. + * + * @param {Function} callback Callback to call when flushing is complete. + */ _flush(callback: () => void): void { this._flushQueue(callback); } } +/** + * Rows returned from queries may be chunked, requiring them to be stitched + * together. This function returns a stream that will properly assemble these + * rows, as well as retry after an error. Rows are only emitted if they hit a + * "checkpoint", which is when a `resumeToken` is returned from the API. Without + * that token, it's unsafe for the query to be retried, as we wouldn't want to + * emit the same data multiple times. + * + * @private + * + * @param {RequestFunction} requestFn The function that makes an API request. It + * will receive one argument, `resumeToken`, which should be used however is + * necessary to send to the API for additional requests. + * @param {RowOptions} [options] Options for formatting rows. + * @returns {PartialResultStream} + */ export function partialResultStream( requestFn: RequestFunction, options?: RowOptions, @@ -691,7 +730,7 @@ export function partialResultStream( if (lastRequestStream) { lastRequestStream.removeListener('end', endListener); lastRequestStream.removeAllListeners('error'); - lastRequestStream.on('error', () => { }); + lastRequestStream.on('error', () => {}); lastRequestStream.unpipe(requestsStream); lastRequestStream.destroy(); } From d1a1f7f46b10bd154301ea02aeee89fe8f970157 Mon Sep 17 00:00:00 2001 From: Alka Trivedi Date: Sat, 5 Sep 2026 14:53:53 +0530 Subject: [PATCH 7/7] review comment --- .../spanner/src/partial-result-stream.ts | 35 ++++------- .../spanner/test/partial-result-stream.ts | 61 +++++++++++++++++++ 2 files changed, 73 insertions(+), 23 deletions(-) diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index bacbef6f775b..f43f0384bdb5 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -580,7 +580,7 @@ class CheckpointStream extends Transform { } /** - * Buffers chunks and flushes queue asynchronously on checkpoints or when max limit is reached. + * Buffers chunks and flushes queue synchronously on checkpoints or when max limit is reached. * * @param {google.spanner.v1.PartialResultSet} chunk The chunk to transform. * @param {string} enc Encoding (unused). @@ -605,30 +605,19 @@ class CheckpointStream extends Transform { return callback(); } - this._flushQueue(callback); + this._flushQueue(); + callback(); } /** - * Flushes queued chunks asynchronously using `setImmediate` to avoid blocking the event loop. + * Flushes queued chunks synchronously to prevent state races on retry/reset. * * @private - * @param {Function} callback Callback to call when all chunks are flushed. */ - private _flushQueue(callback: () => void): void { - const loop = () => { - if (this.destroyed) { - return callback(); - } - - if (this.queue.length > 0) { - this.push(this.queue.shift()); - setImmediate(loop); - } else { - callback(); - } - }; - - loop(); + private _flushQueue(): void { + while (this.queue.length > 0 && !this.destroyed) { + this.push(this.queue.shift()); + } } /** @@ -637,9 +626,8 @@ class CheckpointStream extends Transform { * @param {Error} err The error to destroy the stream with. */ flushAndDestroy(err: Error): void { - this._flushQueue(() => { - this.destroy(err); - }); + this._flushQueue(); + this.destroy(err); } /** @@ -655,7 +643,8 @@ class CheckpointStream extends Transform { * @param {Function} callback Callback to call when flushing is complete. */ _flush(callback: () => void): void { - this._flushQueue(callback); + this._flushQueue(); + callback(); } } diff --git a/handwritten/spanner/test/partial-result-stream.ts b/handwritten/spanner/test/partial-result-stream.ts index 603f399be256..1aafbfca6a85 100644 --- a/handwritten/spanner/test/partial-result-stream.ts +++ b/handwritten/spanner/test/partial-result-stream.ts @@ -979,6 +979,67 @@ describe('PartialResultStream', () => { fakeStream.push(RESULT_WITH_TOKEN); }); + + it('should not drop buffered checkpointed chunks when a retry occurs during flush', done => { + const firstStream = through.obj(); + const secondStream = through.obj(); + const requestFnStub = sandbox.stub(); + + const token1 = 'token1'; + // Chunks 1 to 3 have no token; Chunk 4 has token1 + const chunk1 = Object.assign({}, RESULT, {resumeToken: ''}); + const chunk2 = Object.assign({}, RESULT, {resumeToken: ''}); + const chunk3 = Object.assign({}, RESULT, {resumeToken: ''}); + const chunk4 = Object.assign({}, RESULT, {resumeToken: token1}); + // Chunk 5 is returned after retry + const chunk5 = Object.assign({}, RESULT, {resumeToken: 'token2'}); + + requestFnStub.onCall(0).callsFake(() => { + setImmediate(() => { + firstStream.push(chunk1); + firstStream.push(chunk2); + firstStream.push(chunk3); + firstStream.push(chunk4); // Checkpoint hit: queue has 4 items + // Simulate network blip immediately after sending chunk4 + firstStream.emit('error', { + code: grpc.status.UNAVAILABLE, + message: 'Unavailable', + } as grpc.ServiceError); + }); + return firstStream; + }); + + requestFnStub.onCall(1).callsFake(resumeToken => { + try { + assert.strictEqual(resumeToken, token1, 'Must resume from token1'); + } catch (e) { + done(e); + } + setImmediate(() => { + secondStream.push(chunk5); + secondStream.end(); + }); + return secondStream; + }); + + const receivedRows: Row[] = []; + partialResultStream(requestFnStub) + .on('data', (row: any) => receivedRows.push(row)) + .on('error', done) + .on('end', () => { + try { + // Must receive all 4 rows from the checkpointed batch + 1 from retry = 5 total + assert.strictEqual( + receivedRows.length, + 5, + 'All checkpointed rows must be delivered without being dropped by retry reset()', + ); + done(); + } catch (e) { + done(e); + } + }); + }); }); });