From cbf8ff34c4d23b3d832742fb4073938c6a96bb41 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Tue, 8 Sep 2026 13:49:16 +0200 Subject: [PATCH] perf(spanner): complete query stream immediately on PartialResultSet.last Optimizes query streaming completion and server-side RPC status by leveraging PartialResultSet.last: 1. Latency: Completes user stream immediately upon receiving chunk.last = true via this.push(null) and next(), bypassing the wait for gRPC trailers, the gRPC 'end' event, and secondary flush ticks. 2. Graceful Drain: Detaches the underlying gRPC request stream from the pipeline and calls resume() when chunk.last is true. This allows trailing gRPC metadata and EOF to drain asynchronously in the background without canceling the call, preventing successful queries from being reported as CANCELLED in Cloud Monitoring. 3. Natural Lifecycle: Signals requestsStream.end() on chunk.last so EOF cascades through CheckpointStream and PartialResultStream naturally, allowing 'finish', 'close', and stream.finished() to resolve cleanly without ad-hoc stream destruction. 4. Error Safety: Symmetrically detaches error/end listeners and attaches a no-op error handler during background drain to prevent late socket resets from triggering retries or unhandled exceptions. --- .../spanner/src/partial-result-stream.ts | 45 +- .../spanner/test/partial-result-stream.ts | 404 +++++++++++++++++- 2 files changed, 441 insertions(+), 8 deletions(-) diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index 18f492e577c0..720620c61e27 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -275,6 +275,11 @@ export class PartialResultStream extends Transform implements ResultEvents { if (chunk.last) { this.push(null); + // Calling next() notifies Node's stream machinery that processing of this + // chunk is complete on the Writable side of this Transform stream. + // This is a local, synchronous callback to Node's internal buffer; it does + // not block the event loop or wait for upstream network I/O or gRPC trailers. + next(); return; } @@ -667,7 +672,7 @@ export function partialResultStream( const retryableCodes = [grpc.status.UNAVAILABLE]; const maxQueued = 10; let lastResumeToken: ResumeToken; - let lastRequestStream: Readable; + let lastRequestStream: Readable | undefined; let errorListener: (err: grpc.ServiceError) => void; const startTime = Date.now(); const timeout = options?.gaxOptions?.timeout ?? Infinity; @@ -686,10 +691,16 @@ export function partialResultStream( // resume token, as that is an indication whether it is safe to retry the // stream halfway. let withoutCheckpointCount = 0; + let receivedLast = false; const batchAndSplitOnTokenStream = new CheckpointStream({ maxQueued, isCheckpointFn: (chunk: google.spanner.v1.PartialResultSet): boolean => { - const withCheckpoint = _hasResumeToken(chunk); + if (chunk.last) { + receivedLast = true; + destroyRequestStream(); + requestsStream.end(); + } + const withCheckpoint = _hasResumeToken(chunk) || Boolean(chunk.last); if (withCheckpoint) { withoutCheckpointCount = 0; } else { @@ -702,7 +713,13 @@ export function partialResultStream( // This listener ensures that the last request that executed successfully // after one or more retries will end the requestsStream. const endListener = () => { + if (receivedLast) { + return; + } setImmediate(() => { + if (receivedLast) { + return; + } // Push a fake PartialResultSet without any values but with a resume token // into the stream to ensure that the checkpoint stream is emptied, and // then push `null` to end the stream. @@ -713,11 +730,22 @@ export function partialResultStream( const destroyRequestStream = (): void => { if (lastRequestStream) { - lastRequestStream.removeListener('end', endListener); - lastRequestStream.removeAllListeners('error'); - lastRequestStream.on('error', () => {}); - lastRequestStream.unpipe(requestsStream); - lastRequestStream.destroy(); + const streamToClean = lastRequestStream; + lastRequestStream = undefined; + streamToClean.removeListener('end', endListener); + if (errorListener) { + streamToClean.removeListener('error', errorListener); + } + streamToClean.on('error', () => {}); + streamToClean.unpipe(requestsStream); + if (receivedLast) { + // Query completed successfully. Do not cancel the gRPC call; allow it + // to drain remaining trailers/EOF in the background so it is not marked + // CANCELLED by Spanner or Cloud Monitoring. + streamToClean.resume(); + } else { + streamToClean.destroy(); + } } }; @@ -728,6 +756,9 @@ export function partialResultStream( lastRequestStream = requestFn(lastResumeToken); lastRequestStream.on('end', endListener); errorListener = (err: grpc.ServiceError) => { + if (receivedLast) { + return; + } destroyRequestStream(); setImmediate(() => retry(err)); }; diff --git a/handwritten/spanner/test/partial-result-stream.ts b/handwritten/spanner/test/partial-result-stream.ts index 1aafbfca6a85..322588dc2cb7 100644 --- a/handwritten/spanner/test/partial-result-stream.ts +++ b/handwritten/spanner/test/partial-result-stream.ts @@ -21,7 +21,7 @@ import {before, beforeEach, afterEach, describe, it} from 'mocha'; const concat = require('concat-stream'); import * as proxyquire from 'proxyquire'; import * as sinon from 'sinon'; -import {Transform} from 'stream'; +import {Transform, finished} from 'stream'; import * as through from 'through2'; import {codec} from '../src/codec'; @@ -1040,6 +1040,408 @@ describe('PartialResultStream', () => { } }); }); + + it('should immediately flush rows and emit end when chunk.last is true without waiting for gRPC stream end', done => { + let dataEmitted = false; + + stream + .on('data', () => { + dataEmitted = true; + }) + .on('end', () => { + try { + assert.strictEqual(dataEmitted, true); + // Underlying fakeRequestStream has NOT received null/end yet (simulating pending trailers) + // and MUST NOT be destroyed! + assert.strictEqual(fakeRequestStream.destroyed, false); + // Now simulate trailers arriving asynchronously + fakeRequestStream.push(null); + done(); + } catch (error) { + done(error); + } + }) + .on('error', done); + + // Push chunk with last: true but NO resumeToken and without ending fakeRequestStream + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: true, + }), + ); + }); + + it('should not destroy request stream when user stream closes after chunk.last is true', done => { + stream + .on('data', () => {}) + .on('end', () => { + stream.destroy(); + }) + .on('close', () => { + setImmediate(() => { + try { + // The request stream must remain open to drain trailers asynchronously + assert.strictEqual( + fakeRequestStream.destroyed, + false, + 'Request stream should not be destroyed when chunk.last is true', + ); + fakeRequestStream.push(null); + done(); + } catch (error) { + done(error); + } + }); + }) + .on('error', done); + + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: true, + }), + ); + }); + + it('should clean up request stream when error occurs after chunk.last without retrying', done => { + let dataEmitted = false; + + stream + .on('data', () => { + dataEmitted = true; + }) + .on('end', () => { + try { + assert.strictEqual(dataEmitted, true); + // Simulate transport error occurring while trailers were in flight + fakeRequestStream.emit('error', new Error('Late transport error')); + setImmediate(() => { + // Should not throw unhandled error and should have detached listeners + done(); + }); + } catch (error) { + done(error); + } + }) + .on('error', err => { + done( + new Error( + `Stream should not emit error after chunk.last: ${err.message}`, + ), + ); + }); + + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: true, + }), + ); + }); + + it('should automatically clean up request stream on stream end without explicit destroy()', done => { + const resumeSpy = sandbox.spy(fakeRequestStream, 'resume'); + + stream + .on('data', () => {}) + .on('end', () => { + setImmediate(() => { + try { + // resume() must be called to drain background trailers without requiring stream.destroy() + assert.strictEqual( + resumeSpy.called, + true, + 'resume() should be called on fakeRequestStream', + ); + fakeRequestStream.push(null); + done(); + } catch (error) { + done(error); + } + }); + }) + .on('error', done); + + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: true, + }), + ); + }); + + it('should emit finish and close, and resolve stream.finished() naturally on chunk.last without manual destroy', done => { + let finishedCalled = false; + let closeEmitted = false; + let finishEmitted = false; + + finished(stream, err => { + try { + assert.ifError(err); + finishedCalled = true; + if (closeEmitted && finishEmitted) { + done(); + } + } catch (error) { + done(error); + } + }); + + stream + .on('data', () => {}) + .on('finish', () => { + finishEmitted = true; + }) + .on('close', () => { + closeEmitted = true; + if (finishedCalled && finishEmitted) { + done(); + } + }) + .on('error', done); + + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: true, + }), + ); + }); + + it('should correctly handle zero-row result set with chunk.last', done => { + let dataEmitted = false; + + stream + .on('data', () => { + dataEmitted = true; + }) + .on('end', () => { + try { + assert.strictEqual( + dataEmitted, + false, + 'No rows should be emitted for empty result set', + ); + done(); + } catch (error) { + done(error); + } + }) + .on('error', done); + + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [], + last: true, + }); + }); + + it('should handle multi-chunk stream ending with chunk.last', done => { + const receivedRows: Row[] = []; + + stream + .on('data', (row: Row) => { + receivedRows.push(row); + }) + .on('end', () => { + try { + assert.strictEqual(receivedRows.length, 2); + done(); + } catch (error) { + done(error); + } + }) + .on('error', done); + + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: false, + resumeToken: 'token1', + }), + ); + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: true, + }), + ); + }); + + it('should emit stats event before end when chunk.last contains stats', done => { + let statsEmitted = false; + let endEmitted = false; + const fakeStats = {queryStats: {rowsReturned: '1'}}; + + stream + .on('data', () => {}) + .on('stats', (stats: any) => { + statsEmitted = true; + assert.strictEqual( + endEmitted, + false, + 'stats must be emitted before end', + ); + assert.deepStrictEqual(stats, fakeStats); + }) + .on('end', () => { + endEmitted = true; + try { + assert.strictEqual(statsEmitted, true); + done(); + } catch (error) { + done(error); + } + }) + .on('error', done); + + fakeRequestStream.push( + Object.assign({}, RESULT, { + stats: fakeStats, + last: true, + }), + ); + }); + + it('should clean up request stream when error occurs on request stream while receivedLast is true before stream ends', done => { + let dataEmitted = false; + + stream + .on('data', () => { + dataEmitted = true; + // Emit error on the request stream while receivedLast is true and before stream ends + fakeRequestStream.emit( + 'error', + new Error('Immediate transport error'), + ); + }) + .on('end', () => { + try { + assert.strictEqual(dataEmitted, true); + done(); + } catch (error) { + done(error); + } + }) + .on('error', err => { + done( + new Error( + `Stream should not emit error after chunk.last: ${err.message}`, + ), + ); + }); + + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: true, + }), + ); + }); + + it('should flush all uncheckpointed chunks queued in CheckpointStream when chunk.last is true', done => { + const receivedRows: Row[] = []; + + stream + .on('data', (row: Row) => { + receivedRows.push(row); + }) + .on('end', () => { + try { + assert.strictEqual(receivedRows.length, 2); + done(); + } catch (error) { + done(error); + } + }) + .on('error', done); + + // Chunk 1 has NO resumeToken and last: false (gets buffered in CheckpointStream) + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: false, + resumeToken: undefined, + }), + ); + // Chunk 2 has last: true and NO resumeToken (must trigger flush of chunk 1 and chunk 2) + fakeRequestStream.push( + Object.assign({}, RESULT, { + last: true, + resumeToken: undefined, + }), + ); + }); + + it('should successfully retry on retryable error and complete on chunk.last', done => { + const unavailableError = new Error('Unavailable') as grpc.ServiceError; + unavailableError.code = grpc.status.UNAVAILABLE; + + let attempts = 0; + const retryRequestFunction = () => { + const requestStream = through.obj(); + attempts++; + if (attempts === 1) { + setImmediate(() => requestStream.emit('error', unavailableError)); + } else { + setImmediate(() => { + requestStream.push(Object.assign({}, RESULT, {last: true})); + }); + } + return requestStream; + }; + + const retryStream = partialResultStream(retryRequestFunction); + let rowsCount = 0; + retryStream + .on('data', () => rowsCount++) + .on('end', () => { + try { + assert.strictEqual(attempts, 2); + assert.strictEqual(rowsCount, 1); + done(); + } catch (error) { + done(error); + } + }) + .on('error', done); + }); + + it('should emit error when decoding fails on chunk.last', done => { + const failingStream = partialResultStream(() => fakeRequestStream, { + json: true, + jsonOptions: {wrapNumbers: false}, + }); + + failingStream + .on('data', () => {}) + .on('end', () => { + done(new Error('Stream should not emit end when decoding fails')); + }) + .on('error', error => { + try { + assert( + error.message.includes( + 'Serializing column "large_id" encountered an error:', + ), + ); + done(); + } catch (assertionError) { + done(assertionError); + } + }); + + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [ + { + name: 'large_id', + type: {code: 'INT64'}, + }, + ], + }, + }, + values: [convertToIValue('9223372036854775807')], + last: true, + }); + }); }); });