From 9bcc699697030408ee0d635be316f33352ee1fae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Thu, 17 Sep 2026 09:43:16 +0200 Subject: [PATCH] perf: synchronous multiplexed session hand-off for non-stream callers Adds `getSessionSync()` to `MultiplexedSession` and `SessionFactory` to return the cached multiplexed session directly when available, avoiding an extra `process.nextTick` tick on every query. Non-stream callers (`database.run` and `Database.prototype.makePooledRequest_`) now use `getSessionSync()` to start query execution in the same tick on cache hit, falling back to asynchronous `getSession()` on cold start or when multiplexed sessions are disabled. Stream-returning callers (`runStream`, `batchWriteAtLeastOnce`, `makePooledStreamingRequest_`) keep asynchronous session acquisition so callers can attach event listeners before stream events are emitted. Also removes redundant callback wrapping in `SessionFactory.prototype.getSession` and moves `isAFEServerTimingEnabled` to `src/common.ts` to break a module load cycle between `common.ts` and `index.ts`. --- handwritten/spanner/src/common.ts | 5 +- handwritten/spanner/src/database.ts | 46 +++++- .../spanner/src/multiplexed-session.ts | 36 ++++- handwritten/spanner/src/session-factory.ts | 24 ++- handwritten/spanner/test/common-headers.ts | 82 +++++++++++ handwritten/spanner/test/database.ts | 139 +++++++++++++++++- .../spanner/test/multiplexed-session.ts | 12 ++ handwritten/spanner/test/session-factory.ts | 42 ++++++ 8 files changed, 362 insertions(+), 24 deletions(-) create mode 100644 handwritten/spanner/test/common-headers.ts diff --git a/handwritten/spanner/src/common.ts b/handwritten/spanner/src/common.ts index eed86499e4a6..ecfe2d1df739 100644 --- a/handwritten/spanner/src/common.ts +++ b/handwritten/spanner/src/common.ts @@ -18,7 +18,6 @@ import {grpc, CallOptions, Operation as GaxOperation} from 'google-gax'; import {protos} from '@google-cloud/spanner-api'; import instanceAdmin = protos.google; import databaseAdmin = protos.google; -import {Spanner} from '.'; export type IOperation = instanceAdmin.longrunning.IOperation; @@ -113,13 +112,13 @@ export function getCommonHeaders( const headers: {[k: string]: string} = {}; if ( - process.env.SPANNER_ENABLE_END_TO_END_TRACING === 'true' || + process.env.SPANNER_ENABLE_END_TO_END_TRACING?.toLowerCase() === 'true' || enableTracing ) { headers[END_TO_END_TRACING_HEADER] = 'true'; } - if (Spanner.isAFEServerTimingEnabled()) { + if (process.env.SPANNER_DISABLE_AFE_SERVER_TIMING?.toLowerCase() !== 'true') { headers[AFE_SERVER_TIMING_HEADER] = 'true'; } diff --git a/handwritten/spanner/src/database.ts b/handwritten/spanner/src/database.ts index b6ba8347e4c5..1eee2d4b8add 100644 --- a/handwritten/spanner/src/database.ts +++ b/handwritten/spanner/src/database.ts @@ -38,7 +38,11 @@ import { } from 'google-gax'; import {Backup} from './backup'; import {BatchTransaction, TransactionIdentifier} from './batch-transaction'; -import {SessionFactory, SessionFactoryInterface} from './session-factory'; +import { + GetSessionCallback, + SessionFactory, + SessionFactoryInterface, +} from './session-factory'; import {protos} from '@google-cloud/spanner-api'; import google = protos.google; import databaseAdmin = protos.google; @@ -2531,7 +2535,7 @@ class Database extends common.GrpcServiceObject { callback?: PoolRequestCallback, ): void | Promise { const sessionFactory_ = this.sessionFactory_; - sessionFactory_.getSessionForReadWrite((err, session) => { + const onSession: GetSessionCallback = (err, session) => { if (err) { callback!(err as ServiceError, null); return; @@ -2544,7 +2548,17 @@ class Database extends common.GrpcServiceObject { sessionFactory_.release(session!); callback!(err, ...args); }); - }); + }; + + const session = sessionFactory_.isMultiplexedEnabledForRW?.() + ? sessionFactory_.getSessionSync?.() + : null; + if (session) { + onSession(null, session); + return; + } + + sessionFactory_.getSessionForReadWrite(onSession); } /** @@ -3025,16 +3039,30 @@ class Database extends common.GrpcServiceObject { callback!(error, rows, stats!, metadata!); }; - this.sessionFactory_.getSession((error, session) => { + const onSession: GetSessionCallback = (error, session) => { if (error) { complete(error as grpc.ServiceError); return; } streamSpan.addEvent('Using Session', {'session.id': session?.id}); - snapshot = session!.snapshot(options, this.queryOptions_); - this._runOnSnapshot(snapshot, session!, query, complete); - }); + try { + snapshot = session!.snapshot(options, this.queryOptions_); + this._runOnSnapshot(snapshot, session!, query, complete); + } catch (syncError) { + // Defer error delivery via nextTick so callback callers never experience + // synchronous callback execution (Zalgo) when getSessionSync() returns synchronously. + process.nextTick(() => complete(syncError as grpc.ServiceError)); + } + }; + + const session = this.sessionFactory_.getSessionSync?.(); + if (session) { + onSession(null, session); + return; + } + + this.sessionFactory_.getSession(onSession); } /** @@ -3081,7 +3109,9 @@ class Database extends common.GrpcServiceObject { snapshot.run(query, callback as RunCallback); } } catch (syncError) { - callback(syncError as grpc.ServiceError); + // Defer error delivery via nextTick so callback callers never experience + // synchronous callback execution (Zalgo) when getSessionSync() returns synchronously. + process.nextTick(() => callback(syncError as grpc.ServiceError)); } } /** diff --git a/handwritten/spanner/src/multiplexed-session.ts b/handwritten/spanner/src/multiplexed-session.ts index fc8fd959918d..a9eadce6c56c 100644 --- a/handwritten/spanner/src/multiplexed-session.ts +++ b/handwritten/spanner/src/multiplexed-session.ts @@ -51,6 +51,13 @@ export interface MultiplexedSessionInterface extends EventEmitter { * @param {GetSessionCallback} callback The callback function. */ getSession(callback: GetSessionCallback): void; + + /** + * When called returns a cached multiplexed session synchronously if available. + * + * @name MultiplexedSessionInterface#getSessionSync + */ + getSessionSync(): Session | null; } /** @@ -195,6 +202,21 @@ export class MultiplexedSession this._refreshHandle.unref(); } + /** + * Synchronously returns the cached multiplexed session if available, + * or null if no session is currently cached. + * + * @returns {Session|null} The cached multiplexed session or null. + */ + getSessionSync(): Session | null { + if (this._multiplexedSession !== null) { + const span = getActiveOrNoopSpan(); + span.addEvent('Cache hit: has usable multiplexed session'); + return this._multiplexedSession; + } + return null; + } + /** * Retrieves a session asynchronously and invokes a callback with the session details. * Note: The callback receives `(null, session)`. To prevent unnecessary allocations on @@ -207,10 +229,8 @@ export class MultiplexedSession * */ getSession(callback: GetSessionCallback): void { - if (this._multiplexedSession !== null) { - const session = this._multiplexedSession; - const span = getActiveOrNoopSpan(); - span.addEvent('Cache hit: has usable multiplexed session'); + const session = this.getSessionSync(); + if (session !== null) { // Use process.nextTick to guarantee asynchronous callback execution ("never release Zalgo"). // This avoids microtask and Promise allocation overhead while preventing race conditions // where callers (such as Database.prototype.runStream) need to return their stream and @@ -249,13 +269,13 @@ export class MultiplexedSession * */ async _getSession(): Promise { - const span = getActiveOrNoopSpan(); // Check if the multiplexed session is already available - if (this._multiplexedSession !== null) { - span.addEvent('Cache hit: has usable multiplexed session'); - return this._multiplexedSession; + const cachedSession = this.getSessionSync(); + if (cachedSession !== null) { + return cachedSession; } + const span = getActiveOrNoopSpan(); span.addEvent('Waiting for a multiplexed session to become available'); // If initialization is ALREADY in progress, join the existing line! diff --git a/handwritten/spanner/src/session-factory.ts b/handwritten/spanner/src/session-factory.ts index 2fe7fa6679bb..603b92658230 100644 --- a/handwritten/spanner/src/session-factory.ts +++ b/handwritten/spanner/src/session-factory.ts @@ -58,6 +58,13 @@ export interface SessionFactoryInterface { */ getSession(callback: GetSessionCallback): void; + /** + * When called returns a cached multiplexed session synchronously if available. + * + * @name SessionFactoryInterface#getSessionSync + */ + getSessionSync(): Session | null; + /** * When called returns a session for paritioned dml. * @@ -168,6 +175,18 @@ export class SessionFactory } } + /** + * Synchronously returns a cached multiplexed session if multiplexed sessions + * are enabled and one is available, otherwise null. + * + * @returns {Session|null} The cached multiplexed session or null. + */ + getSessionSync(): Session | null { + return this.isMultiplexed + ? (this.multiplexedSession_?.getSessionSync?.() ?? null) + : null; + } + /** * Retrieves a session, either a regular session or a multiplexed session, based on the environment variable configuration. * @@ -176,15 +195,12 @@ export class SessionFactory * * @param {GetSessionCallback} callback The callback function. */ - getSession(callback: GetSessionCallback): void { const sessionHandler = this.isMultiplexed ? this.multiplexedSession_ : this.pool_; - sessionHandler!.getSession((err, session, transaction) => - callback(err, session, transaction), - ); + sessionHandler!.getSession(callback); } /** diff --git a/handwritten/spanner/test/common-headers.ts b/handwritten/spanner/test/common-headers.ts new file mode 100644 index 000000000000..c75ca45273b8 --- /dev/null +++ b/handwritten/spanner/test/common-headers.ts @@ -0,0 +1,82 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import * as assert from 'assert'; +import {describe, it, beforeEach, afterEach} from 'mocha'; +import { + AFE_SERVER_TIMING_HEADER, + CLOUD_RESOURCE_HEADER, + END_TO_END_TRACING_HEADER, + getCommonHeaders, +} from '../src/common'; + +describe('getCommonHeaders', () => { + const resource = 'projects/p/instances/i/databases/d'; + const originalDisableAfe = process.env.SPANNER_DISABLE_AFE_SERVER_TIMING; + const originalEnableE2E = process.env.SPANNER_ENABLE_END_TO_END_TRACING; + + beforeEach(() => { + delete process.env.SPANNER_DISABLE_AFE_SERVER_TIMING; + delete process.env.SPANNER_ENABLE_END_TO_END_TRACING; + }); + + afterEach(() => { + if (originalDisableAfe !== undefined) { + process.env.SPANNER_DISABLE_AFE_SERVER_TIMING = originalDisableAfe; + } else { + delete process.env.SPANNER_DISABLE_AFE_SERVER_TIMING; + } + if (originalEnableE2E !== undefined) { + process.env.SPANNER_ENABLE_END_TO_END_TRACING = originalEnableE2E; + } else { + delete process.env.SPANNER_ENABLE_END_TO_END_TRACING; + } + }); + + it('should include resource prefix and AFE timing header by default', () => { + const headers = getCommonHeaders(resource); + assert.strictEqual(headers[CLOUD_RESOURCE_HEADER], resource); + assert.strictEqual(headers[AFE_SERVER_TIMING_HEADER], 'true'); + assert.strictEqual(headers[END_TO_END_TRACING_HEADER], undefined); + }); + + it('should omit AFE timing header when disabled via case-insensitive env var', () => { + process.env.SPANNER_DISABLE_AFE_SERVER_TIMING = 'TRUE'; + const headers = getCommonHeaders(resource); + assert.strictEqual(headers[AFE_SERVER_TIMING_HEADER], undefined); + }); + + it('should include AFE timing header when env var is false', () => { + process.env.SPANNER_DISABLE_AFE_SERVER_TIMING = 'false'; + const headers = getCommonHeaders(resource); + assert.strictEqual(headers[AFE_SERVER_TIMING_HEADER], 'true'); + }); + + it('should include end-to-end tracing header when enabled via case-insensitive env var', () => { + process.env.SPANNER_ENABLE_END_TO_END_TRACING = 'True'; + const headers = getCommonHeaders(resource); + assert.strictEqual(headers[END_TO_END_TRACING_HEADER], 'true'); + }); + + it('should include end-to-end tracing header when enableTracing option is true', () => { + const headers = getCommonHeaders(resource, true); + assert.strictEqual(headers[END_TO_END_TRACING_HEADER], 'true'); + }); + + it('should omit end-to-end tracing header when env var is false and enableTracing option is false', () => { + process.env.SPANNER_ENABLE_END_TO_END_TRACING = 'FALSE'; + const headers = getCommonHeaders(resource, false); + assert.strictEqual(headers[END_TO_END_TRACING_HEADER], undefined); + }); +}); diff --git a/handwritten/spanner/test/database.ts b/handwritten/spanner/test/database.ts index 1539b572de0e..94d846de3bae 100644 --- a/handwritten/spanner/test/database.ts +++ b/handwritten/spanner/test/database.ts @@ -141,6 +141,9 @@ export class FakeMultiplexedSession extends EventEmitter { } createSession() {} getSession() {} + getSessionSync(): FakeSession | null { + return null; + } } export class FakeSessionFactory extends EventEmitter { @@ -150,6 +153,9 @@ export class FakeSessionFactory extends EventEmitter { this.calledWith_ = arguments; } getSession() {} + getSessionSync(): FakeSession | null { + return null; + } getSessionForPartitionedOps() {} getSessionForReadWrite() {} getPool(): FakeSessionPool { @@ -1626,6 +1632,9 @@ describe('Database', () => { callback(null, SESSION); }; + SESSIONFACTORY.isMultiplexedEnabledForRW = () => false; + SESSIONFACTORY.getSessionSync = () => null; + SESSIONFACTORY.release = util.noop; }); @@ -1693,6 +1702,51 @@ describe('Database', () => { done(); }); }); + + it('should use synchronous multiplexed session hand-off when available', done => { + const getSessionForReadWriteStub = sandbox.stub( + SESSIONFACTORY, + 'getSessionForReadWrite', + ); + sandbox.stub(SESSIONFACTORY, 'isMultiplexedEnabledForRW').returns(true); + sandbox.stub(SESSIONFACTORY, 'getSessionSync').returns(SESSION); + + let requestCalledSynchronously = false; + database.request = (config, callback) => { + requestCalledSynchronously = true; + assert.strictEqual(config.reqOpts.session, SESSION.formattedName_); + callback(null, 'response'); + }; + + database.makePooledRequest_(CONFIG, (err, res) => { + assert.ifError(err); + assert.strictEqual(res, 'response'); + assert.strictEqual(getSessionForReadWriteStub.callCount, 0); + done(); + }); + assert.strictEqual(requestCalledSynchronously, true); + }); + + it('should fall back to getSessionForReadWrite when getSessionSync returns null', done => { + const getSessionForReadWriteSpy = sandbox.spy( + SESSIONFACTORY, + 'getSessionForReadWrite', + ); + sandbox.stub(SESSIONFACTORY, 'isMultiplexedEnabledForRW').returns(true); + sandbox.stub(SESSIONFACTORY, 'getSessionSync').returns(null); + + database.request = (config, callback) => { + assert.strictEqual(config.reqOpts.session, SESSION.formattedName_); + callback(null, 'response'); + }; + + database.makePooledRequest_(CONFIG, (err, res) => { + assert.ifError(err); + assert.strictEqual(res, 'response'); + assert.strictEqual(getSessionForReadWriteSpy.callCount, 1); + done(); + }); + }); }); describe('makePooledStreamingRequest_', () => { @@ -2076,6 +2130,26 @@ describe('Database', () => { }); }); + it('should use synchronous multiplexed session hand-off via getSessionSync when available', done => { + sandbox.stub(fakeSessionFactory, 'getSessionSync').returns(fakeSession); + + let snapshotCreatedSynchronously = false; + snapshotStub.callsFake(() => { + snapshotCreatedSynchronously = true; + return fakeSnapshot; + }); + + database.run(QUERY, (err, rows) => { + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 1}]); + assert.strictEqual(getSessionStub.callCount, 0); + assert.strictEqual(snapshotStub.callCount, 1); + done(); + }); + + assert.strictEqual(snapshotCreatedSynchronously, true); + }); + it('should fall back to streaming path when multiplexed session is disabled', done => { (fakeSessionFactory.isMultiplexedEnabled as sinon.SinonStub).returns( false, @@ -2121,6 +2195,15 @@ describe('Database', () => { assert.strictEqual(runStub.callCount, 1); }); + it('should use synchronous multiplexed session hand-off with Promise-based run', async () => { + sandbox.stub(fakeSessionFactory, 'getSessionSync').returns(fakeSession); + const runPromise = pfy.promisify(database.run.bind(database)); + const [rows] = await runPromise(QUERY); + assert.deepStrictEqual(rows, [{id: 1}]); + assert.strictEqual(getSessionStub.callCount, 0); + assert.strictEqual(snapshotStub.callCount, 1); + }); + it('should fall back to snapshot.run when snapshot.runStream is overridden', done => { fakeSnapshot.runStream = () => through.obj() as any; const snapshotRunStub = sandbox.stub(fakeSnapshot, '_run'); @@ -2135,17 +2218,71 @@ describe('Database', () => { }); }); - it('should catch synchronous error in runMethod, end snapshot and propagate error', done => { + it('should fall back to getSession when getSessionSync returns null', done => { + sandbox.stub(fakeSessionFactory, 'getSessionSync').returns(null); + + database.run(QUERY, (err, rows) => { + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 1}]); + assert.strictEqual(getSessionStub.callCount, 1); + assert.strictEqual(snapshotStub.callCount, 1); + done(); + }); + }); + + it('should catch synchronous error in runMethod, end snapshot and propagate error asynchronously', done => { + sandbox.stub(fakeSessionFactory, 'getSessionSync').returns(fakeSession); const syncError = new Error('Synchronous parameter failure'); const endStub = sandbox.stub(fakeSnapshot, 'end'); runStub.throws(syncError); + let isSynchronous = true; database.run(QUERY, (err, rows) => { + assert.strictEqual(isSynchronous, false); assert.strictEqual(err, syncError); assert.deepStrictEqual(rows, []); assert.strictEqual(endStub.callCount, 1); done(); }); + isSynchronous = false; + }); + + it('should catch synchronous error in session.snapshot and propagate error asynchronously', done => { + sandbox.stub(fakeSessionFactory, 'getSessionSync').returns(fakeSession); + const snapshotError = new Error('Invalid timestamp bounds'); + snapshotStub.throws(snapshotError); + + let isSynchronous = true; + database.run(QUERY, (err, rows) => { + assert.strictEqual(isSynchronous, false); + assert.strictEqual(err, snapshotError); + assert.deepStrictEqual(rows, []); + done(); + }); + isSynchronous = false; + }); + + it('should dispatch query synchronously on getSessionSync fast-path while invoking callback asynchronously', done => { + sandbox.stub(fakeSessionFactory, 'getSessionSync').returns(fakeSession); + let queryDispatchedSynchronously = false; + runStub.callsFake((query, optionsOrCallback, cb) => { + queryDispatchedSynchronously = true; + const callback = + typeof optionsOrCallback === 'function' ? optionsOrCallback : cb; + if (callback) { + process.nextTick(() => callback(null, [{id: 1}])); + } + }); + + let isSynchronous = true; + database.run(QUERY, (err, rows) => { + assert.strictEqual(isSynchronous, false); + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 1}]); + done(); + }); + assert.strictEqual(queryDispatchedSynchronously, true); + isSynchronous = false; }); }); diff --git a/handwritten/spanner/test/multiplexed-session.ts b/handwritten/spanner/test/multiplexed-session.ts index 656b0794cf35..5c75b635f38c 100644 --- a/handwritten/spanner/test/multiplexed-session.ts +++ b/handwritten/spanner/test/multiplexed-session.ts @@ -169,6 +169,18 @@ describe('MultiplexedSession', () => { }); }); + describe('getSessionSync', () => { + it('should return null when no session is cached', () => { + multiplexedSession._multiplexedSession = null; + assert.strictEqual(multiplexedSession.getSessionSync(), null); + }); + + it('should synchronously return the cached session when available', () => { + multiplexedSession._multiplexedSession = fakeMuxSession; + assert.strictEqual(multiplexedSession.getSessionSync(), fakeMuxSession); + }); + }); + describe('getSession', () => { let restoreProcessListeners: (() => void) | null = null; diff --git a/handwritten/spanner/test/session-factory.ts b/handwritten/spanner/test/session-factory.ts index 476b471d402c..6ea047c65923 100644 --- a/handwritten/spanner/test/session-factory.ts +++ b/handwritten/spanner/test/session-factory.ts @@ -177,6 +177,48 @@ describe('SessionFactory', () => { }); }); + describe('getSessionSync', () => { + describe('when multiplexed session is disabled', () => { + before(() => { + process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false'; + }); + + after(() => { + delete process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS; + }); + + it('should return null even if multiplexedSession_ has a session', () => { + const factory = new SessionFactory(DATABASE, NAME, POOL_OPTIONS); + sandbox + .stub(factory.multiplexedSession_, 'getSessionSync') + .returns(fakeMuxSession); + assert.strictEqual(factory.getSessionSync(), null); + }); + }); + + describe('when multiplexed session is default', () => { + it('should return the cached multiplexed session when available', () => { + sandbox + .stub(sessionFactory.multiplexedSession_, 'getSessionSync') + .returns(fakeMuxSession); + assert.strictEqual(sessionFactory.getSessionSync(), fakeMuxSession); + }); + + it('should return null when no multiplexed session is cached', () => { + sandbox + .stub(sessionFactory.multiplexedSession_, 'getSessionSync') + .returns(null); + assert.strictEqual(sessionFactory.getSessionSync(), null); + }); + + it('should safely return null if multiplexedSession_ does not implement getSessionSync', () => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + sessionFactory.multiplexedSession_ = {} as any; + assert.strictEqual(sessionFactory.getSessionSync(), null); + }); + }); + }); + describe('getSession', () => { describe('when multiplexed session is disabled', () => { before(() => {