perf(spanner): distribute single-use queries across gRPC channel pool

handwritten/spanner/src/channel-factory.ts

@@ -0,0 +1,93 @@
+/*!
+ * Copyright 2026 Google LLC. All Rights Reserved.
+ *
+ * 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 {grpc} from 'google-gax';
+import grpcGcpModule = require('grpc-gcp');
+
+const grpcGcp = grpcGcpModule(grpc);
+
+export interface ChannelFactoryWithWithoutAffinity {
+  _withoutAffinity?: unknown;
+}
+
+/**
+ * Creates a delegate for GcpChannelFactory that overrides getAffinityConfig
+ * to return undefined. This causes grpc-gcp to skip affinity lookups and select
+ * a channel using its native stream load balancer (getActiveStreamsCount).
+ */
+export function createChannelFactoryWithoutAffinity(
+  channelFactory: object,
+): object {
+  const channelFactoryWithoutAffinity = Object.create(channelFactory);
+  (
+    channelFactoryWithoutAffinity as {getAffinityConfig?: () => undefined}
+  ).getAffinityConfig = () => undefined;
+  return channelFactoryWithoutAffinity;
+}
+
+/**
+ * Custom channel factory override that pre-allocates a static delegate
+ * without affinity lookup on the factory instance. This avoids any object
+ * or closure allocations per request.
+ */
+export function spannerChannelFactoryOverride(
+  address: string,
+  credentials: grpc.ChannelCredentials,
+  options: object,
+) {
+  const channelFactory = grpcGcp.gcpChannelFactoryOverride(
+    address,
+    credentials,
+    options,
+  );
+  if (channelFactory) {
+    (channelFactory as ChannelFactoryWithWithoutAffinity)._withoutAffinity =
+      createChannelFactoryWithoutAffinity(channelFactory);
+  }
+  return channelFactory;
+}
+
+interface SingleUseTransactionArgument {
+  transaction?: {
+    singleUse?: unknown;
+    single_use?: unknown;
+  };
+}
+
+/**
+ * Intercepts calls before dispatch. For single-use transactions (e.g. single queries),
+ * routes through the pre-allocated delegate to distribute across channels in the pool.
+ * For read/write and multi-use transactions, uses the standard channel factory so
+ * requests adhere to session-to-channel affinity.
+ */
+export function spannerCallInvocationTransformer<RequestType, ResponseType>(
+  callProperties: grpc.CallProperties<RequestType, ResponseType>,
+): grpc.CallProperties<RequestType, ResponseType> {
+  if (!callProperties) {
+    return callProperties;
+  }
+  const argument = callProperties.argument as
+    SingleUseTransactionArgument | undefined;
+  if (argument?.transaction?.singleUse || argument?.transaction?.single_use) {
+    const channelFactory =
+      callProperties.channel as ChannelFactoryWithWithoutAffinity;
+    if (channelFactory?._withoutAffinity) {
+      callProperties.channel =
+        channelFactory._withoutAffinity as typeof callProperties.channel;
+    }
+  }
+  return grpcGcp.gcpCallInvocationTransformer(callProperties);
+}

handwritten/spanner/src/index.ts

@@ -107,6 +107,11 @@
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const gcpApiConfig = require('./spanner_grpc_config.json');
 
+import {
+  spannerCallInvocationTransformer,
+  spannerChannelFactoryOverride,
+} from './channel-factory';
+
 export type IOperation = instanceAdmin.longrunning.IOperation;
 
 export type GetInstancesOptions = PagedOptionsWithFilter;
@@ -169,7 +174,7 @@
   >;
   observabilityOptions?: ObservabilityOptions;
   disableBuiltInMetrics?: boolean;
   interceptors?: any[];
   sessionLabels?: {[key: string]: string};
   /**
    * The Trusted Cloud Domain (TPC) DNS of the service used to make requests.
@@ -424,8 +429,8 @@
         // Add grpc keep alive setting
         'grpc.keepalive_time_ms': 120000,
         // Enable grpc-gcp support
-        'grpc.callInvocationTransformer': grpcGcp.gcpCallInvocationTransformer,
-        'grpc.channelFactoryOverride': grpcGcp.gcpChannelFactoryOverride,
+        'grpc.callInvocationTransformer': spannerCallInvocationTransformer,
+        'grpc.channelFactoryOverride': spannerChannelFactoryOverride,
         'grpc.gcpApiConfig': grpcGcp.createGcpApiConfig(gcpApiConfig),
         grpc,
       },
@@ -534,7 +539,7 @@
     if (!this.clients_.has(clientName)) {
       this.clients_.set(
         clientName,
         new v1[clientName](this.options as ClientOptions),
       );
     }
     return this.clients_.get(clientName)! as v1.InstanceAdminClient;
@@ -558,7 +563,7 @@
     if (!this.clients_.has(clientName)) {
       this.clients_.set(
         clientName,
         new v1[clientName](this.options as ClientOptions),
       );
     }
     return this.clients_.get(clientName)! as v1.DatabaseAdminClient;
@@ -615,9 +620,9 @@
 
     if (callback) {
       // process.nextTick prevents Unhandled Promise Rejections if callback throws
       res.then(
         () => process.nextTick(() => callback(null)),
         err => process.nextTick(() => callback(err)),
       );
     } else {
       return res;
@@ -1727,7 +1732,7 @@
       const clientName = config.client;
       try {
         if (!this.clients_.has(clientName)) {
           this.clients_.set(clientName, new v1[clientName](this.options));
         }
       } catch (err) {
         callback(err, null);
@@ -1773,7 +1778,7 @@
         // Attach the x-goog-spanner-request-id to the currently active span.
         attributeXGoogSpannerRequestIdToActiveSpan(config);
       }
       const interceptors: any[] = [];
       if (this._metricsEnabled) {
         interceptors.push(MetricInterceptor);
       }
@@ -1828,7 +1833,7 @@
             }
 
             return new Promise((resolve, reject) => {
               requestFn(...args)
                 .then(resolve)
                 .catch(err => {
                   injectRequestIDIntoError(config, err as Error);
@@ -1892,9 +1897,9 @@
           } else {
             const result = requestFn();
             if (result && typeof result.then === 'function') {
               result
                 .then(val => {
                   metricsTracer?.recordOperationCompletion();
                   resolve(val);
                 })
                 .catch(error => {

handwritten/spanner/test/channel-factory.ts

@@ -0,0 +1,241 @@
+/*!
+ * Copyright 2026 Google LLC. All Rights Reserved.
+ *
+ * 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 * as sinon from 'sinon';
+import * as proxyquire from 'proxyquire';
+import {grpc} from 'google-gax';
+import {
+  createChannelFactoryWithoutAffinity,
+  spannerCallInvocationTransformer,
+  spannerChannelFactoryOverride,
+  ChannelFactoryWithWithoutAffinity,
+} from '../src/channel-factory';
+
+describe('ChannelFactory and Transformer', () => {
+  let sandbox: sinon.SinonSandbox;
+
+  beforeEach(() => {
+    sandbox = sinon.createSandbox();
+  });
+
+  afterEach(() => {
+    sandbox.restore();
+  });
+
+  describe('createChannelFactoryWithoutAffinity', () => {
+    it('should create a prototype delegate that returns undefined for getAffinityConfig', () => {
+      const originalFactory: {
+        options: {foo: string};
+        getAffinityConfig(path?: string): {command: string} | undefined;
+      } = {
+        options: {foo: 'bar'},
+        getAffinityConfig() {
+          return {command: 'BOUND'};
+        },
+      };
+
+      const delegate = createChannelFactoryWithoutAffinity(
+        originalFactory,
+      ) as typeof originalFactory;
+
+      assert.strictEqual(delegate.getAffinityConfig('/test.Method'), undefined);
+      assert.strictEqual(
+        originalFactory.getAffinityConfig('/test.Method')?.command,
+        'BOUND',
+      );
+      assert.strictEqual(delegate.options, originalFactory.options);
+      assert.strictEqual(Object.getPrototypeOf(delegate), originalFactory);
+    });
+  });
+
+  describe('spannerChannelFactoryOverride', () => {
+    it('should attach _withoutAffinity delegate to the created channel factory', () => {
+      const channelFactory = spannerChannelFactoryOverride(
+        'localhost:443',
+        grpc.credentials.createInsecure(),
+        {},
+      ) as ChannelFactoryWithWithoutAffinity & {
+        getAffinityConfig: (path: string) => unknown;
+      };
+
+      assert.ok(channelFactory);
+      assert.ok(channelFactory._withoutAffinity);
+      const withoutAffinity = channelFactory._withoutAffinity as {
+        getAffinityConfig: (path: string) => unknown;
+      };
+      assert.strictEqual(typeof withoutAffinity.getAffinityConfig, 'function');
+      assert.strictEqual(
+        withoutAffinity.getAffinityConfig('/test.Method'),
+        undefined,
+      );
+    });
+
+    it('should return nullish channel factory when gcpChannelFactoryOverride returns nullish', () => {
+      const proxiedModule = proxyquire('../src/channel-factory', {
+        'grpc-gcp': () => ({
+          gcpChannelFactoryOverride: () => null,
+          gcpCallInvocationTransformer: (props: unknown) => props,
+        }),
+      });
+
+      const channelFactory = proxiedModule.spannerChannelFactoryOverride(
+        'localhost:443',
+        grpc.credentials.createInsecure(),
+        {},
+      );
+
+      assert.strictEqual(channelFactory, null);
+    });
+  });
+
+  describe('spannerCallInvocationTransformer', () => {
+    function createMockCallProperties(argument?: unknown, channel?: unknown) {
+      return {
+        argument,
+        channel: channel ?? {
+          _withoutAffinity: {name: 'withoutAffinityChannel'},
+        },
+      } as unknown as grpc.CallProperties<unknown, unknown>;
+    }
+
+    it('should route through _withoutAffinity for camelCase singleUse transaction', () => {
+      const mockChannelFactory = {
+        name: 'originalChannel',
+        _withoutAffinity: {name: 'withoutAffinityChannel'},
+      };
+      const callProps = createMockCallProperties(
+        {transaction: {singleUse: {readOnly: {}}}},
+        mockChannelFactory,
+      );
+
+      const transformed = spannerCallInvocationTransformer(callProps);
+      assert.strictEqual(
+        transformed.channel,
+        mockChannelFactory._withoutAffinity,
+      );
+    });
+
+    it('should route through _withoutAffinity for snake_case single_use transaction', () => {
+      const mockChannelFactory = {
+        name: 'originalChannel',
+        _withoutAffinity: {name: 'withoutAffinityChannel'},
+      };
+      const callProps = createMockCallProperties(
+        {transaction: {single_use: {read_only: {}}}},
+        mockChannelFactory,
+      );
+
+      const transformed = spannerCallInvocationTransformer(callProps);
+      assert.strictEqual(
+        transformed.channel,
+        mockChannelFactory._withoutAffinity,
+      );
+    });
+
+    it('should not change channel for multi-use transaction with id', () => {
+      const mockChannelFactory = {
+        name: 'originalChannel',
+        _withoutAffinity: {name: 'withoutAffinityChannel'},
+      };
+      const callProps = createMockCallProperties(
+        {transaction: {id: Buffer.from('tx-123')}},
+        mockChannelFactory,
+      );
+
+      const transformed = spannerCallInvocationTransformer(callProps);
+      assert.strictEqual(transformed.channel, mockChannelFactory);
+    });
+
+    it('should not change channel for read-write transaction with begin', () => {
+      const mockChannelFactory = {
+        name: 'originalChannel',
+        _withoutAffinity: {name: 'withoutAffinityChannel'},
+      };
+      const callProps = createMockCallProperties(
+        {transaction: {begin: {readWrite: {}}}},
+        mockChannelFactory,
+      );
+
+      const transformed = spannerCallInvocationTransformer(callProps);
+      assert.strictEqual(transformed.channel, mockChannelFactory);
+    });
+
+    it('should not change channel when argument is undefined', () => {
+      const mockChannelFactory = {
+        name: 'originalChannel',
+        _withoutAffinity: {name: 'withoutAffinityChannel'},
+      };
+      const callProps = createMockCallProperties(undefined, mockChannelFactory);
+
+      const transformed = spannerCallInvocationTransformer(callProps);
+      assert.strictEqual(transformed.channel, mockChannelFactory);
+    });
+
+    it('should not change channel when argument is null', () => {
+      const mockChannelFactory = {
+        name: 'originalChannel',
+        _withoutAffinity: {name: 'withoutAffinityChannel'},
+      };
+      const callProps = createMockCallProperties(null, mockChannelFactory);
+
+      const transformed = spannerCallInvocationTransformer(callProps);
+      assert.strictEqual(transformed.channel, mockChannelFactory);
+    });
+
+    it('should not change channel when argument has no transaction property', () => {
+      const mockChannelFactory = {
+        name: 'originalChannel',
+        _withoutAffinity: {name: 'withoutAffinityChannel'},
+      };
+      const callProps = createMockCallProperties(
+        {name: 'some-session'},
+        mockChannelFactory,
+      );
+
+      const transformed = spannerCallInvocationTransformer(callProps);
+      assert.strictEqual(transformed.channel, mockChannelFactory);
+    });
+
+    it('should fall back gracefully when _withoutAffinity is not present on channel', () => {
+      const mockChannelFactory = {
+        name: 'originalChannel',
+      };
+      const callProps = createMockCallProperties(
+        {transaction: {singleUse: {readOnly: {}}}},
+        mockChannelFactory,
+      );
+
+      const transformed = spannerCallInvocationTransformer(callProps);
+      assert.strictEqual(transformed.channel, mockChannelFactory);
+    });
+
+    it('should return callProperties directly when callProperties is undefined or null', () => {
+      assert.strictEqual(
+        spannerCallInvocationTransformer(
+          undefined as unknown as grpc.CallProperties<unknown, unknown>,
+        ),
+        undefined,
+      );
+      assert.strictEqual(
+        spannerCallInvocationTransformer(
+          null as unknown as grpc.CallProperties<unknown, unknown>,
+        ),
+        null,
+      );
+    });
+  });
+});