Skip to content

perf: optimize stream pipeline by eliminating events-intercept #9221

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alkatrivedi wants to merge 7 commits into
base: main
Choose a base branch
from
Open

perf: optimize stream pipeline by eliminating events-intercept #9221

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
f452d6f
perf: optimize stream pipeline by eliminating events-intercept
alkatrivedi Aug 27, 2026
700e088
gemini review comments
alkatrivedi Aug 28, 2026
a706f0e
refactor
alkatrivedi Aug 28, 2026
21e9811
review comments
alkatrivedi Sep 3, 2026
152d1ab
test: PartialResultStream
alkatrivedi Sep 3, 2026
6289364
doc: class CheckpointStream
alkatrivedi Sep 4, 2026
d1a1f7f
review comment
alkatrivedi Sep 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions handwritten/spanner/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand Down
158 changes: 131 additions & 27 deletions handwritten/spanner/src/partial-result-stream.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,15 @@
*/

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';

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;
Expand Down Expand Up @@ -253,7 +249,7 @@
this._options.columnsMetadata,
name,
)
? (this._options.columnsMetadata as any)[name]

Check warning on line 252 in handwritten/spanner/src/partial-result-stream.ts

View workflow job for this annotation

GitHub Actions / lint 

Unexpected any. Specify a different type
: undefined;
if (codec.decode !== originalDecode) {
return val =>
Expand Down Expand Up @@ -554,6 +550,104 @@
}
}

/**
* 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
Expand All @@ -578,24 +672,25 @@
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);
Expand All @@ -617,26 +712,42 @@
// 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));

@olavloite olavloite Aug 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deferring the retry here with setImmediate(..) means that other events could be handled before the retry. If for example lastRequestStream emits 'end' before the retry, the retry will fail. It could also cause multiple errors to trigger multiple retries.

Verification test cases:

it('should successfully retry when the failed stream emits an error followed by end', 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 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);
      fakeCheckpointStream.emit('checkpoint', 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 fakeCheckpointStream = through.obj();
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  (fakeCheckpointStream as any).reset = () => {};
  sandbox.stub(checkpointStream, 'obj').returns(fakeCheckpointStream);

  const firstStream = through.obj();
  const secondStream = through.obj();

  const requestFnStub = sandbox.stub();

  // 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);
      fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN);
      secondStream.end();
    });
    return secondStream;
  });

  partialResultStream(requestFnStub)
    .on('error', done)
    .pipe(
      concat(rows => {
        try {
          // Exactly 1 initial request + 1 retry request = 2 calls total
          assert.strictEqual(requestFnStub.callCount, 2, 'Should only trigger one retry request');
          assert.strictEqual(rows.length, 1);
          done();
        } catch (e) {
          done(e);
        }
      }),
    );
});

Suggested fix:

--- a/handwritten/spanner/src/partial-result-stream.ts
+++ b/handwritten/spanner/src/partial-result-stream.ts
@@ -618,17 +618,27 @@ 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();
     }
     lastRequestStream = requestFn(lastResumeToken);
     lastRequestStream.on('end', endListener);
     errorListener = (err: grpc.ServiceError) => {
+      destroyRequestStream();
       setImmediate(() => retry(err));
     };
     lastRequestStream.on('error', errorListener);
     lastRequestStream.pipe(requestsStream, {end: false});
   };

@alkatrivedi alkatrivedi Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

};
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;
}
Expand All @@ -653,14 +764,10 @@
// 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
Expand All @@ -674,15 +781,12 @@
};

userStream.once('reading', makeRequest);
eventsIntercept.patch(requestsStream);

@olavloite olavloite Aug 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the application calls stream.destroy(), then that is not propagated into 'our' stream. We should add a listener on the userStream for close and make sure that we clean up then the user stream is closed.

Verification test case:

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);
});

Suggested fix:

--- a/handwritten/spanner/src/partial-result-stream.ts
+++ b/handwritten/spanner/src/partial-result-stream.ts
@@ -677,6 +677,12 @@ export function partialResultStream(
   };
 
   userStream.once('reading', makeRequest);
+  userStream.once('close', () => {
+    destroyRequestStream();
+    requestsStream.destroy();
+    flushStream.destroy();
+    batchAndSplitOnTokenStream.destroy();
+  });
 
   return (
     requestsStream

@alkatrivedi alkatrivedi Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done


// 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
Expand Down
Loading
Loading