perf: optimize stream pipeline by eliminating events-intercept

[alkatrivedi]

Optimizes performance by removing the events-intercept library and its associated CPU penalty caused by monkey-patching EventEmitter.emit.

Changes:

• Replaced merge-stream with standard Node.js PassThrough streams.
• Replaced error interception hacks with direct error listeners (.on('error')) on the request stream for retry handling.
• Removed events-intercept and merge-stream dependencies from package.json.

Impact: Eliminates artificial CPU tax paid on every stream event and simplifies pipeline architecture without breaking retry functionality.

[gemini-code-assist]

Code Review

This pull request removes the external dependencies events-intercept and merge-stream, replacing them with Node's native stream.PassThrough to handle stream merging and error interception. However, the current implementation introduces critical runtime issues: the stream module is not imported, which will lead to a ReferenceError when instantiating PassThrough, and piping flushStream with {end: false} will cause the stream to hang indefinitely instead of terminating properly.

[olavloite]

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

[olavloite]

This cleanup code is skipped if the early return above for non-retriable errors is used (the one on line 657). That causes the lastRequestStream open and with all event listeners attached.

Verification test:

it('should destroy the request stream and detach listeners on non-retryable errors', done => {
  const fakeCheckpointStream = through.obj();
  sandbox.stub(checkpointStream, 'obj').returns(fakeCheckpointStream);

  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'), 0, 'errorListener should be removed');
        done();
      } catch (e) {
        done(e);
      }
    });
});

Suggested fix:

--- a/handwritten/spanner/src/partial-result-stream.ts
+++ b/handwritten/spanner/src/partial-result-stream.ts
@@ -633,6 +633,7 @@ export function partialResultStream(
   };
 
   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
@@ -659,13 +660,6 @@ export function partialResultStream(
       return;
     }
 
-    if (lastRequestStream) {
-      lastRequestStream.removeListener('end', endListener);
-      lastRequestStream.removeAllListeners('error');
-      lastRequestStream.on('error', () => {}); // Prevent unhandled exception crash
-      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

[olavloite]

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

[olavloite]

This actually also imports events-intercept as a transitive dependency. So the removal of the import below does not actually remove it entirely from this file. And checkpointStream also monkey-patches the emit method. So while this PR gets rid of some of the monkey-patching of the query stream, it does not get rid of all of it.

Test:

it('should not have events-intercept monkey-patching the stream pipeline', () => {
  const stream = checkpointStream.obj();
  // @ts-ignore
  const hasIntercept = typeof stream.intercept === 'function';
  assert.strictEqual(
    hasIntercept,
    false,
    'events-intercept is still monkey-patching stream pipeline via checkpoint-stream',
  );
});

We could instead implement our own specific 'checkpointStream' without monkey-patching and remove the entire dependency on checkpointStream here:

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

  _transform(
    chunk: google.spanner.v1.PartialResultSet,
    enc: string,
    callback: () => void,
  ): void {
    this.queue.push(chunk);
    const isCheckpoint = this.isCheckpointFn(chunk);
    if (isCheckpoint) {
      this.emit('checkpoint', chunk);
      this._flushQueue();
    } else if (this.queue.length > this.maxQueued) {
      this._flushQueue();
    }
    callback();
  }

  private _flushQueue(): void {
    while (this.queue.length > 0) {
      this.push(this.queue.shift());
    }
  }

  reset(): void {
    this.queue = [];
  }

  _flush(callback: () => void): void {
    this._flushQueue();
    callback();
  }
}