feat(spanner): make affinity keys configurable via spanner_grpc_confi…

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;

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 || {},

handwritten/spanner/src/spanner_grpc_config.json

@@ -90,3 +90,4 @@
     }
   ]
 }
+

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<string, any>();
+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<Row | Json>;
 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) => {

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', () => {

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();
       });
     });