diff --git a/handwritten/spanner/package.json b/handwritten/spanner/package.json index 717f82b1538..33c01cdf35d 100644 --- a/handwritten/spanner/package.json +++ b/handwritten/spanner/package.json @@ -71,15 +71,12 @@ "@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", - "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 +95,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 2a1533c5419..f43f0384bdb 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -15,11 +15,8 @@ */ 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 {PassThrough, Readable, Transform} from 'stream'; import * as streamEvents from 'stream-events'; import {grpc, CallOptions} from 'google-gax'; import {DeadlineError, isRetryableInternalError} from './transaction-runner'; @@ -27,7 +24,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; @@ -554,6 +550,104 @@ export class PartialResultStream extends Transform implements ResultEvents { } } +/** + * A custom Transform stream that buffers PartialResultSet chunks and flushes them + * asynchronously to prevent blocking the event loop. + * + * It holds chunks in a queue until a "checkpoint" is reached (as determined by + * `isCheckpointFn`) or until the queue exceeds `maxQueued` items. + * + * 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[] = []; + 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; + } + + /** + * 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). + * @param {Function} callback Callback to signal completion of transformation. + */ + _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(); + } + + /** + * Flushes queued chunks synchronously to prevent state races on retry/reset. + * + * @private + */ + private _flushQueue(): void { + while (this.queue.length > 0 && !this.destroyed) { + this.push(this.queue.shift()); + } + } + + /** + * 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 @@ -578,24 +672,25 @@ 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 flushStream = new stream.PassThrough({objectMode: true}); - requestsStream.add(flushStream); + 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 // 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,26 +712,42 @@ export function partialResultStream( // then push `null` to end the stream. flushStream.push({resumeToken: '_'}); flushStream.push(null); - requestsStream.end(); }); }; + + 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(); } lastRequestStream = requestFn(lastResumeToken); lastRequestStream.on('end', endListener); - requestsStream.add(lastRequestStream); + errorListener = (err: grpc.ServiceError) => { + destroyRequestStream(); + setImmediate(() => retry(err)); + }; + lastRequestStream.on('error', errorListener); + lastRequestStream.pipe(requestsStream, {end: false}); }; 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; } @@ -653,14 +764,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.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 @@ -674,15 +781,12 @@ 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)), - ); + userStream.once('close', () => { + destroyRequestStream(); + requestsStream.destroy(); + flushStream.destroy(); + batchAndSplitOnTokenStream.destroy(); + }); return ( requestsStream diff --git a/handwritten/spanner/test/partial-result-stream.ts b/handwritten/spanner/test/partial-result-stream.ts index 38700a57d24..1aafbfca6a8 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,236 @@ 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); + }); + + 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); + } + }); + }); }); });