From 7bdeff8d27941c7dacc7321c385d9c697ecfe289 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Tue, 15 Sep 2026 15:11:50 +0200 Subject: [PATCH] chore(spanner): de-duplicate multiplexed session affinity gax options Snapshot#request and Snapshot#requestStream each contained their own copy of the same logic for adding the multiplexed session affinity key to the gax options of a request. Both now call a single private helper. The helper assigns the affinity key to the request descriptor instead of copying it first. That descriptor is built by the caller for one specific RPC, and prepareGapicRequest_ already modifies it in place, so the copy served no purpose. It removes one object allocation per request. The caller's own gax options are still never modified. --- handwritten/spanner/src/transaction.ts | 67 ++++++++++++++----------- handwritten/spanner/test/transaction.ts | 46 +++++++++++++++++ 2 files changed, 84 insertions(+), 29 deletions(-) diff --git a/handwritten/spanner/src/transaction.ts b/handwritten/spanner/src/transaction.ts index f84b73362eb..a41aeab27cf 100644 --- a/handwritten/spanner/src/transaction.ts +++ b/handwritten/spanner/src/transaction.ts @@ -438,35 +438,12 @@ export class Snapshot extends EventEmitter { }, }, }; - this.request = (config: any, callback?: Function) => { - let gaxOpts; - if (!config.gaxOpts || Object.keys(config.gaxOpts).length === 0) { - gaxOpts = this._bindGaxOpts as any; - } else { - gaxOpts = injectGaxOpt( - config.gaxOpts, - 'affinityKey', - this._affinityKey, - ); - } - config = Object.assign({}, config, {gaxOpts}); - return session.request(config, callback); - }; - - this.requestStream = (config: any) => { - let gaxOpts; - if (!config.gaxOpts || Object.keys(config.gaxOpts).length === 0) { - gaxOpts = this._bindGaxOpts as any; - } else { - gaxOpts = injectGaxOpt( - config.gaxOpts, - 'affinityKey', - this._affinityKey, - ); - } - config = Object.assign({}, config, {gaxOpts}); - return session.requestStream(config); - }; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + this.request = (config: any, callback?: Function) => + session.request(this._applyAffinityGaxOpts(config), callback); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + this.requestStream = (config: any) => + session.requestStream(this._applyAffinityGaxOpts(config)); } else { this.request = session.request.bind(session); this.requestStream = session.requestStream.bind(session); @@ -487,6 +464,38 @@ export class Snapshot extends EventEmitter { this._mutationKey = null; } + /** + * Binds the multiplexed session affinity key to the gax options of an + * outgoing request, so that all requests of this transaction are routed to + * the same gRPC channel. + * + * `config` is always a request descriptor that was freshly constructed by the + * caller for this one RPC (and {@link Spanner#prepareGapicRequest_} already + * modifies `config.headers` in place), so the affinity key is assigned + * directly instead of allocating a copy of the descriptor per request. + * + * @private + * + * @param {object} config The request configuration. + * @returns {object} The same request configuration. + */ + // eslint-disable-next-line @typescript-eslint/no-explicit-any + private _applyAffinityGaxOpts(config: any): any { + if (!config) { + return config; + } + if (!config.gaxOpts || Object.keys(config.gaxOpts).length === 0) { + config.gaxOpts = this._bindGaxOpts; + } else { + config.gaxOpts = injectGaxOpt( + config.gaxOpts, + 'affinityKey', + this._affinityKey, + ); + } + return config; + } + protected _updatePrecommitToken(resp: PrecommitTokenProvider): void { if ( this._latestPreCommitToken === null || diff --git a/handwritten/spanner/test/transaction.ts b/handwritten/spanner/test/transaction.ts index aad6c9517b4..647f6c763d8 100644 --- a/handwritten/spanner/test/transaction.ts +++ b/handwritten/spanner/test/transaction.ts @@ -171,6 +171,25 @@ describe('Transaction', () => { assert.strictEqual(REQUEST_STREAM.callCount, 1); }); + it('should keep `request` and `requestStream` usable when detached', () => { + REQUEST.resetHistory(); + REQUEST_STREAM.resetHistory(); + const multiplexedSession = Object.assign({}, SESSION, { + metadata: {multiplexed: true}, + }); + const txn = new Snapshot(multiplexedSession); + + // `TransactionRunner#_interceptErrors` and user code hold on to these + // methods without their receiver, so they must be pre-bound. + const {request, requestStream} = txn; + + request({client: 'SpannerClient'}, () => {}); + requestStream({client: 'SpannerClient'}); + + assert.strictEqual(REQUEST.callCount, 1); + assert.strictEqual(REQUEST_STREAM.callCount, 1); + }); + it('should generate _affinityKey for multiplexed sessions', () => { const multiplexedSession = Object.assign({}, SESSION, { metadata: {multiplexed: true}, @@ -214,6 +233,33 @@ describe('Transaction', () => { assert.deepStrictEqual(arg.gaxOpts, txn._bindGaxOpts); }); + it('should merge the affinity key into caller supplied gaxOpts', () => { + REQUEST.resetHistory(); + const multiplexedSession = Object.assign({}, SESSION, { + metadata: {multiplexed: true}, + }); + const txn = new Snapshot(multiplexedSession); + const gaxOpts = { + timeout: 1000, + otherArgs: {options: {unbind: true}}, + }; + + txn.request({client: 'SpannerClient', gaxOpts}, () => {}); + + const arg = REQUEST.lastCall.args[0]; + assert.deepStrictEqual(arg.gaxOpts, { + timeout: 1000, + otherArgs: { + options: {unbind: true, affinityKey: txn._affinityKey}, + }, + }); + // The caller supplied gax options must not be modified. + assert.deepStrictEqual(gaxOpts, { + timeout: 1000, + otherArgs: {options: {unbind: true}}, + }); + }); + it('should set the commonHeaders_', () => { assert.deepStrictEqual(snapshot.commonHeaders_, { [CLOUD_RESOURCE_HEADER]: snapshot.session.parent.formattedName_,