diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index 8b5b03e4052..53ab0ab6219 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -367,7 +367,9 @@ export class PartialResultStream extends Transform implements ResultEvents { private _options: RowOptions; private _pendingValue?: p.IValue; private _pendingValueForResume?: p.IValue; - private _values: p.IValue[]; + private _rowValues?: Value[]; + private _valueIndex: number; + private _valueIndexForResume?: number; private _numPushFailed = 0; private _isFirstChunk = true; constructor(options = {}) { @@ -375,7 +377,7 @@ export class PartialResultStream extends Transform implements ResultEvents { this._destroyed = false; this._options = Object.assign({maxResumeRetries: 20}, options); - this._values = []; + this._valueIndex = 0; this._isFirstChunk = true; } /** @@ -389,6 +391,7 @@ export class PartialResultStream extends Transform implements ResultEvents { } this._destroyed = true; + this._rowValues = undefined; process.nextTick(() => { if (err) { @@ -486,70 +489,117 @@ export class PartialResultStream extends Transform implements ResultEvents { } _resetPendingValues() { - if (this._pendingValueForResume) { - this._pendingValue = this._pendingValueForResume; + this._pendingValue = this._pendingValueForResume; + if (this._valueIndexForResume !== undefined) { + this._valueIndex = this._valueIndexForResume; } else { - delete this._pendingValue; + this._valueIndex = 0; } + this._rowValues?.fill(undefined, this._valueIndex); } /** - * Manages any chunked values. + * Manages stream chunks, chunked value merging across chunk boundaries, + * row assembly, and resume token checkpoints. * - * @private + * Processing follows 5 distinct stages: + * 1. Single-chunk fast path: If the entire stream response is in this first chunk, + * decode rows directly and bypass incremental chunk buffering. + * 2. Pending value merge: If the previous chunk ended with an incomplete chunked + * value, merge it with the incoming continuation value at chunkValues[0]. + * 3. Chunked value hold: If this chunk ends with a chunked value, hold the tail + * value (chunkValues[numValues - 1]) in `this._pendingValue` for the next chunk. + * 4. Complete value decoding: Iterate from `startIndex` to `endIndex`, decoding + * and adding values into row buffers via `_addValue`. + * 5. Checkpoint tracking: If this chunk contains a resume token, save the pending + * chunked value and the current column write cursor (`_valueIndex`) for retries. * - * @param {object} chunk The partial result set. + * @private + * @param {google.spanner.v1.PartialResultSet} chunk The partial result set. + * @returns {boolean} Whether downstream can accept more data. */ private _addChunk(chunk: google.spanner.v1.PartialResultSet): boolean { const isFirstChunk = this._isFirstChunk; this._isFirstChunk = false; - // Fast path for single-chunk stream responses: + // Stage 1: Fast path for single-chunk stream responses: if (isFirstChunk && chunk.last && !chunk.chunkedValue) { return this._addSingleChunk(chunk); } const chunkValues = chunk.values; const numValues = chunkValues.length; - const values: Value[] = new Array(numValues); - for (let i = 0; i < numValues; i++) { - values[i] = GrpcService.decodeValue_(chunkValues[i]); - } + let startIndex = 0; + let endIndex = numValues; + let canAcceptMore = true; - // If we have a chunk to merge, merge the values now. - if (this._pendingValue) { - const currentField = this._values.length % this._fields.length; + // Stage 2: Merge pending chunked value from the previous chunk with the + // incoming continuation value at chunkValues[0]. + if (this._pendingValue && numValues > 0) { + const currentField = this._valueIndex; const field = this._fields[currentField]; + const headValue = this._pendingValue; + const continuationValue = GrpcService.decodeValue_(chunkValues[0]); const merged = PartialResultStream.merge( field.type as google.spanner.v1.Type, - this._pendingValue, - values.shift(), + headValue, + continuationValue, ); - values.unshift(...merged); - delete this._pendingValue; + // We consumed chunkValues[0] as the continuation of the pending value. + startIndex = 1; + + // If this chunk only had 1 value and is still chunked, the last element + // of merged remains pending for the next chunk. + let mergedCount = merged.length; + if (numValues === 1 && chunk.chunkedValue) { + mergedCount--; + this._pendingValue = merged[mergedCount]; + } else { + this._pendingValue = undefined; + } + + for (let i = 0; i < mergedCount; i++) { + if (!this._addValue(merged[i]) && canAcceptMore) { + canAcceptMore = false; + this.emit('paused'); + } + } } - // If the chunk is chunked, store the last value for merging with the next - // chunk to be processed. - if (chunk.chunkedValue) { - this._pendingValue = values.pop(); - if (_hasResumeToken(chunk)) { - this._pendingValueForResume = this._pendingValue; + // Stage 3: If this chunk ends with a chunked value, hold the tail value for + // merging with the next chunk instead of decoding it into the current row now. + if (chunk.chunkedValue && numValues > 0) { + if (numValues > 1 || !startIndex) { + endIndex = numValues - 1; + this._pendingValue = GrpcService.decodeValue_( + chunkValues[numValues - 1], + ); } - } else if (_hasResumeToken(chunk)) { - delete this._pendingValueForResume; } - let res = true; - const len = values.length; - for (let i = 0; i < len; i++) { - res = this._addValue(values[i]) && res; - if (!res) { + // Stage 4: Decode in-place and push complete values into row buffers. + for (let i = startIndex; i < endIndex; i++) { + const value = GrpcService.decodeValue_(chunkValues[i]); + if (!this._addValue(value) && canAcceptMore) { + canAcceptMore = false; this.emit('paused'); } } - return res; + + // Stage 5: If this chunk contains a resume token, record the checkpoint state. + // Both the pending chunked value and the row column index are preserved so that + // a retry can resume at the exact column position. + if (_hasResumeToken(chunk)) { + if (chunk.chunkedValue) { + this._pendingValueForResume = this._pendingValue; + } else { + this._pendingValueForResume = undefined; + } + this._valueIndexForResume = this._valueIndex; + } + + return canAcceptMore; } /** @@ -587,18 +637,20 @@ export class PartialResultStream extends Transform implements ResultEvents { * @param {*} value The complete value. */ private _addValue(value: Value): boolean { - const values = this._values; + if (!this._rowValues) { + this._rowValues = new Array(this._fields.length); + } - values.push(value); + this._rowValues[this._valueIndex++] = value; - if (values.length !== this._fields.length) { + if (this._valueIndex !== this._fields.length) { return true; } - this._values = []; + this._valueIndex = 0; return this.push( - formatRow(this._fields, this._decoders, values, this._options), + formatRow(this._fields, this._decoders, this._rowValues, this._options), ); } @@ -898,9 +950,7 @@ export function partialResultStream( }; const makeRequest = (): void => { - if (isDefined(lastResumeToken) && lastResumeToken.length > 0) { - partialRSStream._resetPendingValues(); - } + partialRSStream._resetPendingValues(); lastRequestStream = requestFn(lastResumeToken); lastRequestStream.on('end', endListener); errorListener = (err: grpc.ServiceError) => { diff --git a/handwritten/spanner/test/partial-result-stream.ts b/handwritten/spanner/test/partial-result-stream.ts index 032f4e1db3d..b8ebfa30c1a 100644 --- a/handwritten/spanner/test/partial-result-stream.ts +++ b/handwritten/spanner/test/partial-result-stream.ts @@ -654,6 +654,146 @@ describe('PartialResultStream', () => { done(); }); + it('should emit paused event exactly once when downstream backpressure is triggered during multi-chunk streaming', done => { + const stream = new PartialResultStream({}); + let pausedCount = 0; + stream.on('paused', () => { + pausedCount++; + }); + + sandbox.stub(stream, 'push').callsFake(data => { + if (data === undefined || data === null) { + return true; + } + return false; + }); + + const fields = [ + {name: 'col1', type: {code: 'STRING'}}, + {name: 'col2', type: {code: 'STRING'}}, + ]; + // First chunk establishes stream and metadata, not last + stream.write({ + metadata: {rowType: {fields}}, + values: [convertToIValue('val1'), convertToIValue('val2')], + }); + // Second chunk emits multiple rows while push() returns false + stream.write({ + values: [ + convertToIValue('val3'), + convertToIValue('val4'), + convertToIValue('val5'), + convertToIValue('val6'), + ], + last: true, + }); + + assert.strictEqual(pausedCount, 1); + done(); + }); + + it('should handle multi-chunk streaming where an intermediate chunk has a single value that remains chunked', done => { + const stream = new PartialResultStream({}); + const rows: prs.Row[] = []; + stream + .on('data', row => rows.push(row)) + .on('end', () => { + try { + assert.strictEqual(rows.length, 2); + assert.deepStrictEqual(rows[0].toJSON(), { + id: 'id1', + text: 'hello-world-again', + }); + assert.deepStrictEqual(rows[1].toJSON(), { + id: 'id2', + text: 'text2', + }); + done(); + } catch (err) { + done(err); + } + }) + .on('error', done); + + const fields = [ + {name: 'id', type: {code: 'STRING'}}, + {name: 'text', type: {code: 'STRING'}}, + ]; + // Chunk 1: starts row 1, ends with partial text 'hello-' + stream.write({ + metadata: {rowType: {fields}}, + values: [convertToIValue('id1'), convertToIValue('hello-')], + chunkedValue: true, + }); + // Chunk 2: only has 1 value, and is still chunked ('world-') + stream.write({ + values: [convertToIValue('world-')], + chunkedValue: true, + }); + // Chunk 3: completes row 1 and provides complete row 2 + stream.write({ + values: [ + convertToIValue('again'), + convertToIValue('id2'), + convertToIValue('text2'), + ], + last: true, + }); + stream.end(); + }); + + it('should reuse pre-allocated row buffer across multiple rows and chunks without cross-row contamination', done => { + const stream = new PartialResultStream({}); + const rows: prs.Row[] = []; + stream + .on('data', row => rows.push(row)) + .on('end', () => { + try { + assert.strictEqual(rows.length, 3); + assert.deepStrictEqual( + rows.map(row => row.toJSON()), + [ + {a: '1', b: '2', c: '3'}, + {a: '4', b: '5', c: '6'}, + {a: '7', b: '8', c: '9'}, + ], + ); + done(); + } catch (err) { + done(err); + } + }) + .on('error', done); + + const fields = [ + {name: 'a', type: {code: 'STRING'}}, + {name: 'b', type: {code: 'STRING'}}, + {name: 'c', type: {code: 'STRING'}}, + ]; + // Chunk 1: completes row 1, starts row 2 + stream.write({ + metadata: {rowType: {fields}}, + values: [ + convertToIValue('1'), + convertToIValue('2'), + convertToIValue('3'), + convertToIValue('4'), + ], + }); + // Chunk 2: completes row 2, completes row 3 + stream.write({ + values: [ + convertToIValue('5'), + convertToIValue('6'), + convertToIValue('7'), + convertToIValue('8'), + convertToIValue('9'), + ], + last: true, + }); + stream.end(); + }); + it('should route first chunk with last=true to _addSingleChunk', done => { const stream = new PartialResultStream({}); const addSingleChunkSpy = sandbox.spy(stream as any, '_addSingleChunk'); @@ -711,6 +851,80 @@ describe('PartialResultStream', () => { }); stream.end(); }); + + describe('_resetPendingValues', () => { + it('should reset pending values and row buffer to beginning if no resume token was received', () => { + const stream = new PartialResultStream({}); + const internalStream = stream as unknown as { + _valueIndex: number; + _pendingValue?: unknown; + _rowValues: unknown[]; + }; + const fields = [ + {name: 'col1', type: {code: 'STRING'}}, + {name: 'col2', type: {code: 'STRING'}}, + {name: 'col3', type: {code: 'STRING'}}, + ]; + stream.write({ + metadata: {rowType: {fields}}, + values: [convertToIValue('val1'), convertToIValue('part-')], + chunkedValue: true, + }); + + assert.strictEqual(internalStream._valueIndex, 1); + assert.strictEqual(internalStream._pendingValue, 'part-'); + assert.strictEqual(internalStream._rowValues[0], 'val1'); + + stream._resetPendingValues(); + + assert.strictEqual(internalStream._valueIndex, 0); + assert.strictEqual(internalStream._pendingValue, undefined); + assert.strictEqual(internalStream._rowValues[0], undefined); + assert.strictEqual(internalStream._rowValues[1], undefined); + assert.strictEqual(internalStream._rowValues[2], undefined); + }); + + it('should restore pending values and row buffer to saved checkpoint when a resume token was received', () => { + const stream = new PartialResultStream({}); + const internalStream = stream as unknown as { + _valueIndex: number; + _pendingValue?: unknown; + _rowValues: unknown[]; + }; + const fields = [ + {name: 'col1', type: {code: 'STRING'}}, + {name: 'col2', type: {code: 'STRING'}}, + {name: 'col3', type: {code: 'STRING'}}, + ]; + // Chunk 1 has resume token and partial row (col1 only) + stream.write({ + metadata: {rowType: {fields}}, + values: [convertToIValue('checkpointVal')], + resumeToken: 'token-1', + }); + + assert.strictEqual(internalStream._valueIndex, 1); + assert.strictEqual(internalStream._rowValues[0], 'checkpointVal'); + + // Chunk 2 continues without resume token and has chunkedValue + stream.write({ + values: [convertToIValue('val2-part-')], + chunkedValue: true, + }); + + assert.strictEqual(internalStream._valueIndex, 1); + assert.strictEqual(internalStream._pendingValue, 'val2-part-'); + + // Reset should roll back to Chunk 1 checkpoint + stream._resetPendingValues(); + + assert.strictEqual(internalStream._valueIndex, 1); + assert.strictEqual(internalStream._pendingValue, undefined); + assert.strictEqual(internalStream._rowValues[0], 'checkpointVal'); + assert.strictEqual(internalStream._rowValues[1], undefined); + assert.strictEqual(internalStream._rowValues[2], undefined); + }); + }); }); describe('partialResultStream', () => { @@ -892,6 +1106,72 @@ describe('PartialResultStream', () => { ); }); + it('should correctly resume and preserve incomplete row state when resumed stream first chunk contains metadata', done => { + const firstStream = through.obj(); + const secondStream = through.obj(); + const requestFnStub = sandbox.stub(); + + const metadata = { + rowType: { + fields: [ + {name: 'col1', type: {code: 'STRING'}}, + {name: 'col2', type: {code: 'STRING'}}, + ], + }, + }; + + requestFnStub.onCall(0).callsFake(() => { + setImmediate(() => { + firstStream.push({ + metadata, + values: [convertToIValue('val1')], + resumeToken: 'checkpoint-token', + }); + + setImmediate(() => { + firstStream.emit('error', { + code: grpc.status.UNAVAILABLE, + message: 'Unavailable', + } as grpc.ServiceError); + }); + }); + + return firstStream; + }); + + requestFnStub.onCall(1).callsFake(resumeToken => { + assert.strictEqual(resumeToken, 'checkpoint-token'); + + setImmediate(() => { + secondStream.push({ + metadata, + values: [convertToIValue('val2')], + last: true, + }); + secondStream.end(); + }); + + return secondStream; + }); + + const rows: Row[] = []; + partialResultStream(requestFnStub) + .on('data', row => rows.push(row)) + .on('end', () => { + try { + assert.strictEqual(rows.length, 1); + assert.deepStrictEqual(rows[0].toJSON(), { + col1: 'val1', + col2: 'val2', + }); + done(); + } catch (err) { + done(err); + } + }) + .on('error', done); + }); + it('should emit non-retryable error', done => { // This test will emit two rows and then an error. const fakeRequestStream = through.obj();