From 2d150bff00c06c5228bdec7d30bbd8fb5e6db4c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Wed, 16 Sep 2026 15:05:22 +0200 Subject: [PATCH] perf(spanner): replace setTimeout backpressure with event-driven resumption Replace the setTimeout backpressure polling loop in PartialResultStream with native Node.js stream flow control: - Hold the transform completion callback when downstream push() returns false, and invoke it upon stream drain via _read() to resume upstream ingestion. - Remove _tryResume, _numPushFailed, and the artificial 'Stream is still not ready to receive data' timeout failures. - Guard 'paused' event emission in _addChunk so it only emits once per backpressure transition instead of on every subsequent value. - Mark maxResumeRetries as deprecated on streaming RowOptions, clarifying that it is only used by non-streaming Snapshot.run() for RPC retry limits. - Replace setTimeout with setImmediate in test/spanner.ts slow-writer test, eliminating a 5-second sleep in compliance with the zero-sleep test policy. - Add comprehensive backpressure unit tests covering single-chunk direct decoding, multi-chunk pause/resume cycles, cross-chunk value stitching, stream destroy/error propagation, and full-pipeline flow control. --- .../spanner/src/partial-result-stream.ts | 63 +-- handwritten/spanner/src/transaction.ts | 10 +- .../spanner/test/partial-result-stream.ts | 473 ++++++++++++++++++ handwritten/spanner/test/spanner.ts | 17 +- 4 files changed, 506 insertions(+), 57 deletions(-) diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index 8b5b03e40528..8ad655e668c0 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -45,11 +45,9 @@ interface RequestFunction { * @property {boolean} [json=false] Indicates if the Row objects should be * formatted into JSON. * @property {JSONOptions} [jsonOptions] JSON options. - * @property {number} [maxResumeRetries=20] The maximum number of times that the - * stream will retry to push data downstream, when the downstream indicates - * that it is not ready for any more data. Increase this value if you - * experience 'Stream is still not ready to receive data' errors as a - * result of a slow writer in your receiving stream. + * @property {number} [maxResumeRetries] @deprecated Backpressure is now handled + * automatically via stream flow control. This option is no longer used by + * streaming queries and is retained only for backward compatibility. * @property {object} [columnsMetadata] An object map that can be used to pass * additional properties for each column type which can help in deserializing * the data coming from backend. (Eg: We need to pass Proto Function and Enum @@ -368,13 +366,13 @@ export class PartialResultStream extends Transform implements ResultEvents { private _pendingValue?: p.IValue; private _pendingValueForResume?: p.IValue; private _values: p.IValue[]; - private _numPushFailed = 0; + private _resumeCallback?: () => void; private _isFirstChunk = true; constructor(options = {}) { super({objectMode: true}); this._destroyed = false; - this._options = Object.assign({maxResumeRetries: 20}, options); + this._options = Object.assign({}, options); this._values = []; this._isFirstChunk = true; } @@ -389,6 +387,7 @@ export class PartialResultStream extends Transform implements ResultEvents { } this._destroyed = true; + this._resumeCallback = undefined; process.nextTick(() => { if (err) { @@ -447,42 +446,21 @@ export class PartialResultStream extends Transform implements ResultEvents { if (res) { next(); } else { - // Wait a little before we push any more data into the pipeline as a - // component downstream has indicated that a break is needed. Pause the - // request stream to prevent it from filling up the buffer while we are - // waiting. - // The stream will initially pause for 2ms, and then double the pause time - // for each new pause. - const initialPauseMs = 2; - setTimeout(() => { - this._tryResume(next, 2 * initialPauseMs); - }, initialPauseMs); + // Downstream buffer has reached highWaterMark and cannot accept more data + // at the moment. Hold the completion callback until the downstream consumer + // drains and Node invokes _read(), resuming the upstream request stream. + this._resumeCallback = next as () => void; } } - private _tryResume(next: Function, timeout: number) { - // Try to push an empty chunk to check whether more data can be accepted. - if (this.push(undefined)) { - this._numPushFailed = 0; + _read(size: number): void { + if (this._resumeCallback) { + const callback = this._resumeCallback; + this._resumeCallback = undefined; this.emit('resumed'); - next(); - } else { - // Downstream returned false indicating that it is still not ready for - // more data. - this._numPushFailed++; - if (this._numPushFailed === this._options.maxResumeRetries) { - this.destroy( - new Error( - `Stream is still not ready to receive data after ${this._numPushFailed} attempts to resume.`, - ), - ); - return; - } - setTimeout(() => { - const nextTimeout = Math.min(timeout * 2, 1024); - this._tryResume(next, nextTimeout); - }, timeout); + callback(); } + super._read(size); } _resetPendingValues() { @@ -541,15 +519,16 @@ export class PartialResultStream extends Transform implements ResultEvents { delete this._pendingValueForResume; } - let res = true; + let canAcceptMore = true; const len = values.length; for (let i = 0; i < len; i++) { - res = this._addValue(values[i]) && res; - if (!res) { + const accepted = this._addValue(values[i]); + if (!accepted && canAcceptMore) { + canAcceptMore = false; this.emit('paused'); } } - return res; + return canAcceptMore; } /** diff --git a/handwritten/spanner/src/transaction.ts b/handwritten/spanner/src/transaction.ts index f84b73362eb7..0290ae42a21e 100644 --- a/handwritten/spanner/src/transaction.ts +++ b/handwritten/spanner/src/transaction.ts @@ -1848,11 +1848,11 @@ export class Snapshot extends EventEmitter { * @property {object} [gaxOptions] Request configuration options, * See {@link https://googleapis.dev/nodejs/google-gax/latest/interfaces/CallOptions.html|CallOptions} * for more details. - * @property {number} [maxResumeRetries] The maximum number of times that the - * stream will retry to push data downstream, when the downstream indicates - * that it is not ready for any more data. Increase this value if you - * experience 'Stream is still not ready to receive data' errors as a - * result of a slow writer in your receiving stream. + * @property {number} [maxResumeRetries] The maximum number of times that the + * query will retry on retryable errors (such as UNAVAILABLE). Only + * applicable to non-streaming queries executed via {@link Snapshot#run}. + * For streaming queries ({@link Snapshot#runStream}), this option is + * deprecated as backpressure is managed automatically. * @property {object} [directedReadOptions] * Indicates which replicas or regions should be used for non-transactional reads or queries. */ diff --git a/handwritten/spanner/test/partial-result-stream.ts b/handwritten/spanner/test/partial-result-stream.ts index 032f4e1db3d5..698d92827e05 100644 --- a/handwritten/spanner/test/partial-result-stream.ts +++ b/handwritten/spanner/test/partial-result-stream.ts @@ -711,6 +711,424 @@ describe('PartialResultStream', () => { }); stream.end(); }); + + describe('event-driven backpressure', () => { + it('should hold completion callback and emit paused when downstream push returns false', done => { + const stream = new PartialResultStream({}); + let pausedEmitted = false; + stream.on('paused', () => { + pausedEmitted = true; + }); + + // Stub push to simulate downstream backpressure on rows. + const pushStub = sandbox.stub(stream, 'push'); + // Accept first row, reject second row to trigger backpressure. + pushStub.onFirstCall().returns(true); + pushStub.onSecondCall().returns(false); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + let writeCallbackCalled = false; + + stream.write( + { + metadata: {rowType: {fields}}, + values: [convertToIValue('row1'), convertToIValue('row2')], + }, + () => { + writeCallbackCalled = true; + }, + ); + + // The write callback should NOT have been called because the stream is paused. + assert.strictEqual(writeCallbackCalled, false); + assert.strictEqual(pausedEmitted, true); + assert.strictEqual( + typeof (stream as unknown as {_resumeCallback?: Function}) + ._resumeCallback, + 'function', + ); + done(); + }); + + it('should resume and invoke held callback when _read is called', done => { + const stream = new PartialResultStream({}); + let resumedEmitted = false; + stream.on('resumed', () => { + resumedEmitted = true; + }); + + const pushStub = sandbox.stub(stream, 'push'); + pushStub.returns(false); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + let writeCallbackCalled = false; + + stream.write( + { + metadata: {rowType: {fields}}, + values: [convertToIValue('row1')], + }, + () => { + writeCallbackCalled = true; + }, + ); + + assert.strictEqual(writeCallbackCalled, false); + + // Simulate Node readable machinery calling _read when readable buffer drains. + (stream as unknown as {_read: (size: number) => void})._read(1); + + assert.strictEqual(writeCallbackCalled, true); + assert.strictEqual(resumedEmitted, true); + assert.strictEqual( + (stream as unknown as {_resumeCallback?: Function})._resumeCallback, + undefined, + ); + done(); + }); + + it('should clear _resumeCallback when stream is destroyed while paused', done => { + const stream = new PartialResultStream({}); + const pushStub = sandbox.stub(stream, 'push'); + pushStub.returns(false); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + stream.write({ + metadata: {rowType: {fields}}, + values: [convertToIValue('row1')], + }); + + assert.strictEqual( + typeof (stream as unknown as {_resumeCallback?: Function}) + ._resumeCallback, + 'function', + ); + + stream.destroy(); + + assert.strictEqual( + (stream as unknown as {_resumeCallback?: Function})._resumeCallback, + undefined, + ); + done(); + }); + + it('should stream all rows to a slow writable stream with backpressure', done => { + const stream = new PartialResultStream({}); + const rows: Row[] = []; + let pausedCount = 0; + let resumedCount = 0; + + stream.on('paused', () => { + pausedCount++; + }); + stream.on('resumed', () => { + resumedCount++; + }); + + const slowSink = new Transform({ + objectMode: true, + highWaterMark: 1, + transform(chunk, encoding, callback) { + rows.push(chunk); + setImmediate(callback); + }, + }); + + stream.pipe(slowSink); + + const totalRows = 25; + slowSink.on('finish', () => { + try { + assert.strictEqual(rows.length, totalRows); + assert.ok(pausedCount > 0, 'should have paused at least once'); + assert.ok(resumedCount > 0, 'should have resumed at least once'); + done(); + } catch (err) { + done(err); + } + }); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + const values1: Array> = []; + for (let i = 0; i < 20; i++) { + values1.push(convertToIValue(`row${i}`)); + } + const values2: Array> = []; + for (let i = 20; i < totalRows; i++) { + values2.push(convertToIValue(`row${i}`)); + } + + stream.write({ + metadata: {rowType: {fields}}, + values: values1, + last: false, + }); + stream.write({ + values: values2, + last: true, + }); + stream.end(); + }); + + it('should emit paused only once per transition even with multiple unaccepted rows in a chunk', done => { + const stream = new PartialResultStream({}); + let pausedCount = 0; + stream.on('paused', () => { + pausedCount++; + }); + + const pushStub = sandbox.stub(stream, 'push'); + // Accept first row, reject all subsequent rows + pushStub.onFirstCall().returns(true); + pushStub.returns(false); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + const values = [ + convertToIValue('row1'), + convertToIValue('row2'), + convertToIValue('row3'), + convertToIValue('row4'), + convertToIValue('row5'), + ]; + + stream.write( + { + metadata: {rowType: {fields}}, + values, + }, + () => {}, + ); + + // Even though rows 2..5 were rejected by push(), paused should only be emitted ONCE + assert.strictEqual(pausedCount, 1); + done(); + }); + + it('should handle backpressure cleanly during single-chunk optimization', done => { + const stream = new PartialResultStream({}); + let pausedCount = 0; + stream.on('paused', () => pausedCount++); + + const rows: Row[] = []; + const slowSink = new Transform({ + objectMode: true, + highWaterMark: 1, + transform(chunk, encoding, callback) { + rows.push(chunk); + setImmediate(callback); + }, + }); + + stream.pipe(slowSink); + + const totalRows = 25; + slowSink.on('finish', () => { + try { + assert.strictEqual(rows.length, totalRows); + assert.ok(pausedCount > 0, 'should emit paused on backpressure'); + done(); + } catch (err) { + done(err); + } + }); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + const values: Array> = []; + for (let i = 0; i < totalRows; i++) { + values.push(convertToIValue(`single_chunk_row_${i}`)); + } + + // First chunk with last=true triggers _addSingleChunk + stream.write({ + metadata: {rowType: {fields}}, + values, + last: true, + }); + stream.end(); + }); + + it('should handle multiple sequential pause and resume cycles across chunks', done => { + const stream = new PartialResultStream({}); + let pausedCount = 0; + let resumedCount = 0; + stream.on('paused', () => pausedCount++); + stream.on('resumed', () => resumedCount++); + + const rows: Row[] = []; + const slowSink = new Transform({ + objectMode: true, + highWaterMark: 1, + transform(chunk, encoding, callback) { + rows.push(chunk); + setImmediate(callback); + }, + }); + + stream.pipe(slowSink); + + const totalRows = 45; + slowSink.on('finish', () => { + try { + assert.strictEqual(rows.length, totalRows); + for (let i = 0; i < totalRows; i++) { + assert.strictEqual(rows[i][0].value, `val_${i}`); + } + assert.ok( + pausedCount >= 2, + `expected at least 2 pauses, got ${pausedCount}`, + ); + assert.ok( + resumedCount >= 2, + `expected at least 2 resumes, got ${resumedCount}`, + ); + done(); + } catch (err) { + done(err); + } + }); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + const chunkSizes = [20, 20, 5]; + let offset = 0; + for (let chunkIndex = 0; chunkIndex < chunkSizes.length; chunkIndex++) { + const count = chunkSizes[chunkIndex]; + const values: Array> = []; + for (let i = 0; i < count; i++) { + values.push(convertToIValue(`val_${offset + i}`)); + } + offset += count; + stream.write({ + ...(chunkIndex === 0 ? {metadata: {rowType: {fields}}} : {}), + values, + last: chunkIndex === chunkSizes.length - 1, + }); + } + stream.end(); + }); + + it('should preserve row assembly when backpressure occurs across chunked values', done => { + const stream = new PartialResultStream({}); + const rows: Row[] = []; + + const slowSink = new Transform({ + objectMode: true, + highWaterMark: 1, + transform(chunk, encoding, callback) { + rows.push(chunk); + setImmediate(callback); + }, + }); + + stream.pipe(slowSink); + + slowSink.on('finish', () => { + try { + assert.strictEqual(rows.length, 2); + assert.strictEqual(rows[0][0].value, 'first_row'); + assert.strictEqual(rows[1][0].value, 'chunked_part1_part2'); + done(); + } catch (err) { + done(err); + } + }); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + + // Chunk 1: first complete row + start of chunked row + stream.write({ + metadata: {rowType: {fields}}, + values: [ + convertToIValue('first_row'), + convertToIValue('chunked_part1_'), + ], + chunkedValue: true, + last: false, + }); + + // Chunk 2: continuation of chunked row + stream.write({ + values: [convertToIValue('part2')], + chunkedValue: false, + last: true, + }); + stream.end(); + }); + + it('should propagate error and clean up held callback when destroyed with error while paused', done => { + const stream = new PartialResultStream({}); + const pushStub = sandbox.stub(stream, 'push'); + pushStub.returns(false); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + stream.write({ + metadata: {rowType: {fields}}, + values: [convertToIValue('row1')], + }); + + assert.strictEqual( + typeof (stream as unknown as {_resumeCallback?: Function}) + ._resumeCallback, + 'function', + ); + + const testError = new Error('simulated failure'); + stream.on('error', err => { + try { + assert.strictEqual(err, testError); + assert.strictEqual( + (stream as unknown as {_resumeCallback?: Function}) + ._resumeCallback, + undefined, + ); + done(); + } catch (assertionErr) { + done(assertionErr); + } + }); + + stream.destroy(testError); + }); + + it('should not emit paused or resumed when consumer is fast', done => { + const stream = new PartialResultStream({}); + let pausedEmitted = false; + let resumedEmitted = false; + + stream.on('paused', () => { + pausedEmitted = true; + }); + stream.on('resumed', () => { + resumedEmitted = true; + }); + + const rows: Row[] = []; + stream.on('data', row => rows.push(row)); + stream.on('end', () => { + try { + assert.strictEqual(rows.length, 10); + assert.strictEqual(pausedEmitted, false); + assert.strictEqual(resumedEmitted, false); + done(); + } catch (err) { + done(err); + } + }); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + const values: Array> = []; + for (let i = 0; i < 10; i++) { + values.push(convertToIValue(`row_${i}`)); + } + + stream.write({ + metadata: {rowType: {fields}}, + values, + last: true, + }); + stream.end(); + }); + }); }); describe('partialResultStream', () => { @@ -1578,6 +1996,61 @@ describe('PartialResultStream', () => { last: true, }); }); + + it('should handle downstream backpressure through the full pipeline without dropping rows', done => { + const rows: Row[] = []; + let pausedCount = 0; + let resumedCount = 0; + + stream.on('paused', () => pausedCount++); + stream.on('resumed', () => resumedCount++); + + const slowSink = new Transform({ + objectMode: true, + highWaterMark: 1, + transform(chunk, encoding, callback) { + rows.push(chunk); + setImmediate(callback); + }, + }); + + stream.pipe(slowSink); + + const totalRows = 25; + slowSink.on('finish', () => { + try { + assert.strictEqual(rows.length, totalRows); + assert.ok(pausedCount > 0, 'pipeline should pause on backpressure'); + assert.ok(resumedCount > 0, 'pipeline should resume on drain'); + done(); + } catch (err) { + done(err); + } + }); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + const values1: Array> = []; + for (let i = 0; i < 20; i++) { + values1.push(convertToIValue(`pipeline_row_${i}`)); + } + const values2: Array> = []; + for (let i = 20; i < totalRows; i++) { + values2.push(convertToIValue(`pipeline_row_${i}`)); + } + + fakeRequestStream.push({ + metadata: {rowType: {fields}}, + values: values1, + resumeToken: 'token1', + last: false, + }); + fakeRequestStream.push({ + values: values2, + resumeToken: 'token2', + last: true, + }); + fakeRequestStream.push(null); + }); }); describe('decodeRowsDirect & createFieldDecoders', () => { diff --git a/handwritten/spanner/test/spanner.ts b/handwritten/spanner/test/spanner.ts index 926aeb8ca63b..ba3149117885 100644 --- a/handwritten/spanner/test/spanner.ts +++ b/handwritten/spanner/test/spanner.ts @@ -1001,13 +1001,14 @@ describe('Spanner with mock server', () => { } }); - it('should fail on slow writer when maxResumeRetries has been exceeded', async () => { + it('should not fail on slow writer even when maxResumeRetries is small', async () => { const largeSelect = 'select * from large_table'; spannerMock.putStatementResult( largeSelect, mock.StatementResult.resultSet(mock.createLargeResultSet()), ); const database = newTestDatabase(); + let rowCount = 0; try { const rs = database.runStream({ sql: largeSelect, @@ -1022,10 +1023,11 @@ describe('Spanner with mock server', () => { highWaterMark: 1, objectMode: true, transform(chunk, encoding, callback) { - // Simulate a slow flush. - setTimeout(() => { + rowCount++; + // Simulate an asynchronous slow consumer using setImmediate. + setImmediate(() => { callback(undefined, chunk); - }, 50); + }); }, }), new stream.Transform({ @@ -1035,12 +1037,7 @@ describe('Spanner with mock server', () => { }, }), ); - assert.fail('missing expected error'); - } catch (err) { - assert.strictEqual( - (err as ServiceError).message, - 'Stream is still not ready to receive data after 1 attempts to resume.', - ); + assert.strictEqual(rowCount, NUM_ROWS_LARGE_RESULT_SET); } finally { await database.close(); }