diff --git a/handwritten/spanner/src/database.ts b/handwritten/spanner/src/database.ts index 145018db064..0666839b787 100644 --- a/handwritten/spanner/src/database.ts +++ b/handwritten/spanner/src/database.ts @@ -471,6 +471,7 @@ class Database extends common.GrpcServiceObject { this.formattedName_ = formattedName_; this.instance = instance; + this.spanner = instance.parent as Spanner; const poolOpts = typeof poolOptions === 'object' ? poolOptions : null; this.databaseRole = databaseRole || poolOpts?.databaseRole || null; diff --git a/handwritten/spanner/src/index.ts b/handwritten/spanner/src/index.ts index 50d393249fe..68babb21945 100644 --- a/handwritten/spanner/src/index.ts +++ b/handwritten/spanner/src/index.ts @@ -424,7 +424,8 @@ class Spanner extends GrpcService { // Enable grpc-gcp support 'grpc.callInvocationTransformer': grpcGcp.gcpCallInvocationTransformer, 'grpc.channelFactoryOverride': grpcGcp.gcpChannelFactoryOverride, - 'grpc.gcpApiConfig': grpcGcp.createGcpApiConfig(gcpApiConfig), + // Bypass createGcpApiConfig to preserve our custom metadata fields + 'grpc.gcpApiConfig': gcpApiConfig, grpc, }, options || {}, diff --git a/handwritten/spanner/src/spanner_grpc_config.json b/handwritten/spanner/src/spanner_grpc_config.json index bd640bdbe0b..426aa617d5e 100644 --- a/handwritten/spanner/src/spanner_grpc_config.json +++ b/handwritten/spanner/src/spanner_grpc_config.json @@ -90,3 +90,4 @@ } ] } + diff --git a/handwritten/spanner/src/transaction.ts b/handwritten/spanner/src/transaction.ts index 75d4b2d0079..1a81462917c 100644 --- a/handwritten/spanner/src/transaction.ts +++ b/handwritten/spanner/src/transaction.ts @@ -55,6 +55,21 @@ import { } from './instrument'; import {RunTransactionOptions} from './transaction-runner'; import {injectRequestIDIntoHeaders, nextNthRequest} from './request_id_header'; +import * as uuid from 'uuid'; + +const gcpApiConfig = require('./spanner_grpc_config.json'); + +// Pre-compute a map for O(1) affinity lookups +const methodToAffinityMap = new Map(); +if (gcpApiConfig && gcpApiConfig.method) { + gcpApiConfig.method.forEach((m: any) => { + if (m.name && m.affinity) { + m.name.forEach((name: string) => { + methodToAffinityMap.set(name, m.affinity); + }); + } + }); +} export type Rows = Array; const RETRY_INFO_TYPE = 'type.googleapis.com/google.rpc.retryinfo'; @@ -296,6 +311,7 @@ export class Snapshot extends EventEmitter { | undefined | null; id?: Uint8Array | string; + public _affinityKey?: string; multiplexedSessionPreviousTransactionId?: Uint8Array | string; ended: boolean; metadata?: spannerClient.spanner.v1.ITransaction; @@ -365,8 +381,49 @@ export class Snapshot extends EventEmitter { this.ended = false; this.session = session; this.queryOptions = Object.assign({}, queryOptions); - this.request = session.request.bind(session); - this.requestStream = session.requestStream.bind(session); + // If the session is multiplexed, generate a unique affinity key (UUID) for this + // specific transaction/snapshot. This allows requests using the same shared + // multiplexed session to be distributed across different gRPC channels. + if (session.metadata && session.metadata.multiplexed) { + this._affinityKey = uuid.v4(); + } + this.request = (config: any, callback: Function) => { + if (this._affinityKey) { + config = { + ...config, + gaxOpts: { + ...(config.gaxOpts || {}), + otherArgs: { + ...(config.gaxOpts?.otherArgs || {}), + options: { + ...(config.gaxOpts?.otherArgs?.options || {}), + affinityKey: this._affinityKey, + }, + }, + }, + }; + } + return session.request(config, callback); + }; + + this.requestStream = (config: any) => { + if (this._affinityKey) { + config = { + ...config, + gaxOpts: { + ...(config.gaxOpts || {}), + otherArgs: { + ...(config.gaxOpts?.otherArgs || {}), + options: { + ...(config.gaxOpts?.otherArgs?.options || {}), + affinityKey: this._affinityKey, + }, + }, + }, + }; + } + return session.requestStream(config); + }; const readOnly = Snapshot.encodeTimestampBounds(options || {}); this._options = {readOnly}; @@ -1031,6 +1088,24 @@ export class Snapshot extends EventEmitter { this.ended = true; process.nextTick(() => this.emit('end')); + + if (this._affinityKey) { + const database = this.session.parent as Database; + const spanner = database.spanner; + const client = spanner.clients_.get('SpannerClient') as any; + if (client && client.spannerStub) { + client.spannerStub + .then((stub: any) => { + if (stub && typeof stub.getChannel === 'function') { + const channel = stub.getChannel(); + if (channel && typeof channel.unbind === 'function') { + channel.unbind(this._affinityKey); + } + } + }) + .catch(() => {}); + } + } } /** @@ -2456,12 +2531,25 @@ export class Transaction extends Dml { span.addEvent('Starting Commit'); const database = this.session.parent as Database; + let newGaxOpts = gaxOpts; + if (this._affinityKey) { + newGaxOpts = Object.assign({}, gaxOpts, { + otherArgs: { + ...((gaxOpts as any)?.otherArgs || {}), + options: { + ...((gaxOpts as any)?.otherArgs?.options || {}), + unbind: true, + }, + }, + }); + } + this.request( { client: 'SpannerClient', method: 'commit', reqOpts, - gaxOpts: gaxOpts, + gaxOpts: newGaxOpts, headers: injectRequestIDIntoHeaders( headers, this.session, @@ -2819,12 +2907,25 @@ export class Transaction extends Dml { addLeaderAwareRoutingHeader(headers); } + let newGaxOpts = gaxOpts; + if (this._affinityKey) { + newGaxOpts = Object.assign({}, gaxOpts, { + otherArgs: { + ...((gaxOpts as any)?.otherArgs || {}), + options: { + ...((gaxOpts as any)?.otherArgs?.options || {}), + unbind: true, + }, + }, + }); + } + this.request( { client: 'SpannerClient', method: 'rollback', reqOpts, - gaxOpts, + gaxOpts: newGaxOpts, headers: headers, }, (err: null | ServiceError) => { diff --git a/handwritten/spanner/test/index.ts b/handwritten/spanner/test/index.ts index 9a47f920d99..f2e57fd7e56 100644 --- a/handwritten/spanner/test/index.ts +++ b/handwritten/spanner/test/index.ts @@ -232,9 +232,7 @@ describe('Spanner', () => { 'grpc.callInvocationTransformer': fakeGrpcGcp().gcpCallInvocationTransformer, 'grpc.channelFactoryOverride': fakeGrpcGcp().gcpChannelFactoryOverride, - 'grpc.gcpApiConfig': { - calledWith_: apiConfig, - }, + 'grpc.gcpApiConfig': apiConfig, }); it('should localize a cached gapic client map', () => { diff --git a/handwritten/spanner/test/multiplexed-session.ts b/handwritten/spanner/test/multiplexed-session.ts index d21f912f5e8..ba934dd25d5 100644 --- a/handwritten/spanner/test/multiplexed-session.ts +++ b/handwritten/spanner/test/multiplexed-session.ts @@ -43,7 +43,11 @@ describe('MultiplexedSession', () => { return Object.assign(new Session(DATABASE, name), props, { create: sandbox.stub().resolves(), - transaction: sandbox.stub().returns(new FakeTransaction()), + transaction: sandbox.stub().callsFake(() => { + const txn = new FakeTransaction(); + (txn as any)._affinityKey = 'mock-uuid'; + return txn; + }), }); }; @@ -185,13 +189,15 @@ describe('MultiplexedSession', () => { }); }); - it('should pass back the session and txn', done => { - const fakeTxn = new FakeTransaction() as unknown as Transaction; + it('should pass back the session and txn with affinity key', done => { sandbox.stub(multiplexedSession, '_getSession').resolves(fakeMuxSession); multiplexedSession.getSession((err, session, txn) => { assert.ifError(err); assert.strictEqual(session, fakeMuxSession); - assert.deepStrictEqual(txn, fakeTxn); + assert(txn); + assert(txn._affinityKey); + assert.strictEqual(typeof txn._affinityKey, 'string'); + assert(txn._affinityKey.length > 0); done(); }); });