perf: optimize stream pipeline by eliminating events-intercept - #9221
Conversation
There was a problem hiding this comment.
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.
42b3354 to
d69a38a
Compare
d69a38a to
a4c89d2
Compare
| if (lastRequestStream) { | ||
| lastRequestStream.removeListener('end', endListener); | ||
| lastRequestStream.removeAllListeners('error'); | ||
| lastRequestStream.on('error', () => {}); // Prevent unhandled exception crash | ||
| lastRequestStream.destroy(); | ||
| } |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
since we are having:
const destroyRequestStream = (): void => {
if (lastRequestStream) {
lastRequestStream.removeListener('end', endListener);
lastRequestStream.removeAllListeners('error');
lastRequestStream.on('error', () => {});
lastRequestStream.unpipe(requestsStream);
lastRequestStream.destroy();
}
};
where we are listening on the error event to prevent uncaught exceptions as per lastRequestStream.on('error', () => {});. Now keeping this code to prevent uncaught exceptions is resulting in registering of a new event hence, assert.strictEqual(fakeStream.listenerCount('error'), 0, 'errorListener should be removed'); is failing since the error listener count is 1.
Currently, to keep the production safe I am asserting on value 1 in the test. But, can there be a better way?
| @@ -16,18 +16,15 @@ | |||
|
|
|||
| import {GrpcService} from './common-grpc/service'; | |||
| import * as checkpointStream from 'checkpoint-stream'; | |||
There was a problem hiding this comment.
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();
}
}There was a problem hiding this comment.
thanks for pointing this out! I have removed the checkpoint-stream package and have added our custom CheckpointStream class with some modification.
TL-DR:
The suggested class is synchronous and blocking. To make it asynchronous and non-blocking and to ensure it behaves like a production Node stream component I have done some modifications. I noticed few differences against production checkpoint-stream:
1. Synchronous vs. Asynchronous Flushing (Blocking the Event Loop)
In production: The original checkpoint-stream used a helper called split-array-stream which pushes items to the user asynchronously, one by one, using setImmediate(). This yielded control back to the event loop between each item.
Custom class: Our CheckpointStream flushes the queue synchronously in a while loop:
while (this.queue.length > 0) {
this.push(this.queue.shift());
}
This could block the event loop until all items are pushed.
Mitigation: To mitigate, I have implemented an asynchronous recursive loop (loop) using setImmediate() which matches behavior identical to split-array-stream.
2. Stream Flow Control and Backpressure
Custom class: In _transform, the callback() was called synchronously at the very end, regardless of whether a flush was happening asynchronously (or even synchronously).
_transform(...) {
// ... flush logic ...
callback(); // Called immediately
}
Mitigation: With the modification, we tied the callback directly to the completion of the flush. The callback is passed to _flushQueue and is only called when the queue will be completely empty.
if (!shouldFlush) {
return callback(); // Call immediately if not flushing
}
this._flushQueue(callback); // Pass callback to be called AFTER async flush
It will ensure that Node streams will be aware of when this transform step is done processing the current chunk, which will take care of flow control and backpressure.
3. Potential Data Loss on Errors
Custom class: Misses mechanism to handle shutdown on errors. If a non-retryable error occurred and standard stream.destroy(err) was called, any buffered data in this.queue would be dropped.
We added the flushAndDestroy(err) method:
flushAndDestroy(err: Error): void {
this._flushQueue(() => {
this.destroy(err);
});
}
It will guarantees that if a fatal error occurs, we first push all buffered rows to the user asynchronously, and only destroy the stream once the queue is empty. It will prevent prevents data loss right before a stream failure.
a4c89d2 to
575f998
Compare
|
/gcbrun |
f05f854 to
77509c8
Compare
77509c8 to
d1a1f7f
Compare
olavloite
left a comment
There was a problem hiding this comment.
This looks great! I think that this will significantly reduce the overhead for running streaming queries, and also cleans up the code quite a bit.
| * 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. |
There was a problem hiding this comment.
nit: I think that it would be good to just remove this section. In 6 months, no one will know what 'legacy behavior' is referring to.
There was a problem hiding this comment.
nit: I think that we can remove this dependency as well now (or move it to dev-dependencies, as it is currently used in a test, but I don't think that test is relevant anymore due to all the other changes here).
There was a problem hiding this comment.
thanks for pointing this out! I have removed the dependency and have updated the test accordingly.
Optimizes performance by removing the events-intercept library and its associated CPU penalty caused by monkey-patching EventEmitter.emit.
Changes:
Impact: Eliminates artificial CPU overhead on every stream event and simplifies pipeline architecture without breaking retry functionality.