Skip to content

fix(spanner): complete query stream on PartialResultSet.last without resumeToken #9258

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

Merged
olavloite merged 2 commits into from
Sep 10, 2026
+441 −8
Merged

fix(spanner): complete query stream on PartialResultSet.last without resumeToken #9258

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cbf8ff3
perf(spanner): complete query stream immediately on PartialResultSet.…
olavloite Sep 8, 2026
41cd695
Merge branch 'main' into fix-partial-result-set-last
olavloite Sep 9, 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
45 changes: 38 additions & 7 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 @@ -249,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 All @@ -275,6 +275,11 @@

if (chunk.last) {
this.push(null);
// Calling next() notifies Node's stream machinery that processing of this
// chunk is complete on the Writable side of this Transform stream.
// This is a local, synchronous callback to Node's internal buffer; it does
// not block the event loop or wait for upstream network I/O or gRPC trailers.
next();
return;
}

Expand Down Expand Up @@ -667,7 +672,7 @@
const retryableCodes = [grpc.status.UNAVAILABLE];
const maxQueued = 10;
let lastResumeToken: ResumeToken;
let lastRequestStream: Readable;
let lastRequestStream: Readable | undefined;
let errorListener: (err: grpc.ServiceError) => void;
const startTime = Date.now();
const timeout = options?.gaxOptions?.timeout ?? Infinity;
Expand All @@ -686,10 +691,16 @@
// resume token, as that is an indication whether it is safe to retry the
// stream halfway.
let withoutCheckpointCount = 0;
let receivedLast = false;
const batchAndSplitOnTokenStream = new CheckpointStream({
maxQueued,
isCheckpointFn: (chunk: google.spanner.v1.PartialResultSet): boolean => {
const withCheckpoint = _hasResumeToken(chunk);
if (chunk.last) {
receivedLast = true;
destroyRequestStream();
requestsStream.end();
}
const withCheckpoint = _hasResumeToken(chunk) || Boolean(chunk.last);
if (withCheckpoint) {
withoutCheckpointCount = 0;
} else {
Expand All @@ -702,7 +713,13 @@
// This listener ensures that the last request that executed successfully
// after one or more retries will end the requestsStream.
const endListener = () => {
if (receivedLast) {
return;
}
setImmediate(() => {
if (receivedLast) {
return;
}
// Push a fake PartialResultSet without any values but with a resume token
// into the stream to ensure that the checkpoint stream is emptied, and
// then push `null` to end the stream.
Expand All @@ -713,11 +730,22 @@

const destroyRequestStream = (): void => {
if (lastRequestStream) {
lastRequestStream.removeListener('end', endListener);
lastRequestStream.removeAllListeners('error');
lastRequestStream.on('error', () => {});
lastRequestStream.unpipe(requestsStream);
lastRequestStream.destroy();
const streamToClean = lastRequestStream;
lastRequestStream = undefined;
streamToClean.removeListener('end', endListener);
if (errorListener) {
streamToClean.removeListener('error', errorListener);
}
streamToClean.on('error', () => {});
streamToClean.unpipe(requestsStream);
if (receivedLast) {
// Query completed successfully. Do not cancel the gRPC call; allow it
// to drain remaining trailers/EOF in the background so it is not marked
// CANCELLED by Spanner or Cloud Monitoring.
streamToClean.resume();

@sakthivelmanii sakthivelmanii Sep 10, 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.

Should we add some drain timeout here incase if trailers are not received for any reasons?

  const drainTimeout = setTimeout(() => streamToClean.destroy(), 5000);
  streamToClean.once('end', () => clearTimeout(drainTimeout));
  streamToClean.once('close', () => clearTimeout(drainTimeout));

@olavloite olavloite Sep 10, 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.

No, I don't think we should:

  1. Calling setTimeout(..) for each query introduces additional event loop congestion, which is exactly the kind of thing that we are trying to get rid of with this optimization (and other optimzations)
  2. The probability that the stream gets stuck after having sent a PartialResultSet with last=true, but before receiving the trailers, is extremely unlikely.
  3. In the extremely unlikely case that this happens, then the stream will eventually be cleaned up once the query timeout is reached (but that is 1hr by default)
  4. We do not add an extra timeout for the same optimization in Java and Go.
  5. The Node Firestore client in this monorepo also contains a similar optimization, and does also not add an additional timeout for receiving the trailers.

@sakthivelmanii sakthivelmanii Sep 10, 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.

My concern was if it stuck for some reasons and if we are not cleaning up the stream for next 1hr then in channel capacity will be reduced due to MAX_CONCURRENT_STREAMS(max concurrent streams is 100).

Since this is unlikely high, we can skip.

} else {
streamToClean.destroy();
}
}
};

Expand All @@ -728,6 +756,9 @@
lastRequestStream = requestFn(lastResumeToken);
lastRequestStream.on('end', endListener);
errorListener = (err: grpc.ServiceError) => {
if (receivedLast) {
return;
}
destroyRequestStream();
setImmediate(() => retry(err));
};
Expand Down
Loading
Loading