From 4d33e49a7925ec49bcc1f618f0928244ba5811c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 21 Sep 2026 11:55:47 +0200 Subject: [PATCH] chore(spanner): add a channel pool prototype --- .../spanner/src/channel-pool/affinity.ts | 53 + .../src/channel-pool/channel-adapter.ts | 116 ++ .../spanner/src/channel-pool/dynamic-pool.ts | 587 ++++++ handwritten/spanner/src/channel-pool/index.ts | 55 + handwritten/spanner/src/channel-pool/p2c.ts | 57 + .../spanner/src/channel-pool/static-pool.ts | 192 ++ .../spanner/src/channel-pool/transformer.ts | 132 ++ handwritten/spanner/src/channel-pool/types.ts | 138 ++ handwritten/spanner/src/index.ts | 189 +- .../src/metrics/metrics-tracer-factory.ts | 56 +- .../spanner/src/multiplexed-session.ts | 7 + handwritten/spanner/src/transaction-runner.ts | 3 + handwritten/spanner/src/transaction.ts | 57 +- handwritten/spanner/test/channel-pool.ts | 1796 +++++++++++++++++ handwritten/spanner/test/index.ts | 7 + .../test/metrics/metrics-tracer-factory.ts | 19 + handwritten/spanner/test/spanner.ts | 21 +- .../spanner/test/transaction-runner.ts | 14 +- handwritten/spanner/test/transaction.ts | 56 +- 19 files changed, 3525 insertions(+), 30 deletions(-) create mode 100644 handwritten/spanner/src/channel-pool/affinity.ts create mode 100644 handwritten/spanner/src/channel-pool/channel-adapter.ts create mode 100644 handwritten/spanner/src/channel-pool/dynamic-pool.ts create mode 100644 handwritten/spanner/src/channel-pool/index.ts create mode 100644 handwritten/spanner/src/channel-pool/p2c.ts create mode 100644 handwritten/spanner/src/channel-pool/static-pool.ts create mode 100644 handwritten/spanner/src/channel-pool/transformer.ts create mode 100644 handwritten/spanner/src/channel-pool/types.ts create mode 100644 handwritten/spanner/test/channel-pool.ts diff --git a/handwritten/spanner/src/channel-pool/affinity.ts b/handwritten/spanner/src/channel-pool/affinity.ts new file mode 100644 index 000000000000..699daed6ff8e --- /dev/null +++ b/handwritten/spanner/src/channel-pool/affinity.ts @@ -0,0 +1,53 @@ +/*! + * 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 {ChannelEntry} from './types'; + +/** + * Kind of affinity required by an operation. + */ +export enum AffinityKind { + /** Multi-statement Read/Write transactions require hard stickiness. */ + ReadWrite = 0, + /** Read-Only transactions prefer soft stickiness for warmth. */ + ReadOnly = 1, +} + +/** + * Caller-owned handle managing channel affinity across multi-statement transactions. + */ +export class TransactionAffinity { + pinnedEntry: ChannelEntry | null = null; + onReset?: (entry: ChannelEntry) => void; + + constructor(readonly kind: AffinityKind = AffinityKind.ReadWrite) {} + + /** + * Resets affinity upon transaction commit, rollback, or close. + */ + reset(): void { + const entry = this.pinnedEntry; + if (entry && this.kind === AffinityKind.ReadWrite) { + entry.activeRwTransactions = Math.max(0, entry.activeRwTransactions - 1); + } + this.pinnedEntry = null; + if (entry && this.onReset) { + const callback = this.onReset; + this.onReset = undefined; + callback(entry); + } + } +} diff --git a/handwritten/spanner/src/channel-pool/channel-adapter.ts b/handwritten/spanner/src/channel-pool/channel-adapter.ts new file mode 100644 index 000000000000..c3ff46be1f9a --- /dev/null +++ b/handwritten/spanner/src/channel-pool/channel-adapter.ts @@ -0,0 +1,116 @@ +/*! + * 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 grpc from '@grpc/grpc-js'; +import {ChannelPool} from './types'; + +/** + * Adapter that presents a ChannelPool as a grpc.Channel to the gRPC client runtime. + */ +export class ChannelPoolChannelAdapter implements grpc.ChannelInterface { + constructor(readonly pool: ChannelPool) {} + + close(): void { + // No-op: the underlying ChannelPool lifecycle is managed by its owner + // (e.g. Spanner or caller), not by individual client stubs using this adapter. + } + + getTarget(): string { + return this.pool.getTarget(); + } + + getConnectivityState(tryToConnect?: boolean): grpc.connectivityState { + return this.pool.getConnectivityState(tryToConnect); + } + + watchConnectivityState( + currentState: grpc.connectivityState, + deadline: Date | number, + callback: (error?: Error) => void, + ): void { + this.pool.watchConnectivityState(currentState, deadline, callback); + } + + getChannelzRef(): any { + return null; + } + + get channelRefs(): Array<{channel: grpc.Channel}> { + return this.pool.getChannels().map(channel => ({channel})); + } + + createCall( + method: string, + deadline: grpc.Deadline, + host?: string | null, + parentCall?: any, + propagateFlags?: number | null, + ): any { + // Fallback if call is initiated outside of callInvocationTransformer + const lease = this.pool.acquire(); + let released = false; + const releaseOnce = () => { + if (!released) { + released = true; + lease.release(); + } + }; + let call: any; + try { + call = (lease.entry.channel as any).createCall( + method, + deadline, + host, + parentCall, + propagateFlags, + ); + } catch (error) { + releaseOnce(); + throw error; + } + return new grpc.InterceptingCall(call, { + start: (metadata, listener, next) => { + try { + next(metadata, { + onReceiveMetadata: (receivedMetadata, nextMetadata) => { + nextMetadata(receivedMetadata); + }, + onReceiveMessage: (message, nextMessage) => { + nextMessage(message); + }, + onReceiveStatus: (status, nextStatus) => { + releaseOnce(); + nextStatus(status); + }, + }); + } catch (error) { + releaseOnce(); + throw error; + } + }, + sendMessage: (message, next) => { + next(message); + }, + halfClose: next => { + next(); + }, + cancel: next => { + releaseOnce(); + next(); + }, + }); + } +} diff --git a/handwritten/spanner/src/channel-pool/dynamic-pool.ts b/handwritten/spanner/src/channel-pool/dynamic-pool.ts new file mode 100644 index 000000000000..c22aa0473f57 --- /dev/null +++ b/handwritten/spanner/src/channel-pool/dynamic-pool.ts @@ -0,0 +1,587 @@ +/*! + * 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 grpc from '@grpc/grpc-js'; +import {AffinityKind, TransactionAffinity} from './affinity'; +import {selectPowerOfTwo} from './p2c'; +import { + ChannelEntry, + ChannelLease, + ChannelPool, + DynamicChannelPoolOptions, +} from './types'; + +/** + * Dynamic load-based channel pool that scales up under high concurrency + * and gracefully drains channels during sustained idle periods. + */ +export class DynamicChannelPool implements ChannelPool { + private readonly activeEntries: ChannelEntry[] = []; + private readonly drainingEntries: ChannelEntry[] = []; + private nextPhysicalId = 1; + private readonly target: string; + private readonly credentials: grpc.ChannelCredentials; + private readonly channelOptions: Record; + private readonly minChannels: number; + private readonly maxChannels: number; + private readonly maxRpcPerChannel: number; + private readonly minRpcPerChannel: number; + private readonly maxScaleUpPercent: number; + private readonly maxRemoveChannels: number; + private readonly scaleUpCooldownMs: number; + private readonly primeTimeoutMs: number; + private readonly primeFn?: ( + channel: grpc.Channel, + sessionName: string, + ) => Promise; + + private scaleDownTimer?: NodeJS.Timeout; + private consecutiveLowLoadChecks = 0; + private lastScaleUpTime = 0; + private isScalingUp = false; + private scaleUpPromise?: Promise; + private primeSessionName?: string; + private isClosed = false; + + constructor( + target: string, + credentials: grpc.ChannelCredentials, + channelOptions: Record, + options?: DynamicChannelPoolOptions, + ) { + this.target = target; + this.credentials = credentials; + this.channelOptions = Object.assign({}, channelOptions); + + const rawMin = options?.minChannels; + let initialCount: number; + if (typeof rawMin === 'number' && !Number.isNaN(rawMin)) { + initialCount = Math.max(1, rawMin); + } else if (process.env.SPANNER_NUM_CHANNELS) { + const parsed = parseInt(process.env.SPANNER_NUM_CHANNELS, 10); + initialCount = !Number.isNaN(parsed) && parsed > 0 ? parsed : 4; + } else { + initialCount = 4; + } + + const rawMax = options?.maxChannels; + const maxChannels = + typeof rawMax === 'number' && !Number.isNaN(rawMax) + ? Math.max(initialCount, rawMax) + : Math.max(initialCount, 256); + + this.minChannels = initialCount; + this.maxChannels = maxChannels; + const rawMaxRpc = options?.maxRpcPerChannel; + this.maxRpcPerChannel = + typeof rawMaxRpc === 'number' && !Number.isNaN(rawMaxRpc) + ? Math.max(1, rawMaxRpc) + : 8; + + const defaultMinRpc = Math.min(2, Math.max(0, this.maxRpcPerChannel - 1)); + const rawMinRpc = options?.minRpcPerChannel; + this.minRpcPerChannel = + typeof rawMinRpc === 'number' && !Number.isNaN(rawMinRpc) + ? Math.min(Math.max(0, rawMinRpc), this.maxRpcPerChannel - 1) + : defaultMinRpc; + this.maxScaleUpPercent = + typeof options?.maxScaleUpPercent === 'number' && + !Number.isNaN(options.maxScaleUpPercent) + ? options.maxScaleUpPercent + : 100; + this.maxRemoveChannels = + typeof options?.maxRemoveChannels === 'number' && + !Number.isNaN(options.maxRemoveChannels) + ? options.maxRemoveChannels + : 4; + this.scaleUpCooldownMs = + typeof options?.scaleUpCooldownMs === 'number' && + !Number.isNaN(options.scaleUpCooldownMs) + ? options.scaleUpCooldownMs + : 1000; + this.primeTimeoutMs = + typeof options?.primeTimeoutMs === 'number' && + !Number.isNaN(options.primeTimeoutMs) + ? options.primeTimeoutMs + : 5000; + this.primeFn = options?.primeFn; + + // Eagerly establish initial startup channels (unprimed, as session does not exist yet) + for (let index = 0; index < this.minChannels; index++) { + this.activeEntries.push(this.createEntry(index + 1)); + } + + const intervalMs = options?.scaleDownIntervalMs ?? 60000; + if (intervalMs > 0) { + this.scaleDownTimer = setInterval(() => { + this.evaluateScaleDown(); + }, intervalMs); + this.scaleDownTimer.unref(); + } + } + + /** + * Sets or updates the active multiplexed session name used for channel priming. + */ + setPrimeSession(sessionName: string): void { + this.primeSessionName = sessionName; + } + + acquire(affinity?: TransactionAffinity): ChannelLease { + if (this.isClosed) { + throw new Error('Channel pool is closed.'); + } + + let entry: ChannelEntry; + + if ( + affinity?.pinnedEntry && + (affinity.pinnedEntry.state === 'ACTIVE' || + (affinity.pinnedEntry.state === 'DRAINING' && + affinity.kind === AffinityKind.ReadWrite)) + ) { + entry = affinity.pinnedEntry; + } else { + if (affinity?.pinnedEntry && affinity.kind === AffinityKind.ReadWrite) { + affinity.pinnedEntry.activeRwTransactions = Math.max( + 0, + affinity.pinnedEntry.activeRwTransactions - 1, + ); + } + entry = selectPowerOfTwo(this.activeEntries); + if (affinity) { + affinity.pinnedEntry = entry; + affinity.onReset = (entryToDrain: ChannelEntry) => { + if (entryToDrain.state === 'DRAINING') { + this.checkDrainedEntry(entryToDrain); + } + }; + if (affinity.kind === AffinityKind.ReadWrite) { + entry.activeRwTransactions++; + } + } + } + + entry.inFlightRpcs++; + entry.lastActivity = Date.now(); + + // 3. Event-driven scale-up check + if (entry.inFlightRpcs > this.maxRpcPerChannel) { + this.maybeScaleUp(); + } + + let released = false; + return { + entry, + release: () => { + if (!released) { + released = true; + entry.inFlightRpcs = Math.max(0, entry.inFlightRpcs - 1); + entry.lastActivity = Date.now(); + if (entry.state === 'DRAINING') { + this.checkDrainedEntry(entry); + } + if (this.drainingEntries.length > 0) { + for ( + let index = this.drainingEntries.length - 1; + index >= 0; + index-- + ) { + this.checkDrainedEntry(this.drainingEntries[index]); + } + } + } + }, + }; + } + + get size(): number { + return this.activeEntries.length + this.drainingEntries.length; + } + + get activeCount(): number { + return this.activeEntries.length; + } + + getTarget(): string { + return this.target; + } + + getChannels(): grpc.Channel[] { + return [...this.activeEntries, ...this.drainingEntries].map( + entry => entry.channel, + ); + } + + getConnectivityState(tryToConnect?: boolean): grpc.connectivityState { + if (this.isClosed || this.activeEntries.length === 0) { + return grpc.connectivityState.SHUTDOWN; + } + + let hasReady = false; + let hasConnecting = false; + let hasTransientFailure = false; + + for (const entry of this.activeEntries) { + const state = entry.channel.getConnectivityState(tryToConnect ?? false); + if (state === grpc.connectivityState.READY) { + hasReady = true; + } else if (state === grpc.connectivityState.CONNECTING) { + hasConnecting = true; + } else if (state === grpc.connectivityState.TRANSIENT_FAILURE) { + hasTransientFailure = true; + } + } + + if (hasReady) { + return grpc.connectivityState.READY; + } + if (hasConnecting) { + return grpc.connectivityState.CONNECTING; + } + if (hasTransientFailure) { + return grpc.connectivityState.TRANSIENT_FAILURE; + } + return grpc.connectivityState.IDLE; + } + + watchConnectivityState( + currentState: grpc.connectivityState, + deadline: Date | number, + callback: (error?: Error) => void, + ): void { + if (this.activeEntries.length === 0) { + callback(new Error('No channels in pool.')); + return; + } + const channel = this.activeEntries[0].channel; + channel.watchConnectivityState(currentState, deadline, error => { + if (error || this.isClosed) { + callback(error); + return; + } + const channelState = channel.getConnectivityState(false); + if ( + channelState === grpc.connectivityState.SHUTDOWN && + !this.isClosed && + this.activeEntries.length > 0 + ) { + if (this.getConnectivityState() !== currentState) { + callback(); + return; + } + this.watchConnectivityState(currentState, deadline, callback); + return; + } + callback(error); + }); + } + + async close(): Promise { + this.isClosed = true; + if (this.scaleDownTimer) { + clearInterval(this.scaleDownTimer); + } + if (this.scaleUpPromise) { + await this.scaleUpPromise; + } + for (const entry of this.activeEntries) { + entry.state = 'CLOSED'; + entry.channel.close(); + } + for (const entry of this.drainingEntries) { + entry.state = 'CLOSED'; + entry.channel.close(); + } + this.activeEntries.length = 0; + this.drainingEntries.length = 0; + } + + private allocateLogicalSlot(inProgressEntries: ChannelEntry[] = []): number { + const occupiedSlots = new Set(); + for (const entry of this.activeEntries) { + occupiedSlots.add(entry.id); + } + for (const entry of this.drainingEntries) { + if (entry.state !== 'CLOSED') { + occupiedSlots.add(entry.id); + } + } + for (const entry of inProgressEntries) { + occupiedSlots.add(entry.id); + } + + for (let slot = 1; slot <= this.maxChannels; slot++) { + if (!occupiedSlots.has(slot)) { + return slot; + } + } + + let slot = this.maxChannels + 1; + while (occupiedSlots.has(slot)) { + slot++; + } + return slot; + } + + private createEntry(logicalId?: number): ChannelEntry { + const id = logicalId ?? this.allocateLogicalSlot(); + const physicalId = this.nextPhysicalId++; + const entryChannelOptions = Object.assign({}, this.channelOptions, { + 'grpc.channel_id': physicalId, + 'grpc_gcp.client_channel.id': physicalId, + }); + const channel = new grpc.Channel( + this.target, + this.credentials, + entryChannelOptions, + ); + const now = Date.now(); + return { + id, + channel, + inFlightRpcs: 0, + activeRwTransactions: 0, + state: 'ACTIVE', + lastActivity: now, + createdAt: now, + }; + } + + private maybeScaleUp(): void { + if ( + this.isClosed || + this.isScalingUp || + this.activeEntries.length >= this.maxChannels + ) { + return; + } + + const now = Date.now(); + if (now - this.lastScaleUpTime < this.scaleUpCooldownMs) { + return; + } + this.lastScaleUpTime = now; + this.isScalingUp = true; + + this.scaleUpPromise = new Promise(resolve => { + setImmediate(async () => { + try { + if (this.isClosed || this.activeEntries.length >= this.maxChannels) { + return; + } + + const currentLength = this.activeEntries.length; + let totalLoad = 0; + for (const entry of this.activeEntries) { + totalLoad += entry.inFlightRpcs; + } + const targetRpc = Math.max( + 1, + Math.floor((this.minRpcPerChannel + this.maxRpcPerChannel) / 2), + ); + const desiredChannels = Math.ceil(totalLoad / targetRpc); + const maxToAddByPercent = Math.max( + 2, + Math.ceil((currentLength * this.maxScaleUpPercent) / 100), + ); + // Do NOT veto scale-up when average load is below max: if a channel exceeded maxRpcPerChannel, + // we add at least 1 channel (and up to maxToAddByPercent bounded by maxChannels). + const needed = Math.max(1, desiredChannels - currentLength); + const count = Math.min( + needed, + maxToAddByPercent, + this.maxChannels - currentLength, + ); + + if (count <= 0) { + return; + } + + const newEntries: ChannelEntry[] = []; + for (let index = 0; index < count; index++) { + const logicalId = this.allocateLogicalSlot(newEntries); + newEntries.push(this.createEntry(logicalId)); + } + + // Prime newly created channels in parallel + await Promise.all( + newEntries.map(async entry => { + try { + await this.primeChannel(entry); + if ( + !this.isClosed && + this.activeEntries.length < this.maxChannels + ) { + this.activeEntries.push(entry); + } else { + entry.state = 'CLOSED'; + entry.channel.close(); + } + } catch { + // Discard and close channel if priming fails + entry.state = 'CLOSED'; + entry.channel.close(); + } + }), + ); + } catch { + // Prevent unhandled promise rejection in setImmediate callback + } finally { + this.isScalingUp = false; + resolve(); + } + }); + }); + } + + private async primeChannel(entry: ChannelEntry): Promise { + if (!this.primeSessionName || !this.primeFn) { + return; + } + const maxAttempts = 2; + let lastError: unknown; + for (let attempt = 0; attempt < maxAttempts; attempt++) { + let timer: NodeJS.Timeout | undefined; + try { + const primePromise = this.primeFn(entry.channel, this.primeSessionName); + const timeoutPromise = new Promise((_resolve, reject) => { + timer = setTimeout(() => { + reject(new Error('Channel priming timed out.')); + }, this.primeTimeoutMs); + timer.unref(); + }); + await Promise.race([primePromise, timeoutPromise]); + return; + } catch (error) { + lastError = error; + } finally { + if (timer) { + clearTimeout(timer); + } + } + } + throw lastError; + } + + private evaluateScaleDown(): void { + if ( + this.isClosed || + this.isScalingUp || + Date.now() - this.lastScaleUpTime < this.scaleUpCooldownMs + ) { + return; + } + + if (this.drainingEntries.length > 0) { + for (let index = this.drainingEntries.length - 1; index >= 0; index--) { + this.checkDrainedEntry(this.drainingEntries[index]); + } + } + + if (this.activeEntries.length <= this.minChannels) { + this.consecutiveLowLoadChecks = 0; + return; + } + + let totalInFlight = 0; + for (const entry of this.activeEntries) { + totalInFlight += entry.inFlightRpcs; + } + const averageLoad = totalInFlight / this.activeEntries.length; + + if (averageLoad < this.minRpcPerChannel) { + this.consecutiveLowLoadChecks++; + if (this.consecutiveLowLoadChecks >= 3) { + this.consecutiveLowLoadChecks = 0; + this.drainChannels(totalInFlight); + } + } else { + this.consecutiveLowLoadChecks = 0; + } + } + + private drainChannels(totalInFlight: number): void { + if (this.activeEntries.length <= this.minChannels) { + return; + } + + const targetRpc = Math.max( + 1, + Math.floor((this.minRpcPerChannel + this.maxRpcPerChannel) / 2), + ); + const desiredChannels = Math.max( + this.minChannels, + Math.ceil(totalInFlight / targetRpc), + ); + + if (desiredChannels >= this.activeEntries.length) { + return; + } + + const excessChannels = this.activeEntries.length - desiredChannels; + const channelsToRemove = Math.min( + excessChannels, + this.maxRemoveChannels, + this.activeEntries.length - this.minChannels, + ); + + if (channelsToRemove <= 0) { + return; + } + + // Sort candidates so preferred drain victims are placed first: + // 1. Lowest in-flight RPC load + // 2. Lowest active Read/Write transactions + // 3. Newer channels (prefer keeping older, warmer channels) + this.activeEntries.sort((a, b) => { + if (a.inFlightRpcs !== b.inFlightRpcs) { + return a.inFlightRpcs - b.inFlightRpcs; + } + if (a.activeRwTransactions !== b.activeRwTransactions) { + return a.activeRwTransactions - b.activeRwTransactions; + } + const aCreated = a.createdAt ?? a.id; + const bCreated = b.createdAt ?? b.id; + if (bCreated !== aCreated) { + return bCreated - aCreated; + } + return b.id - a.id; + }); + + const draining = this.activeEntries.splice(0, channelsToRemove); + for (const entry of draining) { + entry.state = 'DRAINING'; + this.drainingEntries.push(entry); + this.checkDrainedEntry(entry); + } + } + + private checkDrainedEntry(entry: ChannelEntry): void { + if ( + entry.state === 'DRAINING' && + entry.inFlightRpcs === 0 && + entry.activeRwTransactions === 0 + ) { + entry.state = 'CLOSED'; + entry.channel.close(); + const index = this.drainingEntries.indexOf(entry); + if (index !== -1) { + this.drainingEntries.splice(index, 1); + } + } + } +} diff --git a/handwritten/spanner/src/channel-pool/index.ts b/handwritten/spanner/src/channel-pool/index.ts new file mode 100644 index 000000000000..58e8d98bc2a9 --- /dev/null +++ b/handwritten/spanner/src/channel-pool/index.ts @@ -0,0 +1,55 @@ +/*! + * 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 grpc from '@grpc/grpc-js'; +import {StaticChannelPool} from './static-pool'; +import {DynamicChannelPool} from './dynamic-pool'; +import {ChannelPool, ChannelPoolConfig} from './types'; + +export * from './types'; +export * from './affinity'; +export * from './p2c'; +export * from './static-pool'; +export * from './dynamic-pool'; +export * from './transformer'; +export * from './channel-adapter'; + +/** + * Creates a ChannelPool instance based on user configuration. + */ +export function createChannelPool( + address: string, + credentials: grpc.ChannelCredentials, + channelOptions: Record, + config?: ChannelPoolConfig, +): ChannelPool { + if (config?.type === 'dynamic') { + return new DynamicChannelPool(address, credentials, channelOptions, config); + } + return new StaticChannelPool(address, credentials, channelOptions, config); +} + +/** + * Type guard to check if an object implements ChannelPool. + */ +export function isChannelPool(obj: any): obj is ChannelPool { + return Boolean( + obj && + typeof obj === 'object' && + typeof obj.acquire === 'function' && + typeof obj.close === 'function', + ); +} diff --git a/handwritten/spanner/src/channel-pool/p2c.ts b/handwritten/spanner/src/channel-pool/p2c.ts new file mode 100644 index 000000000000..0884b645cfc9 --- /dev/null +++ b/handwritten/spanner/src/channel-pool/p2c.ts @@ -0,0 +1,57 @@ +/*! + * 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 {ChannelEntry} from './types'; + +/** + * Selects an active channel using the Power of Two Choices (P2C) algorithm. + * Randomly picks two distinct channels, compares their effective load + * (in-flight RPCs + active R/W transactions), and chooses the lesser loaded. + * On tie, prefers the warmer channel (most recent activity). + * + * @param entries List of available active channel entries. + * @returns The chosen ChannelEntry. + */ +export function selectPowerOfTwo(entries: ChannelEntry[]): ChannelEntry { + const length = entries.length; + if (length === 0) { + throw new Error('No active channels available in channel pool.'); + } + if (length === 1) { + return entries[0]; + } + + const index1 = Math.floor(Math.random() * length); + let index2 = Math.floor(Math.random() * (length - 1)); + if (index2 >= index1) { + index2++; + } + + const entry1 = entries[index1]; + const entry2 = entries[index2]; + + const load1 = entry1.inFlightRpcs + entry1.activeRwTransactions; + const load2 = entry2.inFlightRpcs + entry2.activeRwTransactions; + + if (load1 < load2) { + return entry1; + } + if (load2 < load1) { + return entry2; + } + // Tie-breaker: prefer warmer channel + return entry1.lastActivity >= entry2.lastActivity ? entry1 : entry2; +} diff --git a/handwritten/spanner/src/channel-pool/static-pool.ts b/handwritten/spanner/src/channel-pool/static-pool.ts new file mode 100644 index 000000000000..a6076a8b6b14 --- /dev/null +++ b/handwritten/spanner/src/channel-pool/static-pool.ts @@ -0,0 +1,192 @@ +/*! + * 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 grpc from '@grpc/grpc-js'; +import {AffinityKind, TransactionAffinity} from './affinity'; +import {selectPowerOfTwo} from './p2c'; +import { + ChannelEntry, + ChannelLease, + ChannelPool, + StaticChannelPoolOptions, +} from './types'; + +/** + * Fixed-size channel pool that allocates N channels on startup and uses P2C selection. + * Operates with zero background timers and predictable resource usage. + */ +export class StaticChannelPool implements ChannelPool { + private readonly entries: ChannelEntry[]; + private readonly target: string; + private isClosed = false; + + constructor( + target: string, + credentials: grpc.ChannelCredentials, + channelOptions: Record, + options?: StaticChannelPoolOptions, + ) { + this.target = target; + const rawNumChannels = options?.numChannels; + let count: number; + if (typeof rawNumChannels === 'number' && !Number.isNaN(rawNumChannels)) { + count = Math.max(1, rawNumChannels); + } else if (process.env.SPANNER_NUM_CHANNELS) { + const parsed = parseInt(process.env.SPANNER_NUM_CHANNELS, 10); + count = !Number.isNaN(parsed) && parsed > 0 ? parsed : 4; + } else { + count = 4; + } + + this.entries = []; + for (let index = 0; index < count; index++) { + const entryChannelOptions = Object.assign({}, channelOptions, { + 'grpc.channel_id': index + 1, + 'grpc_gcp.client_channel.id': index + 1, + }); + const channel = new grpc.Channel( + target, + credentials, + entryChannelOptions, + ); + const now = Date.now(); + this.entries.push({ + id: index + 1, + channel, + inFlightRpcs: 0, + activeRwTransactions: 0, + state: 'ACTIVE', + lastActivity: now, + createdAt: now, + }); + } + } + + acquire(affinity?: TransactionAffinity): ChannelLease { + if (this.isClosed) { + throw new Error('Channel pool is closed.'); + } + + let entry: ChannelEntry; + + if (affinity?.pinnedEntry && affinity.pinnedEntry.state === 'ACTIVE') { + entry = affinity.pinnedEntry; + } else { + if (affinity?.pinnedEntry && affinity.kind === AffinityKind.ReadWrite) { + affinity.pinnedEntry.activeRwTransactions = Math.max( + 0, + affinity.pinnedEntry.activeRwTransactions - 1, + ); + } + entry = selectPowerOfTwo(this.entries); + if (affinity) { + affinity.pinnedEntry = entry; + if (affinity.kind === AffinityKind.ReadWrite) { + entry.activeRwTransactions++; + } + } + } + + entry.inFlightRpcs++; + entry.lastActivity = Date.now(); + + let released = false; + return { + entry, + release: () => { + if (!released) { + released = true; + entry.inFlightRpcs = Math.max(0, entry.inFlightRpcs - 1); + entry.lastActivity = Date.now(); + } + }, + }; + } + + get size(): number { + return this.entries.length; + } + + get activeCount(): number { + return this.entries.length; + } + + getTarget(): string { + return this.target; + } + + getChannels(): grpc.Channel[] { + return this.entries.map(entry => entry.channel); + } + + getConnectivityState(tryToConnect?: boolean): grpc.connectivityState { + if (this.isClosed || this.entries.length === 0) { + return grpc.connectivityState.SHUTDOWN; + } + + let hasReady = false; + let hasConnecting = false; + let hasTransientFailure = false; + + for (const entry of this.entries) { + const state = entry.channel.getConnectivityState(tryToConnect ?? false); + if (state === grpc.connectivityState.READY) { + hasReady = true; + } else if (state === grpc.connectivityState.CONNECTING) { + hasConnecting = true; + } else if (state === grpc.connectivityState.TRANSIENT_FAILURE) { + hasTransientFailure = true; + } + } + + if (hasReady) { + return grpc.connectivityState.READY; + } + if (hasConnecting) { + return grpc.connectivityState.CONNECTING; + } + if (hasTransientFailure) { + return grpc.connectivityState.TRANSIENT_FAILURE; + } + return grpc.connectivityState.IDLE; + } + + watchConnectivityState( + currentState: grpc.connectivityState, + deadline: Date | number, + callback: (error?: Error) => void, + ): void { + if (this.entries.length === 0) { + callback(new Error('No channels in pool.')); + return; + } + // Watch connectivity on the first channel as representative + this.entries[0].channel.watchConnectivityState( + currentState, + deadline, + callback, + ); + } + + async close(): Promise { + this.isClosed = true; + for (const entry of this.entries) { + entry.state = 'CLOSED'; + entry.channel.close(); + } + this.entries.length = 0; + } +} diff --git a/handwritten/spanner/src/channel-pool/transformer.ts b/handwritten/spanner/src/channel-pool/transformer.ts new file mode 100644 index 000000000000..09f32828cf9b --- /dev/null +++ b/handwritten/spanner/src/channel-pool/transformer.ts @@ -0,0 +1,132 @@ +/*! + * 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 grpc from '@grpc/grpc-js'; +import {TransactionAffinity} from './affinity'; +import {ChannelPool} from './types'; + +/** + * Creates a gRPC CallInvocationTransformer that routes calls through a Spanner ChannelPool. + * + * @param pool The ChannelPool instance or a resolver returning the active ChannelPool. + * @returns A gRPC CallInvocationTransformer function. + */ +export function createCallInvocationTransformer( + poolOrResolver: ChannelPool | (() => ChannelPool | undefined), +) { + return function spannerCallInvocationTransformer( + callProperties: grpc.CallProperties, + ): grpc.CallProperties { + const pool = + typeof poolOrResolver === 'function' ? poolOrResolver() : poolOrResolver; + + if (!pool) { + return callProperties; + } + + const affinity: TransactionAffinity | undefined = + (callProperties.callOptions as any)?.affinity ?? + (callProperties.callOptions as any)?.otherArgs?.options?.affinity; + const lease = pool.acquire(affinity); + + // Update logical channel ID in Spanner Request ID header if present + if (callProperties.metadata) { + const existing = callProperties.metadata.get('x-goog-spanner-request-id'); + if (existing.length > 0 && typeof existing[0] === 'string') { + const parts = existing[0].split('.'); + if (parts.length >= 6) { + parts[3] = String(lease.entry.id); + callProperties.metadata.set( + 'x-goog-spanner-request-id', + parts.join('.'), + ); + } else { + callProperties.metadata.set( + 'x-goog-spanner-request-id', + `${existing[0]}.${lease.entry.id}`, + ); + } + } + } + + const releaseInterceptor: grpc.Interceptor = (options, nextCall) => { + let released = false; + const releaseOnce = () => { + if (!released) { + released = true; + lease.release(); + } + }; + + const requester: grpc.Requester = { + start: (metadata, listener, next) => { + const newListener: grpc.Listener = { + onReceiveMetadata: (receivedMetadata, nextMetadata) => { + nextMetadata(receivedMetadata); + }, + onReceiveMessage: (message, nextMessage) => { + nextMessage(message); + }, + onReceiveStatus: (status, nextStatus) => { + releaseOnce(); + nextStatus(status); + }, + }; + try { + next(metadata, newListener); + } catch (error) { + releaseOnce(); + throw error; + } + }, + sendMessage: (message, next) => { + next(message); + }, + halfClose: next => { + next(); + }, + cancel: next => { + releaseOnce(); + next(); + }, + }; + + let nextCallResult: ReturnType; + try { + nextCallResult = nextCall(options); + } catch (error) { + releaseOnce(); + throw error; + } + return new grpc.InterceptingCall(nextCallResult, requester); + }; + + const callOptions = Object.assign({}, callProperties.callOptions); + callOptions.interceptors = (callOptions.interceptors || []).concat([ + releaseInterceptor, + ]); + + return { + argument: callProperties.argument, + metadata: callProperties.metadata, + call: callProperties.call, + channel: lease.entry.channel, + methodDefinition: callProperties.methodDefinition, + callOptions, + callback: callProperties.callback, + }; + }; +} diff --git a/handwritten/spanner/src/channel-pool/types.ts b/handwritten/spanner/src/channel-pool/types.ts new file mode 100644 index 000000000000..85f3d2571726 --- /dev/null +++ b/handwritten/spanner/src/channel-pool/types.ts @@ -0,0 +1,138 @@ +/*! + * 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 grpc from '@grpc/grpc-js'; +import {TransactionAffinity} from './affinity'; + +/** + * State of a channel entry within the channel pool. + */ +export type ChannelState = 'ACTIVE' | 'DRAINING' | 'CLOSED'; + +/** + * Represents a single gRPC channel within the channel pool. + */ +export interface ChannelEntry { + /** 1-based logical channel identifier (used for request IDs). */ + readonly id: number; + /** Physical gRPC channel. */ + readonly channel: grpc.Channel; + /** Number of currently in-flight RPCs on this channel. */ + inFlightRpcs: number; + /** Number of active Read/Write transactions pinned to this channel. */ + activeRwTransactions: number; + /** Current state of the channel. */ + state: ChannelState; + /** Timestamp (epoch milliseconds) of the most recent activity. */ + lastActivity: number; + /** Timestamp (epoch milliseconds) when the channel was created. */ + createdAt?: number; +} + +/** + * Lease returned upon acquiring a channel from the pool. + */ +export interface ChannelLease { + /** The leased channel entry. */ + readonly entry: ChannelEntry; + /** Releases the in-flight lease when the RPC completes or is cancelled. */ + release(): void; +} + +/** + * Generic interface for a Spanner gRPC channel pool. + */ +export interface ChannelPool { + /** + * Acquires a channel from the pool. + * If an affinity handle is provided and already pinned, routes to the pinned channel. + * Otherwise, selects an active channel via Power of Two Choices (P2C). + */ + acquire(affinity?: TransactionAffinity): ChannelLease; + + /** Total number of channels currently tracked (active + draining). */ + readonly size: number; + + /** Number of active channels available for new selections. */ + readonly activeCount: number; + + /** Closes all physical channels in the pool and shuts down background monitors. */ + close(): Promise; + + /** Gets the overall connectivity state of the pool. */ + getConnectivityState(tryToConnect?: boolean): grpc.connectivityState; + + /** Watches for connectivity state changes across the pool. */ + watchConnectivityState( + currentState: grpc.connectivityState, + deadline: Date | number, + callback: (error?: Error) => void, + ): void; + + /** Returns the target address of the channel pool. */ + getTarget(): string; + + /** + * Sets or updates the active multiplexed session name used for channel priming. + */ + setPrimeSession?(sessionName: string): void; + + /** + * Returns the underlying gRPC channels currently managed by the pool. + */ + getChannels(): grpc.Channel[]; +} + +/** + * Configuration options for StaticChannelPool. + */ +export interface StaticChannelPoolOptions { + /** Number of channels to maintain in the pool. Defaults to 4. */ + numChannels?: number; +} + +/** + * Configuration options for DynamicChannelPool. + */ +export interface DynamicChannelPoolOptions { + /** Minimum number of channels to retain in the pool. Defaults to 4. */ + minChannels?: number; + /** Maximum number of channels allowed. Defaults to 256. */ + maxChannels?: number; + /** Load threshold per channel (in-flight RPCs) to trigger scale-up. Defaults to 8. */ + maxRpcPerChannel?: number; + /** Low-load threshold per channel used for scale-down checks. Defaults to 2. */ + minRpcPerChannel?: number; + /** Maximum percentage of current pool size added per scale-up event. Defaults to 100 (min 2 channels). */ + maxScaleUpPercent?: number; + /** Maximum number of channels marked draining per scale-down cycle. Defaults to 4. */ + maxRemoveChannels?: number; + /** Interval in milliseconds between periodic scale-down evaluations. Defaults to 60,000 (1 min). */ + scaleDownIntervalMs?: number; + /** Cooldown in milliseconds between consecutive scale-up attempts. Defaults to 1,000ms. */ + scaleUpCooldownMs?: number; + /** Timeout in milliseconds for channel priming queries. Defaults to 5,000ms. */ + primeTimeoutMs?: number; + /** Optional custom priming callback executing SELECT 1 on newly dialed channels. */ + primeFn?: (channel: grpc.Channel, sessionName: string) => Promise; +} + +/** + * Union configuration type accepted in Spanner options. + */ +export type ChannelPoolConfig = + | ({type: 'static'} & StaticChannelPoolOptions) + | ({type: 'dynamic'} & DynamicChannelPoolOptions); diff --git a/handwritten/spanner/src/index.ts b/handwritten/spanner/src/index.ts index 297c08c7536e..628847715145 100644 --- a/handwritten/spanner/src/index.ts +++ b/handwritten/spanner/src/index.ts @@ -92,6 +92,14 @@ import { import grpcGcpModule = require('grpc-gcp'); const grpcGcp = grpcGcpModule(grpc); import * as v1 from './v1'; +import { + ChannelPool, + ChannelPoolChannelAdapter, + ChannelPoolConfig, + createCallInvocationTransformer, + createChannelPool, + isChannelPool, +} from './channel-pool'; import { ObservabilityOptions, ensureInitialContextManagerSet, @@ -188,6 +196,7 @@ export interface SpannerOptions extends GrpcClientOptions { */ universe_domain?: string; universeDomain?: string; + channelPool?: boolean | string | ChannelPoolConfig | ChannelPool; /** * Whether to enable gRPC Channelz service tracking. * Defaults to `0` (disabled) to eliminate per-RPC allocation and tracking overhead. @@ -342,6 +351,8 @@ class Spanner extends GrpcService { private _isInSecureCredentials: boolean; private _metricsEnabled = false; readonly _nthClientId: number; + private channelPool_?: ChannelPool; + private _ownsChannelPool = false; /** * Placeholder used to auto populate a column with the commit timestamp. @@ -421,23 +432,149 @@ class Spanner extends GrpcService { } } - options = Object.assign( - { - libName: 'gccl', - libVersion: require('../../package.json').version, - scopes, - // Add grpc keep alive setting - 'grpc.keepalive_time_ms': 120000, - // Disable Channelz by default to reduce per-RPC tracking and allocation overhead - 'grpc.enable_channelz': 0, - // Enable grpc-gcp support - 'grpc.callInvocationTransformer': grpcGcp.gcpCallInvocationTransformer, - 'grpc.channelFactoryOverride': grpcGcp.gcpChannelFactoryOverride, - 'grpc.gcpApiConfig': grpcGcp.createGcpApiConfig(gcpApiConfig), - grpc, - }, - options || {}, - ) as {} as SpannerOptions; + const rawPool = options?.channelPool; + const envChannelPool = process.env.SPANNER_CHANNEL_POOL?.toLowerCase(); + const isLegacyPool = + rawPool === 'grpc-gcp' || + rawPool === 'legacy' || + (rawPool as any) === false || + (rawPool === undefined && + (envChannelPool === 'grpc-gcp' || envChannelPool === 'legacy')); + + let initialChannelPool: ChannelPool | undefined; + + if (!isLegacyPool) { + let channelPoolInstance: ChannelPool | undefined; + let poolConfig: ChannelPoolConfig; + + if (isChannelPool(rawPool)) { + channelPoolInstance = rawPool; + initialChannelPool = channelPoolInstance; + poolConfig = {type: 'dynamic'}; + } else if (typeof rawPool === 'object' && rawPool !== null) { + poolConfig = rawPool as ChannelPoolConfig; + } else { + let poolType: 'dynamic' | 'static' = 'static'; + if (rawPool === 'dynamic' || rawPool === 'static') { + poolType = rawPool; + } else if ( + envChannelPool === 'dynamic' || + envChannelPool === 'static' + ) { + poolType = envChannelPool; + } + poolConfig = + poolType === 'dynamic' + ? { + type: 'dynamic', + minChannels: process.env.SPANNER_NUM_CHANNELS + ? parseInt(process.env.SPANNER_NUM_CHANNELS, 10) + : undefined, + } + : { + type: 'static', + numChannels: process.env.SPANNER_NUM_CHANNELS + ? parseInt(process.env.SPANNER_NUM_CHANNELS, 10) + : 4, + }; + } + + const primeFn = async (channel: grpc.Channel, sessionName: string) => { + const unclosableChannel = new Proxy(channel, { + get(target, property, receiver) { + if (property === 'close') { + return () => {}; + } + const value = Reflect.get(target, property, receiver); + return typeof value === 'function' ? value.bind(target) : value; + }, + }); + const primeClient = new v1.SpannerClient({ + ...options, + 'grpc.channelFactoryOverride': () => unclosableChannel, + 'grpc.callInvocationTransformer': undefined, + 'grpc.gcpApiConfig': undefined, + } as any); + const timeoutMs = + poolConfig.type === 'dynamic' + ? (poolConfig.primeTimeoutMs ?? 5000) + : 5000; + try { + await primeClient.executeSql( + { + session: sessionName, + sql: 'SELECT 1', + }, + { + timeout: timeoutMs, + }, + ); + } finally { + await primeClient.close(); + } + }; + + if (poolConfig.type === 'dynamic' && !poolConfig.primeFn) { + poolConfig.primeFn = primeFn; + } + + const channelFactoryOverride = ( + address: string, + credentials: grpc.ChannelCredentials, + channelOptions: any, + ) => { + if (!channelPoolInstance) { + channelPoolInstance = createChannelPool( + address, + credentials, + channelOptions, + poolConfig, + ); + this.channelPool_ = channelPoolInstance; + this._ownsChannelPool = true; + } + return new ChannelPoolChannelAdapter(channelPoolInstance); + }; + + const callInvocationTransformer = createCallInvocationTransformer( + () => this.channelPool_ || channelPoolInstance, + ); + + options = Object.assign( + { + libName: 'gccl', + libVersion: require('../../package.json').version, + scopes, + 'grpc.keepalive_time_ms': 120000, + // Disable Channelz by default to reduce per-RPC tracking and allocation overhead + 'grpc.enable_channelz': 0, + 'grpc.callInvocationTransformer': callInvocationTransformer, + 'grpc.channelFactoryOverride': channelFactoryOverride, + grpc, + }, + options || {}, + ) as {} as SpannerOptions; + delete (options as any)['grpc.gcpApiConfig']; + } else { + options = Object.assign( + { + libName: 'gccl', + libVersion: require('../../package.json').version, + scopes, + // Add grpc keep alive setting + 'grpc.keepalive_time_ms': 120000, + // Disable Channelz by default to reduce per-RPC tracking and allocation overhead + 'grpc.enable_channelz': 0, + // Enable grpc-gcp support + 'grpc.callInvocationTransformer': + grpcGcp.gcpCallInvocationTransformer, + 'grpc.channelFactoryOverride': grpcGcp.gcpChannelFactoryOverride, + 'grpc.gcpApiConfig': grpcGcp.createGcpApiConfig(gcpApiConfig), + grpc, + }, + options || {}, + ) as {} as SpannerOptions; + } const directedReadOptions = options.directedReadOptions ? options.directedReadOptions @@ -489,6 +626,9 @@ class Spanner extends GrpcService { packageJson: require('../../package.json'), } as {} as GrpcServiceConfig; super(config, options); + if (initialChannelPool) { + this.channelPool_ = initialChannelPool; + } if (options.routeToLeaderEnabled === false) { this.routeToLeaderEnabled = false; @@ -523,6 +663,16 @@ class Spanner extends GrpcService { return this._universeDomain; } + get channelPool(): ChannelPool | undefined { + return this.channelPool_; + } + + setPrimeSession(sessionName: string): void { + if (this.channelPool_ && 'setPrimeSession' in this.channelPool_) { + (this.channelPool_ as any).setPrimeSession(sessionName); + } + } + /** * Gets the InstanceAdminClient object. * The returned InstanceAdminClient object is a shared, managed instance and should not be manually closed. @@ -593,6 +743,10 @@ class Spanner extends GrpcService { } }); + if (this.channelPool_ && this._ownsChannelPool) { + promises.push(Promise.resolve().then(() => this.channelPool_!.close())); + } + // Wait for all close attempts to settle. // Map success to undefined, and failure to the error. const results = await Promise.all( @@ -2500,3 +2654,4 @@ export {v1, protos}; export default {Spanner}; export {Float32, Float, Int, Struct, Numeric, PGNumeric, SpannerDate, Interval}; export {ObservabilityOptions}; +export * from './channel-pool'; diff --git a/handwritten/spanner/src/metrics/metrics-tracer-factory.ts b/handwritten/spanner/src/metrics/metrics-tracer-factory.ts index e103d2a80cfa..9ae0bcca735c 100644 --- a/handwritten/spanner/src/metrics/metrics-tracer-factory.ts +++ b/handwritten/spanner/src/metrics/metrics-tracer-factory.ts @@ -353,13 +353,39 @@ export class MetricsTracerFactory { */ public getCurrentTracer(requestId: string): MetricsTracer | null { const operationRequest: string = this._extractOperationRequest(requestId); - if (!this._currentOperationTracers.has(operationRequest)) { + let tracer = this._currentOperationTracers.get(operationRequest); + let key = operationRequest; + if (!tracer && operationRequest) { + // Channel pool can rewrite the channelId component (parts[3]) of requestId. + // Since the original channel ID is always '1', we can reconstruct the original key in O(1). + const parts = operationRequest.split('.'); + if (parts.length === 5) { + parts[3] = '1'; + const originalKey = parts.join('.'); + const originalTracer = this._currentOperationTracers.get(originalKey); + if (originalTracer) { + tracer = originalTracer; + key = originalKey; + } else { + const prefix = `${parts[0]}.${parts[1]}.${parts[2]}.`; + const suffix = `.${parts[4]}`; + for (const [k, t] of this._currentOperationTracers.entries()) { + if (k.startsWith(prefix) && k.endsWith(suffix)) { + tracer = t; + key = k; + break; + } + } + } + } + } + if (!tracer) { // Attempting to retrieve tracer that doesn't exist. return null; } - this._currentOperationLastUpdatedMs.set(operationRequest, Date.now()); + this._currentOperationLastUpdatedMs.set(key, Date.now()); - return this._currentOperationTracers.get(operationRequest) ?? null; + return tracer; } /** @@ -367,8 +393,30 @@ export class MetricsTracerFactory { * @param requestId The request id of the gRPC call set under 'x-goog-spanner-request-id'. */ public clearCurrentTracer(requestId: string) { - const operationRequest = + let operationRequest = this._extractOperationRequest(requestId) || requestId; + if ( + !this._currentOperationTracers.has(operationRequest) && + operationRequest + ) { + const parts = operationRequest.split('.'); + if (parts.length === 5) { + parts[3] = '1'; + const originalKey = parts.join('.'); + if (this._currentOperationTracers.has(originalKey)) { + operationRequest = originalKey; + } else { + const prefix = `${parts[0]}.${parts[1]}.${parts[2]}.`; + const suffix = `.${parts[4]}`; + for (const k of this._currentOperationTracers.keys()) { + if (k.startsWith(prefix) && k.endsWith(suffix)) { + operationRequest = k; + break; + } + } + } + } + } if (!this._currentOperationTracers.has(operationRequest)) { return; } diff --git a/handwritten/spanner/src/multiplexed-session.ts b/handwritten/spanner/src/multiplexed-session.ts index a9eadce6c56c..dbdcbed9a8ba 100644 --- a/handwritten/spanner/src/multiplexed-session.ts +++ b/handwritten/spanner/src/multiplexed-session.ts @@ -145,6 +145,13 @@ export class MultiplexedSession }, ); this._multiplexedSession = createSessionResponse; + const spanner = (this.database?.parent as any)?.parent; + if ( + spanner?.setPrimeSession && + createSessionResponse?.formattedName_ + ) { + spanner.setPrimeSession(createSessionResponse.formattedName_); + } span.addEvent('Created a multiplexed session'); } catch (e) { setSpanError(span, e as Error); diff --git a/handwritten/spanner/src/transaction-runner.ts b/handwritten/spanner/src/transaction-runner.ts index fe74ca26d354..ea87ef30b3e3 100644 --- a/handwritten/spanner/src/transaction-runner.ts +++ b/handwritten/spanner/src/transaction-runner.ts @@ -246,6 +246,9 @@ export abstract class Runner { } catch (e) { this.session.lastError = e as grpc.ServiceError; lastError = e as grpc.ServiceError; + if ((transaction as any).affinity) { + (transaction as any).affinity.reset(); + } } finally { this.multiplexedSessionPreviousTransactionId = transaction.id; } diff --git a/handwritten/spanner/src/transaction.ts b/handwritten/spanner/src/transaction.ts index b440e7a17ac3..fb05543fe48a 100644 --- a/handwritten/spanner/src/transaction.ts +++ b/handwritten/spanner/src/transaction.ts @@ -20,6 +20,7 @@ import {isEmpty, toArray} from './helper'; import Long = require('long'); import {EventEmitter} from 'events'; import {grpc, CallOptions, ServiceError, Status, GoogleError} from 'google-gax'; +import * as extend from 'extend'; import {common as p} from 'protobufjs'; import {finished, Readable, PassThrough, Stream} from 'stream'; @@ -66,6 +67,7 @@ import { X_GOOG_SPANNER_REQUEST_ID_HEADER, X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR, } from './request_id_header'; +import {AffinityKind, TransactionAffinity} from './channel-pool'; export type Rows = Array; type ResultStats = ResultSetStats | IResultSetStats; @@ -343,6 +345,7 @@ export class Snapshot extends EventEmitter { | null; id?: Uint8Array | string; protected _affinityKey?: string; + protected _affinity?: TransactionAffinity; protected _bindGaxOpts?: CallOptions; protected _unbindGaxOpts?: CallOptions; multiplexedSessionPreviousTransactionId?: Uint8Array | string; @@ -408,12 +411,14 @@ export class Snapshot extends EventEmitter { session: Session, options?: TimestampBounds, queryOptions?: IQueryOptions, + affinityKind: AffinityKind = AffinityKind.ReadOnly, ) { super(); this.ended = false; this.session = session; this.queryOptions = Object.assign({}, queryOptions); + this._affinity = new TransactionAffinity(affinityKind); // If the session is multiplexed, generate a unique affinity key for this // specific transaction/snapshot. This allows requests using the same shared // multiplexed session to be distributed across different gRPC channels. @@ -425,6 +430,7 @@ export class Snapshot extends EventEmitter { otherArgs: { options: { affinityKey: this._affinityKey, + affinity: this._affinity, }, }, }; @@ -435,6 +441,7 @@ export class Snapshot extends EventEmitter { options: { affinityKey: this._affinityKey, unbind: true, + affinity: this._affinity, }, }, }; @@ -464,10 +471,14 @@ export class Snapshot extends EventEmitter { this._mutationKey = null; } + get affinity(): TransactionAffinity | undefined { + return this._affinity; + } + /** - * 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. + * Binds the multiplexed session affinity key and transaction affinity 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 @@ -492,6 +503,7 @@ export class Snapshot extends EventEmitter { 'affinityKey', this._affinityKey, ); + config.gaxOpts = injectGaxOpt(config.gaxOpts, 'affinity', this._affinity); } return config; } @@ -1140,6 +1152,7 @@ export class Snapshot extends EventEmitter { } this.ended = true; + this._affinity?.reset(); this._releaseWaitingRequests(new Error('Transaction has ended.')); process.nextTick(() => this.emit('end')); @@ -2560,13 +2573,48 @@ export class Transaction extends Dml { queryOptions?: IQueryOptions, requestOptions?: Pick, ) { - super(session, undefined, queryOptions); + super(session, undefined, queryOptions, AffinityKind.ReadWrite); this._queuedMutations = []; this._options = {readWrite: options}; this._options.isolationLevel = IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED; this.requestOptions = requestOptions; this._retryCommit = false; + + if (!this._affinityKey && this._affinity) { + const originalRequest = this.request; + this.request = (config: any, callback?: Function) => { + if (!config) { + return originalRequest.call(this, config, callback!); + } + const gaxOptions = injectGaxOpt( + extend(true, {}, config.gaxOpts), + 'affinity', + this._affinity, + ); + return originalRequest.call( + this, + extend(true, {}, config, {gaxOpts: gaxOptions}), + callback!, + ); + }; + + const originalRequestStream = this.requestStream; + this.requestStream = (config: any) => { + if (!config) { + return originalRequestStream.call(this, config); + } + const gaxOptions = injectGaxOpt( + extend(true, {}, config.gaxOpts), + 'affinity', + this._affinity, + ); + return originalRequestStream.call( + this, + extend(true, {}, config, {gaxOpts: gaxOptions}), + ); + }; + } } /** @@ -3344,6 +3392,7 @@ export class Transaction extends Dml { if (!this.id) { span.addEvent('Transaction ID is unknown, nothing to rollback.'); span.end(); + this.end(); callback(null); return; } diff --git a/handwritten/spanner/test/channel-pool.ts b/handwritten/spanner/test/channel-pool.ts new file mode 100644 index 000000000000..afcd95e7577e --- /dev/null +++ b/handwritten/spanner/test/channel-pool.ts @@ -0,0 +1,1796 @@ +/*! + * 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 {describe, it, before, after, beforeEach, afterEach} from 'mocha'; +import * as grpc from '@grpc/grpc-js'; +import * as sinon from 'sinon'; +import {Spanner} from '../src'; +import * as mock from './mockserver/mockspanner'; +import { + AffinityKind, + ChannelEntry, + ChannelPool, + ChannelPoolChannelAdapter, + createCallInvocationTransformer, + DynamicChannelPool, + selectPowerOfTwo, + StaticChannelPool, + TransactionAffinity, +} from '../src/channel-pool'; + +describe('ChannelPool Module', () => { + let sandbox: sinon.SinonSandbox; + + beforeEach(() => { + sandbox = sinon.createSandbox(); + }); + + afterEach(() => { + sandbox.restore(); + }); + + describe('selectPowerOfTwo', () => { + it('should throw if entries list is empty', () => { + assert.throws(() => selectPowerOfTwo([]), /No active channels available/); + }); + + it('should return the only entry if length is 1', () => { + const entry: ChannelEntry = { + id: 1, + channel: {} as grpc.Channel, + inFlightRpcs: 0, + activeRwTransactions: 0, + state: 'ACTIVE', + lastActivity: 100, + }; + assert.strictEqual(selectPowerOfTwo([entry]), entry); + }); + + it('should select the channel with lower load between sampled pair', () => { + const entry1: ChannelEntry = { + id: 1, + channel: {} as grpc.Channel, + inFlightRpcs: 10, + activeRwTransactions: 0, + state: 'ACTIVE', + lastActivity: 100, + }; + const entry2: ChannelEntry = { + id: 2, + channel: {} as grpc.Channel, + inFlightRpcs: 2, + activeRwTransactions: 0, + state: 'ACTIVE', + lastActivity: 100, + }; + + // With 2 entries, both will always be selected + const picked = selectPowerOfTwo([entry1, entry2]); + assert.strictEqual(picked.id, 2); + }); + + it('should break ties using warmer lastActivity', () => { + const coldEntry: ChannelEntry = { + id: 1, + channel: {} as grpc.Channel, + inFlightRpcs: 5, + activeRwTransactions: 0, + state: 'ACTIVE', + lastActivity: 100, + }; + const warmEntry: ChannelEntry = { + id: 2, + channel: {} as grpc.Channel, + inFlightRpcs: 5, + activeRwTransactions: 0, + state: 'ACTIVE', + lastActivity: 500, + }; + + const picked = selectPowerOfTwo([coldEntry, warmEntry]); + assert.strictEqual(picked.id, 2); + }); + }); + + describe('TransactionAffinity', () => { + it('should manage Read/Write affinity pin and decrement activeRwTransactions on reset', () => { + const affinity = new TransactionAffinity(AffinityKind.ReadWrite); + assert.strictEqual(affinity.kind, AffinityKind.ReadWrite); + assert.strictEqual(affinity.pinnedEntry, null); + + const entry: ChannelEntry = { + id: 1, + channel: {} as grpc.Channel, + inFlightRpcs: 0, + activeRwTransactions: 1, + state: 'ACTIVE', + lastActivity: 100, + }; + + affinity.pinnedEntry = entry; + affinity.reset(); + + assert.strictEqual(affinity.pinnedEntry, null); + assert.strictEqual(entry.activeRwTransactions, 0); + }); + + it('should not decrement activeRwTransactions on ReadOnly affinity reset', () => { + const affinity = new TransactionAffinity(AffinityKind.ReadOnly); + const entry: ChannelEntry = { + id: 1, + channel: {} as grpc.Channel, + inFlightRpcs: 0, + activeRwTransactions: 0, + state: 'ACTIVE', + lastActivity: 100, + }; + + affinity.pinnedEntry = entry; + affinity.reset(); + + assert.strictEqual(affinity.pinnedEntry, null); + assert.strictEqual(entry.activeRwTransactions, 0); + }); + }); + + describe('StaticChannelPool', () => { + it('should create default 4 channels and handle lease acquisition and release', async () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + ); + + assert.strictEqual(pool.size, 4); + assert.strictEqual(pool.activeCount, 4); + + const lease1 = pool.acquire(); + assert.strictEqual(lease1.entry.inFlightRpcs, 1); + + const lease2 = pool.acquire(); + assert.strictEqual(lease2.entry.inFlightRpcs >= 1, true); + + lease1.release(); + assert.strictEqual(lease1.entry.inFlightRpcs, 0); + + // Releasing twice should be a no-op + lease1.release(); + assert.strictEqual(lease1.entry.inFlightRpcs, 0); + + await pool.close(); + }); + + it('should pin Read/Write transactions to the same channel', async () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: 4}, + ); + + const affinity = new TransactionAffinity(AffinityKind.ReadWrite); + + const lease1 = pool.acquire(affinity); + const pinnedChannelId = lease1.entry.id; + assert.strictEqual(lease1.entry.activeRwTransactions, 1); + + // Subsequent acquisitions using the same affinity must hit the pinned channel + for (let i = 0; i < 5; i++) { + const subsequentLease = pool.acquire(affinity); + assert.strictEqual(subsequentLease.entry.id, pinnedChannelId); + subsequentLease.release(); + } + + lease1.release(); + assert.strictEqual(lease1.entry.activeRwTransactions, 1); + + affinity.reset(); + assert.strictEqual(lease1.entry.activeRwTransactions, 0); + + await pool.close(); + }); + + it('should respect SPANNER_NUM_CHANNELS environment variable', async () => { + const prev = process.env.SPANNER_NUM_CHANNELS; + try { + process.env.SPANNER_NUM_CHANNELS = '6'; + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + ); + assert.strictEqual(pool.size, 6); + await pool.close(); + } finally { + if (prev !== undefined) { + process.env.SPANNER_NUM_CHANNELS = prev; + } else { + delete process.env.SPANNER_NUM_CHANNELS; + } + } + }); + + it('should prioritize options.numChannels over SPANNER_NUM_CHANNELS environment variable', async () => { + const prev = process.env.SPANNER_NUM_CHANNELS; + try { + process.env.SPANNER_NUM_CHANNELS = '6'; + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: 2}, + ); + assert.strictEqual(pool.size, 2); + await pool.close(); + } finally { + if (prev !== undefined) { + process.env.SPANNER_NUM_CHANNELS = prev; + } else { + delete process.env.SPANNER_NUM_CHANNELS; + } + } + }); + + it('should return channels array from getChannels()', async () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: 2}, + ); + const channels = pool.getChannels(); + assert.strictEqual(channels.length, 2); + assert.ok(channels[0] instanceof grpc.Channel); + await pool.close(); + }); + + it('should clamp numChannels to at least 1 when 0 is configured', async () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: 0}, + ); + assert.strictEqual(pool.size, 1); + assert.strictEqual(pool.getTarget(), 'localhost:9010'); + await pool.close(); + }); + + it('should fallback to default 4 channels when numChannels is NaN', async () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: NaN}, + ); + assert.strictEqual(pool.size, 4); + await pool.close(); + }); + + it('should throw when acquiring from closed pool and report SHUTDOWN connectivity state', async () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: 2}, + ); + await pool.close(); + assert.throws(() => pool.acquire(), /Channel pool is closed\./); + assert.strictEqual( + pool.getConnectivityState(), + grpc.connectivityState.SHUTDOWN, + ); + }); + + it('should report TRANSIENT_FAILURE connectivity state when channel fails and none are ready or connecting', async () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: 2}, + ); + const entries = (pool as any).entries; + sinon + .stub(entries[0].channel, 'getConnectivityState') + .returns(grpc.connectivityState.TRANSIENT_FAILURE); + sinon + .stub(entries[1].channel, 'getConnectivityState') + .returns(grpc.connectivityState.IDLE); + + assert.strictEqual( + pool.getConnectivityState(), + grpc.connectivityState.TRANSIENT_FAILURE, + ); + await pool.close(); + }); + }); + + describe('DynamicChannelPool', () => { + it('should scale up under concurrency and prime new channels with SELECT 1', async () => { + let primedSession = ''; + let primeCalls = 0; + + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 2, + maxChannels: 5, + maxRpcPerChannel: 2, + scaleUpCooldownMs: 0, + primeFn: async (_channel, sessionName) => { + primeCalls++; + primedSession = sessionName; + }, + }, + ); + + pool.setPrimeSession( + 'projects/p/instances/i/databases/d/sessions/s-prime', + ); + + assert.strictEqual(pool.size, 2); + assert.strictEqual(pool.activeCount, 2); + + // Acquire enough leases on one channel to exceed maxRpcPerChannel = 2 + const leases: import('../src/channel-pool').ChannelLease[] = []; + for (let i = 0; i < 6; i++) { + leases.push(pool.acquire()); + } + + // Wait for scale-up setImmediate task to execute + const deadline = Date.now() + 1000; + while (pool.activeCount <= 2 && Date.now() < deadline) { + await new Promise(resolve => setImmediate(resolve)); + } + + assert.strictEqual(pool.activeCount > 2, true); + assert.strictEqual(primeCalls > 0, true); + assert.strictEqual( + primedSession, + 'projects/p/instances/i/databases/d/sessions/s-prime', + ); + + for (const lease of leases) { + lease.release(); + } + + await pool.close(); + }); + + it('should use default configuration options', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + ); + + assert.strictEqual((pool as any).minChannels, 4); + assert.strictEqual((pool as any).maxChannels, 256); + assert.strictEqual((pool as any).maxRpcPerChannel, 8); + assert.strictEqual((pool as any).minRpcPerChannel, 2); + assert.strictEqual((pool as any).maxScaleUpPercent, 100); + assert.strictEqual((pool as any).maxRemoveChannels, 4); + assert.strictEqual((pool as any).scaleUpCooldownMs, 1000); + + await pool.close(); + }); + + it('should scale up aggressively adding multiple channels under heavy load', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 4, + maxChannels: 16, + maxRpcPerChannel: 8, + minRpcPerChannel: 2, + scaleUpCooldownMs: 0, + }, + ); + + assert.strictEqual(pool.activeCount, 4); + + // Total load 40 across 4 channels -> targetRpc 5 -> desiredChannels 8 -> adds 4 channels in parallel + const leases: import('../src/channel-pool').ChannelLease[] = []; + for (let i = 0; i < 40; i++) { + leases.push(pool.acquire()); + } + + const deadline = Date.now() + 1000; + while (pool.activeCount < 8 && Date.now() < deadline) { + await new Promise(resolve => setImmediate(resolve)); + } + + assert.strictEqual(pool.activeCount, 8); + + for (const lease of leases) { + lease.release(); + } + + await pool.close(); + }); + + it('should scale up without veto when a single channel exceeds maxRpcPerChannel even if total load is low', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 4, + maxChannels: 8, + maxRpcPerChannel: 8, + minRpcPerChannel: 2, + scaleUpCooldownMs: 0, + }, + ); + + assert.strictEqual(pool.activeCount, 4); + + // Artificially put 9 RPCs on channel 0 while other channels have 0 load. + // totalLoad = 9. targetRpc = 5. desiredChannels = ceil(9/5) = 2 <= 4. + // If a veto existed, it would not scale up because desiredChannels <= currentLen. + // Without veto, it must still add at least 1 channel because channel 0 exceeded maxRpcPerChannel = 8. + const activeEntries = (pool as any).activeEntries; + activeEntries[0].inFlightRpcs = 9; + (pool as any).maybeScaleUp(); + + const deadline = Date.now() + 1000; + while (pool.activeCount <= 4 && Date.now() < deadline) { + await new Promise(resolve => setImmediate(resolve)); + } + + assert.strictEqual(pool.activeCount > 4, true); + + activeEntries[0].inFlightRpcs = 0; + await pool.close(); + }); + + it('should keep Read/Write transaction on draining channel until completion', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 1, + maxChannels: 3, + }, + ); + + const affinity = new TransactionAffinity(AffinityKind.ReadWrite); + const lease = pool.acquire(affinity); + const entry = lease.entry; + + // Simulate channel entering DRAINING + entry.state = 'DRAINING'; + + // R/W transaction should remain pinned even while DRAINING + const subsequentLease = pool.acquire(affinity); + assert.strictEqual(subsequentLease.entry, entry); + + subsequentLease.release(); + lease.release(); + + affinity.reset(); + assert.strictEqual(entry.activeRwTransactions, 0); + + await pool.close(); + }); + + it('should seamlessly switch Read-Only transaction away from draining channel', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 2, + maxChannels: 4, + }, + ); + + const affinity = new TransactionAffinity(AffinityKind.ReadOnly); + const lease1 = pool.acquire(affinity); + const oldEntry = lease1.entry; + lease1.release(); + + // Simulate draining channel by removing from activeEntries and moving to drainingEntries + const idx = (pool as any).activeEntries.indexOf(oldEntry); + (pool as any).activeEntries.splice(idx, 1); + oldEntry.state = 'DRAINING'; + (pool as any).drainingEntries.push(oldEntry); + + // Soft affinity should re-pick a fresh active channel + const lease2 = pool.acquire(affinity); + assert.notStrictEqual(lease2.entry, oldEntry); + assert.strictEqual(lease2.entry.state, 'ACTIVE'); + + lease2.release(); + affinity.reset(); + + await pool.close(); + }); + + it('should clamp minChannels to at least 1 when 0 is configured', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 0}, + ); + assert.strictEqual((pool as any).minChannels, 1); + assert.strictEqual(pool.activeCount, 1); + await pool.close(); + }); + + it('should fallback to default minChannels and maxChannels when NaN is configured', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: NaN, maxChannels: NaN}, + ); + assert.strictEqual((pool as any).minChannels, 4); + assert.strictEqual((pool as any).maxChannels, 256); + assert.strictEqual(pool.activeCount, 4); + await pool.close(); + }); + + it('should correctly track activeRwTransactions when fallback occurs on closed channel', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 2, maxChannels: 2}, + ); + + const affinity = new TransactionAffinity(AffinityKind.ReadWrite); + const lease1 = pool.acquire(affinity); + const pinnedEntry = lease1.entry; + assert.strictEqual(pinnedEntry.activeRwTransactions, 1); + lease1.release(); + + // Simulate channel closure and removal from active pool + const idx = (pool as any).activeEntries.indexOf(pinnedEntry); + (pool as any).activeEntries.splice(idx, 1); + pinnedEntry.state = 'CLOSED'; + + // Next acquire with same affinity should fall back to another active channel + const lease2 = pool.acquire(affinity); + assert.notStrictEqual(lease2.entry, pinnedEntry); + assert.strictEqual(pinnedEntry.activeRwTransactions, 0); + assert.strictEqual(lease2.entry.activeRwTransactions, 1); + lease2.release(); + + affinity.reset(); + assert.strictEqual(lease2.entry.activeRwTransactions, 0); + + await pool.close(); + }); + + it('should clean up draining channel when its last transaction affinity resets', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 2, maxChannels: 4}, + ); + + const affinity = new TransactionAffinity(AffinityKind.ReadWrite); + const lease1 = pool.acquire(affinity); + const entry = lease1.entry; + assert.strictEqual(entry.activeRwTransactions, 1); + + // Drain entry + const idx = (pool as any).activeEntries.indexOf(entry); + (pool as any).activeEntries.splice(idx, 1); + entry.state = 'DRAINING'; + (pool as any).drainingEntries.push(entry); + + // Release RPC lease while transaction is still open + lease1.release(); + assert.strictEqual(entry.inFlightRpcs, 0); + assert.strictEqual(entry.state, 'DRAINING'); + assert.strictEqual((pool as any).drainingEntries.length, 1); + + // Transaction commits / resets affinity -> immediately triggers draining cleanup hook + affinity.reset(); + assert.strictEqual(entry.activeRwTransactions, 0); + assert.strictEqual(entry.state, 'CLOSED'); + assert.strictEqual((pool as any).drainingEntries.length, 0); + + await pool.close(); + }); + + it('should scale down in batches up to maxRemoveChannels after sustained low load', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 2, + maxChannels: 8, + minRpcPerChannel: 2, + maxRpcPerChannel: 8, + maxRemoveChannels: 3, + }, + ); + + // Manually expand pool to 8 active entries + while ((pool as any).activeEntries.length < 8) { + (pool as any).activeEntries.push((pool as any).createEntry()); + } + assert.strictEqual(pool.activeCount, 8); + + // Debouncing requires 3 consecutive low-load cycles + (pool as any).evaluateScaleDown(); + assert.strictEqual(pool.activeCount, 8); + assert.strictEqual((pool as any).consecutiveLowLoadChecks, 1); + + (pool as any).evaluateScaleDown(); + assert.strictEqual(pool.activeCount, 8); + assert.strictEqual((pool as any).consecutiveLowLoadChecks, 2); + + // 3rd low-load cycle triggers batched scale-down: + // desiredChannels = max(2, ceil(0 / 5)) = 2. + // excess = 8 - 2 = 6. + // maxRemoveChannels = 3. + // Drains 3 channels in this cycle. + (pool as any).evaluateScaleDown(); + assert.strictEqual(pool.activeCount, 5); + assert.strictEqual((pool as any).consecutiveLowLoadChecks, 0); + + // Next 3 low-load cycles drain remaining 3 excess channels down to minChannels = 2 + (pool as any).evaluateScaleDown(); + (pool as any).evaluateScaleDown(); + (pool as any).evaluateScaleDown(); + assert.strictEqual(pool.activeCount, 2); + + // Subsequent cycles do not scale down below minChannels + (pool as any).evaluateScaleDown(); + (pool as any).evaluateScaleDown(); + (pool as any).evaluateScaleDown(); + assert.strictEqual(pool.activeCount, 2); + + await pool.close(); + }); + + it('should prioritize draining channels with lower in-flight RPCs, fewer R/W transactions, and newer creation time', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 1, + maxChannels: 4, + maxRemoveChannels: 2, + }, + ); + + const entries: ChannelEntry[] = (pool as any).activeEntries; + while (entries.length < 4) { + entries.push((pool as any).createEntry()); + } + + // Entry 0: high in-flight load (preserved) + const busyChannel = entries[0]; + busyChannel.inFlightRpcs = 5; + busyChannel.activeRwTransactions = 0; + busyChannel.createdAt = 1000; + + // Entry 1: 0 in-flight, but active R/W transaction (preserved over entries with 0 transactions) + const rwPinnedChannel = entries[1]; + rwPinnedChannel.inFlightRpcs = 0; + rwPinnedChannel.activeRwTransactions = 1; + rwPinnedChannel.createdAt = 2000; + + // Entry 2: 0 in-flight, 0 R/W transaction, older (createdAt = 3000) + const olderIdleChannel = entries[2]; + olderIdleChannel.inFlightRpcs = 0; + olderIdleChannel.activeRwTransactions = 0; + olderIdleChannel.createdAt = 3000; + + // Entry 3: 0 in-flight, 0 R/W transaction, newer (createdAt = 4000) (top drain candidate) + const newerIdleChannel = entries[3]; + newerIdleChannel.inFlightRpcs = 0; + newerIdleChannel.activeRwTransactions = 0; + newerIdleChannel.createdAt = 4000; + + // Trigger scale-down of 2 channels (maxRemoveChannels = 2) + (pool as any).consecutiveLowLoadChecks = 2; + (pool as any).evaluateScaleDown(); + + // Exactly 2 channels removed: newerIdleChannel and olderIdleChannel + assert.strictEqual(pool.activeCount, 2); + assert.strictEqual(entries.includes(busyChannel), true); + assert.strictEqual(entries.includes(rwPinnedChannel), true); + assert.strictEqual(entries.includes(olderIdleChannel), false); + assert.strictEqual(entries.includes(newerIdleChannel), false); + + await pool.close(); + }); + + it('should prefer draining newer channels over older warm channels when load is tied', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 2, + maxChannels: 4, + maxRemoveChannels: 1, + }, + ); + + const entries: ChannelEntry[] = (pool as any).activeEntries; + while (entries.length < 3) { + entries.push((pool as any).createEntry()); + } + + const warmOldChannel = entries[0]; + warmOldChannel.inFlightRpcs = 0; + warmOldChannel.activeRwTransactions = 0; + warmOldChannel.createdAt = 1000; + + const warmMidChannel = entries[1]; + warmMidChannel.inFlightRpcs = 0; + warmMidChannel.activeRwTransactions = 0; + warmMidChannel.createdAt = 2000; + + const freshNewChannel = entries[2]; + freshNewChannel.inFlightRpcs = 0; + freshNewChannel.activeRwTransactions = 0; + freshNewChannel.createdAt = 3000; + + (pool as any).consecutiveLowLoadChecks = 2; + (pool as any).evaluateScaleDown(); + + // Only 1 channel removed (maxRemoveChannels = 1), and it must be the newest channel + assert.strictEqual(pool.activeCount, 2); + assert.strictEqual(entries.includes(warmOldChannel), true); + assert.strictEqual(entries.includes(warmMidChannel), true); + assert.strictEqual(entries.includes(freshNewChannel), false); + + await pool.close(); + }); + + it('should allocate logical channel IDs in range [1, poolSize] and recycle lowest available slots', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 2, + maxChannels: 6, + }, + ); + + // Initially minChannels = 2, IDs should be 1 and 2 + const active: ChannelEntry[] = (pool as any).activeEntries; + assert.strictEqual(active.length, 2); + assert.strictEqual(active[0].id, 1); + assert.strictEqual(active[1].id, 2); + + // Adding 2 more channels allocates slots 3 and 4 + const entry3 = (pool as any).createEntry(); + active.push(entry3); + assert.strictEqual(entry3.id, 3); + + const entry4 = (pool as any).createEntry(); + active.push(entry4); + assert.strictEqual(entry4.id, 4); + + // Drain and close entry 2 (slot 2 is freed) + const entry2 = active[1]; + active.splice(1, 1); + entry2.state = 'CLOSED'; + entry2.channel.close(); + + // Next channel allocation should recycle the lowest available slot (slot 2) + const recycledEntry = (pool as any).createEntry(); + active.push(recycledEntry); + assert.strictEqual(recycledEntry.id, 2); + + // Verify slot held by an unclosed draining channel is not reused + // Drain entry 3 but keep it DRAINING (e.g. in-flight RPC) + active.splice(1, 1); // removes entry3 + entry3.state = 'DRAINING'; + (pool as any).drainingEntries.push(entry3); + + // Next channel allocation: slots 1, 2, 4 are active, slot 3 is draining + // Slot 3 must NOT be reused while still draining; should allocate slot 5 + const overflowEntry = (pool as any).createEntry(); + active.push(overflowEntry); + assert.strictEqual(overflowEntry.id, 5); + + // Now entry 3 finishes draining and is closed + entry3.state = 'CLOSED'; + entry3.channel.close(); + + // Next channel allocation should now recycle slot 3 + const recycledEntry3 = (pool as any).createEntry(); + active.push(recycledEntry3); + assert.strictEqual(recycledEntry3.id, 3); + + await pool.close(); + }); + + it('should throw when acquiring from closed pool and report SHUTDOWN connectivity state', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 2}, + ); + await pool.close(); + assert.throws(() => pool.acquire(), /Channel pool is closed\./); + assert.strictEqual( + pool.getConnectivityState(), + grpc.connectivityState.SHUTDOWN, + ); + }); + + it('should report TRANSIENT_FAILURE connectivity state when channel fails and none are ready or connecting', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 2}, + ); + const activeEntries = (pool as any).activeEntries; + sinon + .stub(activeEntries[0].channel, 'getConnectivityState') + .returns(grpc.connectivityState.TRANSIENT_FAILURE); + sinon + .stub(activeEntries[1].channel, 'getConnectivityState') + .returns(grpc.connectivityState.IDLE); + + assert.strictEqual( + pool.getConnectivityState(), + grpc.connectivityState.TRANSIENT_FAILURE, + ); + await pool.close(); + }); + + it('should respect SPANNER_NUM_CHANNELS as initial channel count while retaining dynamic maxChannels', async () => { + const previousEnv = process.env.SPANNER_NUM_CHANNELS; + try { + process.env.SPANNER_NUM_CHANNELS = '6'; + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + ); + assert.strictEqual(pool.size, 6); + assert.strictEqual((pool as any).minChannels, 6); + assert.strictEqual((pool as any).maxChannels, 256); + await pool.close(); + } finally { + if (previousEnv !== undefined) { + process.env.SPANNER_NUM_CHANNELS = previousEnv; + } else { + delete process.env.SPANNER_NUM_CHANNELS; + } + } + }); + + it('should wait for pending scale-up to complete when close is called', async () => { + let scaleUpFinished = false; + let resolvePrime: () => void; + const primeBlockedPromise = new Promise(resolve => { + resolvePrime = resolve; + }); + + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 1, + maxChannels: 4, + maxRpcPerChannel: 1, + scaleUpCooldownMs: 0, + primeFn: async () => { + await primeBlockedPromise; + scaleUpFinished = true; + }, + }, + ); + pool.setPrimeSession('projects/p/instances/i/databases/d/sessions/s'); + + // Trigger scale-up + const lease1 = pool.acquire(); + const lease2 = pool.acquire(); + assert.strictEqual((pool as any).isScalingUp, true); + + // Allow setImmediate to run so scale-up starts priming + await new Promise(resolve => setImmediate(resolve)); + + // Close pool while scale-up is running; close() should await scaleUpPromise + const closePromise = pool.close(); + assert.strictEqual(scaleUpFinished, false); + + // Unblock primeFn + resolvePrime!(); + await closePromise; + + assert.strictEqual(scaleUpFinished, true); + assert.strictEqual(pool.size, 0); + + lease1.release(); + lease2.release(); + }); + + it('should prioritize options.minChannels over SPANNER_NUM_CHANNELS environment variable', async () => { + const prev = process.env.SPANNER_NUM_CHANNELS; + try { + process.env.SPANNER_NUM_CHANNELS = '8'; + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 2}, + ); + assert.strictEqual((pool as any).minChannels, 2); + assert.strictEqual(pool.activeCount, 2); + await pool.close(); + } finally { + if (prev !== undefined) { + process.env.SPANNER_NUM_CHANNELS = prev; + } else { + delete process.env.SPANNER_NUM_CHANNELS; + } + } + }); + + it('should clamp minRpcPerChannel when maxRpcPerChannel is small', async () => { + const pool1 = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {maxRpcPerChannel: 1}, + ); + assert.strictEqual((pool1 as any).minRpcPerChannel, 0); + await pool1.close(); + + const pool2 = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {maxRpcPerChannel: 2}, + ); + assert.strictEqual((pool2 as any).minRpcPerChannel, 1); + await pool2.close(); + }); + + it('should abort evaluateScaleDown early when isScalingUp is true or during cooldown', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 1, maxChannels: 4, scaleUpCooldownMs: 60000}, + ); + while ((pool as any).activeEntries.length < 3) { + (pool as any).activeEntries.push((pool as any).createEntry()); + } + assert.strictEqual(pool.activeCount, 3); + + // 1. Abort when isScalingUp is true + (pool as any).isScalingUp = true; + (pool as any).evaluateScaleDown(); + assert.strictEqual((pool as any).consecutiveLowLoadChecks, 0); + + // 2. Abort when within cooldown + (pool as any).isScalingUp = false; + (pool as any).lastScaleUpTime = Date.now(); + (pool as any).evaluateScaleDown(); + assert.strictEqual((pool as any).consecutiveLowLoadChecks, 0); + + // 3. When cooldown has passed, scale-down evaluation proceeds + (pool as any).lastScaleUpTime = Date.now() - 70000; + (pool as any).evaluateScaleDown(); + assert.strictEqual((pool as any).consecutiveLowLoadChecks, 1); + + await pool.close(); + }); + + it('should retry priming a channel up to 2 times on failure', async () => { + let attempts = 0; + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + { + minChannels: 1, + primeFn: async () => { + attempts++; + if (attempts === 1) { + throw new Error('Transient network error'); + } + }, + }, + ); + pool.setPrimeSession('projects/p/instances/i/databases/d/sessions/s'); + + const entry = (pool as any).createEntry(); + await (pool as any).primeChannel(entry); + assert.strictEqual(attempts, 2); + + await pool.close(); + }); + + it('should return channels from getChannels() including draining channels', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 2, maxChannels: 4}, + ); + const drainingEntry = (pool as any).createEntry(); + drainingEntry.state = 'DRAINING'; + (pool as any).drainingEntries.push(drainingEntry); + + const channels = pool.getChannels(); + assert.strictEqual(channels.length, 3); + assert.ok(channels.includes(drainingEntry.channel)); + + await pool.close(); + }); + + it('should re-watch another active channel if watched channel shuts down during scale down while pool remains healthy', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 2, maxChannels: 4}, + ); + const activeEntries = (pool as any).activeEntries; + assert.strictEqual(activeEntries.length, 2); + + const firstChannel = activeEntries[0].channel; + const secondChannel = activeEntries[1].channel; + + let firstChannelCallback: ((error?: Error) => void) | undefined; + sinon + .stub(firstChannel, 'watchConnectivityState') + .callsFake((_state: any, _deadline: any, callback: any) => { + firstChannelCallback = callback; + }); + + let secondChannelWatched = false; + sinon + .stub(secondChannel, 'watchConnectivityState') + .callsFake((_state: any, _deadline: any, callback: any) => { + secondChannelWatched = true; + callback(); + }); + + sinon + .stub(secondChannel, 'getConnectivityState') + .returns(grpc.connectivityState.READY); + + let poolCallbackFired = false; + const watchPromise = new Promise((resolve, reject) => { + pool.watchConnectivityState( + grpc.connectivityState.READY, + Infinity, + error => { + if (error) { + reject(error); + } else { + poolCallbackFired = true; + resolve(); + } + }, + ); + }); + + // Simulate first channel transitioning to SHUTDOWN due to scale down + sinon + .stub(firstChannel, 'getConnectivityState') + .returns(grpc.connectivityState.SHUTDOWN); + // Remove first channel from activeEntries (simulating drain/removal) + activeEntries.shift(); + + // Trigger first channel watch callback without error + firstChannelCallback!(); + await watchPromise; + + assert.strictEqual(poolCallbackFired, true); + assert.strictEqual(secondChannelWatched, true); + await pool.close(); + }); + + it('should notify caller when watched channel shutdown causes pool connectivity state to change', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 1, maxChannels: 2}, + ); + const activeEntries = (pool as any).activeEntries; + const channel = activeEntries[0].channel; + + let channelCallback: ((error?: Error) => void) | undefined; + sinon + .stub(channel, 'watchConnectivityState') + .callsFake((_state: any, _deadline: any, callback: any) => { + channelCallback = callback; + }); + + sinon + .stub(channel, 'getConnectivityState') + .returns(grpc.connectivityState.SHUTDOWN); + + const watchPromise = new Promise((resolve, reject) => { + pool.watchConnectivityState( + grpc.connectivityState.READY, + Infinity, + error => { + if (error) { + reject(error); + } else { + resolve(); + } + }, + ); + }); + + // Channel shuts down and pool state changes from READY to SHUTDOWN + activeEntries.length = 0; + channelCallback!(); + await watchPromise; + await pool.close(); + }); + }); + + describe('createCallInvocationTransformer', () => { + it('should route call to acquired channel and append channel id to request id header', done => { + const fakeChannel: any = { + createCall: sinon.stub(), + }; + const fakeEntry: ChannelEntry = { + id: 3, + channel: fakeChannel, + inFlightRpcs: 0, + activeRwTransactions: 0, + state: 'ACTIVE', + lastActivity: Date.now(), + }; + + const mockPool: ChannelPool = { + acquire: sinon.stub().returns({ + entry: fakeEntry, + release: sinon.stub(), + }), + size: 1, + activeCount: 1, + close: async () => {}, + getConnectivityState: () => grpc.connectivityState.READY, + watchConnectivityState: () => {}, + getTarget: () => 'localhost:9010', + getChannels: () => [fakeChannel], + }; + + const transformer = createCallInvocationTransformer(mockPool); + const metadata = new grpc.Metadata(); + metadata.set('x-goog-spanner-request-id', '1.abcd1234.1.1.5.1'); + + const callProperties: any = { + metadata, + callOptions: {}, + argument: {}, + methodDefinition: {path: '/google.spanner.v1.Spanner/ExecuteSql'}, + }; + + const transformed = transformer(callProperties); + + assert.strictEqual(transformed.channel, fakeChannel); + assert.strictEqual( + metadata.get('x-goog-spanner-request-id')[0], + '1.abcd1234.1.3.5.1', + ); + assert.strictEqual(transformed.callOptions.interceptors!.length, 1); + + done(); + }); + }); + + describe('ChannelPoolChannelAdapter', () => { + it('should release acquired lease when call completes with status', () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: 1}, + ); + const adapter = new ChannelPoolChannelAdapter(pool); + const entry = (pool as any).entries[0]; + assert.strictEqual(entry.inFlightRpcs, 0); + + const fakeCall = { + start: sinon.stub(), + cancel: sinon.stub(), + }; + sinon.stub(entry.channel as any, 'createCall').returns(fakeCall); + + const interceptedCall = adapter.createCall( + '/google.spanner.v1.Spanner/BatchCreateSessions', + Infinity, + 'localhost:9010', + null, + 0, + ); + + assert.strictEqual(entry.inFlightRpcs, 1); + + const fakeListener = { + onReceiveStatus: sinon.stub(), + }; + interceptedCall.start(new grpc.Metadata(), fakeListener); + + const passedListener = fakeCall.start.firstCall.args[1]; + passedListener.onReceiveStatus({code: grpc.status.OK, details: 'OK'}); + + assert.strictEqual(entry.inFlightRpcs, 0); + assert.strictEqual(fakeListener.onReceiveStatus.calledOnce, true); + }); + + it('should release acquired lease when call is cancelled', () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: 1}, + ); + const adapter = new ChannelPoolChannelAdapter(pool); + const entry = (pool as any).entries[0]; + assert.strictEqual(entry.inFlightRpcs, 0); + + const fakeCall = { + start: sinon.stub(), + cancelWithStatus: sinon.stub(), + }; + sinon.stub(entry.channel as any, 'createCall').returns(fakeCall); + + const interceptedCall = adapter.createCall( + '/google.spanner.v1.Spanner/BatchCreateSessions', + Infinity, + 'localhost:9010', + null, + 0, + ); + + assert.strictEqual(entry.inFlightRpcs, 1); + + interceptedCall.cancelWithStatus(grpc.status.CANCELLED, 'Cancelled'); + assert.strictEqual(entry.inFlightRpcs, 0); + assert.strictEqual(fakeCall.cancelWithStatus.calledOnce, true); + }); + + it('should delegate sendMessage and halfClose on InterceptingCall', () => { + const pool = new StaticChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {numChannels: 1}, + ); + const adapter = new ChannelPoolChannelAdapter(pool); + const entry = (pool as any).entries[0]; + + const fakeCall = { + start: sinon.stub(), + sendMessageWithContext: sinon.stub(), + halfClose: sinon.stub(), + }; + sinon.stub(entry.channel as any, 'createCall').returns(fakeCall); + + const interceptedCall = adapter.createCall( + '/google.spanner.v1.Spanner/ExecuteStreamingSql', + Infinity, + 'localhost:9010', + null, + 0, + ); + + interceptedCall.sendMessage({query: 'SELECT 1'}); + assert.strictEqual(fakeCall.sendMessageWithContext.calledOnce, true); + + interceptedCall.halfClose(); + assert.strictEqual(fakeCall.halfClose.calledOnce, true); + }); + + it('should include both active and draining channels in channelRefs for DynamicChannelPool', async () => { + const pool = new DynamicChannelPool( + 'localhost:9010', + grpc.credentials.createInsecure(), + {}, + {minChannels: 2, maxChannels: 4}, + ); + const adapter = new ChannelPoolChannelAdapter(pool); + + assert.strictEqual(adapter.channelRefs.length, 2); + + // Simulate a draining channel + const drainingEntry = (pool as any).createEntry(); + drainingEntry.state = 'DRAINING'; + (pool as any).drainingEntries.push(drainingEntry); + + assert.strictEqual(adapter.channelRefs.length, 3); + assert.strictEqual( + adapter.channelRefs.some(ref => ref.channel === drainingEntry.channel), + true, + ); + + await pool.close(); + }); + }); + + describe('Mock Spanner Channel Pool Integration', () => { + let server: grpc.Server; + let spannerMock: mock.MockSpanner; + let port: number; + + before(async () => { + server = new grpc.Server(); + spannerMock = mock.createMockSpanner(server); + port = await new Promise((resolve, reject) => { + server.bindAsync( + 'localhost:0', + grpc.ServerCredentials.createInsecure(), + (err, assignedPort) => { + if (err) { + reject(err); + } else { + resolve(assignedPort); + } + }, + ); + }); + spannerMock.putStatementResult( + 'SELECT 1', + mock.StatementResult.resultSet(mock.createSelect1ResultSet()), + ); + }); + + after(done => { + server.tryShutdown(done); + }); + + it('should default to StaticChannelPool with 4 channels when no channel pool options are specified', async () => { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + }); + + const database = spanner.instance('instance').database('database'); + const [rows] = await database.run('SELECT 1'); + assert.strictEqual(rows.length, 1); + assert(spanner.channelPool); + assert(spanner.channelPool instanceof StaticChannelPool); + assert.strictEqual(spanner.channelPool.size, 4); + + await spanner.close(); + }); + + it('should default to StaticChannelPool with SPANNER_NUM_CHANNELS channels when env var is set', async () => { + const previousEnv = process.env.SPANNER_NUM_CHANNELS; + process.env.SPANNER_NUM_CHANNELS = '6'; + try { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + }); + + const database = spanner.instance('instance').database('database'); + const [rows] = await database.run('SELECT 1'); + assert.strictEqual(rows.length, 1); + assert(spanner.channelPool); + assert(spanner.channelPool instanceof StaticChannelPool); + assert.strictEqual(spanner.channelPool.size, 6); + + await spanner.close(); + } finally { + if (previousEnv !== undefined) { + process.env.SPANNER_NUM_CHANNELS = previousEnv; + } else { + delete process.env.SPANNER_NUM_CHANNELS; + } + } + }); + + it('should allow falling back to legacy grpc-gcp pool via channelPool option', async () => { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + channelPool: 'grpc-gcp', + }); + + assert.strictEqual(spanner.channelPool, undefined); + await spanner.close(); + }); + + it('should initialize StaticChannelPool and execute queries', async () => { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + channelPool: { + type: 'static', + numChannels: 4, + }, + }); + + const database = spanner.instance('instance').database('database'); + const [rows] = await database.run('SELECT 1'); + assert.strictEqual(rows.length, 1); + assert(spanner.channelPool); + assert.strictEqual(spanner.channelPool.size, 4); + + await spanner.close(); + }); + + it('should initialize DynamicChannelPool, execute queries and support scale-up', async () => { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + channelPool: { + type: 'dynamic', + minChannels: 2, + maxChannels: 6, + }, + }); + + const database = spanner.instance('instance').database('database'); + const [rows] = await database.run('SELECT 1'); + assert.strictEqual(rows.length, 1); + assert(spanner.channelPool); + assert.strictEqual(spanner.channelPool.size, 2); + + await spanner.close(); + }); + + it('should activate dynamic channel pool via SPANNER_CHANNEL_POOL environment variable', async () => { + process.env.SPANNER_CHANNEL_POOL = 'dynamic'; + process.env.SPANNER_NUM_CHANNELS = '3'; + try { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + }); + + const database = spanner.instance('instance').database('database'); + const [rows] = await database.run('SELECT 1'); + assert.strictEqual(rows.length, 1); + assert(spanner.channelPool); + assert.strictEqual(spanner.channelPool.size, 3); + + await spanner.close(); + } finally { + delete process.env.SPANNER_CHANNEL_POOL; + delete process.env.SPANNER_NUM_CHANNELS; + } + }); + + it('should prefer options.channelPool over SPANNER_CHANNEL_POOL environment variable', async () => { + process.env.SPANNER_CHANNEL_POOL = 'dynamic'; + try { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + channelPool: { + type: 'static', + numChannels: 2, + }, + }); + + const database = spanner.instance('instance').database('database'); + const [rows] = await database.run('SELECT 1'); + assert.strictEqual(rows.length, 1); + assert(spanner.channelPool); + assert(spanner.channelPool instanceof StaticChannelPool); + assert.strictEqual(spanner.channelPool.size, 2); + + await spanner.close(); + } finally { + delete process.env.SPANNER_CHANNEL_POOL; + } + }); + + it('should route all RPCs in a read/write transaction to the same channel', async () => { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + channelPool: { + type: 'static', + numChannels: 4, + }, + }); + + const database = spanner.instance('instance').database('database'); + spannerMock.resetRequests(); + + await database.runTransactionAsync(async transaction => { + const [rows1] = await transaction.run('SELECT 1'); + assert.strictEqual(rows1.length, 1); + const [rows2] = await transaction.run('SELECT 1'); + assert.strictEqual(rows2.length, 1); + await transaction.commit(); + }); + + const requests = spannerMock.getRequests(); + const metadataList = spannerMock.getMetadata(); + + const transactionalChannelIds: string[] = []; + for (let index = 0; index < requests.length; index++) { + const request = requests[index] as any; + if ( + request.sql === 'SELECT 1' || + request.mutations !== undefined || + request.transactionId !== undefined + ) { + const channelId = extractChannelId(metadataList[index]); + if (channelId) { + transactionalChannelIds.push(channelId); + } + } + } + + assert.strictEqual(transactionalChannelIds.length >= 3, true); + + const uniqueChannels = new Set(transactionalChannelIds); + assert.strictEqual( + uniqueChannels.size, + 1, + `All RPCs in the read/write transaction should use the same channel, but used: ${transactionalChannelIds.join(', ')}`, + ); + + await spanner.close(); + }); + + it('should distribute non-transactional RPCs across multiple channels', async () => { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + channelPool: { + type: 'static', + numChannels: 4, + }, + }); + + const database = spanner.instance('instance').database('database'); + spannerMock.resetRequests(); + + // Run multiple concurrent queries to exercise P2C channel distribution + const queryPromises: Array> = []; + for (let index = 0; index < 40; index++) { + queryPromises.push(database.run('SELECT 1')); + } + await Promise.all(queryPromises); + + const requests = spannerMock.getRequests(); + const metadataList = spannerMock.getMetadata(); + + const queryChannelIds: string[] = []; + for (let index = 0; index < requests.length; index++) { + const request = requests[index] as any; + if (request.sql === 'SELECT 1') { + const channelId = extractChannelId(metadataList[index]); + if (channelId) { + queryChannelIds.push(channelId); + } + } + } + + assert.strictEqual(queryChannelIds.length, 40); + + const uniqueChannels = new Set(queryChannelIds); + assert.strictEqual( + uniqueChannels.size, + 4, + `Expected non-transactional queries to be distributed across all 4 channels, but got ${uniqueChannels.size}: ${Array.from(uniqueChannels).join(', ')}`, + ); + + await spanner.close(); + }); + + it('should distribute read/write transactions across channels while keeping each transaction pinned', async () => { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + channelPool: { + type: 'static', + numChannels: 4, + }, + }); + + const database = spanner.instance('instance').database('database'); + spannerMock.resetRequests(); + + // Run concurrent read/write transactions to verify P2C distributes transactions + // across channels, while each individual transaction pins all its RPCs to one channel + const transactionCount = 8; + const transactionPromises: Array> = []; + + for ( + let transactionIndex = 0; + transactionIndex < transactionCount; + transactionIndex++ + ) { + const transactionTag = `tag-tx-${transactionIndex}`; + transactionPromises.push( + database.runTransactionAsync( + {requestOptions: {transactionTag}}, + async transaction => { + const [rows1] = await transaction.run('SELECT 1'); + assert.strictEqual(rows1.length, 1); + const [rows2] = await transaction.run('SELECT 1'); + assert.strictEqual(rows2.length, 1); + + // While active, transaction must be pinned to a channel + assert(transaction.affinity); + assert(transaction.affinity.pinnedEntry); + const pinnedChannelId = String( + transaction.affinity.pinnedEntry.id, + ); + + await transaction.commit(); + + // After commit, transaction affinity must be reset + assert.strictEqual(transaction.affinity.pinnedEntry, null); + + return pinnedChannelId; + }, + ), + ); + } + + const assignedChannels = await Promise.all(transactionPromises); + + // Verify transactions are distributed across multiple channels + const distinctAssignedChannels = new Set(assignedChannels); + assert.strictEqual( + distinctAssignedChannels.size > 1, + true, + `Expected concurrent transactions to be distributed across multiple channels, but only used: ${Array.from(distinctAssignedChannels).join(', ')}`, + ); + + // Also verify at the gRPC metadata level: all RPCs belonging to the same transaction used the same channel + const requests = spannerMock.getRequests(); + const metadataList = spannerMock.getMetadata(); + const transactionTagToChannels = new Map(); + + for (let index = 0; index < requests.length; index++) { + const request = requests[index] as any; + const tag = request.requestOptions?.transactionTag; + if ( + tag && + (request.sql === 'SELECT 1' || + request.mutations !== undefined || + request.transactionId !== undefined) + ) { + const channelId = extractChannelId(metadataList[index]); + if (channelId) { + if (!transactionTagToChannels.has(tag)) { + transactionTagToChannels.set(tag, []); + } + transactionTagToChannels.get(tag)!.push(channelId); + } + } + } + + assert.strictEqual(transactionTagToChannels.size, transactionCount); + for (const [tag, channelIds] of transactionTagToChannels.entries()) { + assert.strictEqual( + channelIds.length >= 3, + true, + `Transaction ${tag} should have at least 3 RPCs (2 queries and commit)`, + ); + const uniqueChannelsForTransaction = new Set(channelIds); + assert.strictEqual( + uniqueChannelsForTransaction.size, + 1, + `All RPCs for transaction ${tag} must use the same channel, but used: ${channelIds.join(', ')}`, + ); + } + + await spanner.close(); + }); + + it('should not close user-provided ChannelPool instance on spanner.close()', async () => { + const userPool = new StaticChannelPool( + `localhost:${port}`, + grpc.credentials.createInsecure(), + {}, + {numChannels: 2}, + ); + const closeSpy = sinon.spy(userPool, 'close'); + + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + channelPool: userPool, + }); + + const database = spanner.instance('instance').database('database'); + await database.run('SELECT 1'); + + await spanner.close(); + assert.strictEqual(closeSpy.called, false); + assert.strictEqual(userPool.size, 2); + + await userPool.close(); + assert.strictEqual(closeSpy.calledOnce, true); + }); + + it('should close internally created ChannelPool on spanner.close()', async () => { + const spanner = new Spanner({ + projectId: 'test-project', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + channelPool: {type: 'static', numChannels: 2}, + }); + + const database = spanner.instance('instance').database('database'); + await database.run('SELECT 1'); + + const internalPool = (spanner as any).channelPool_; + assert.ok(internalPool); + const closeSpy = sinon.spy(internalPool, 'close'); + + await spanner.close(); + assert.strictEqual(closeSpy.calledOnce, true); + }); + }); +}); + +function extractChannelId(metadata?: grpc.Metadata): string | null { + if (!metadata) { + return null; + } + const values = metadata.get('x-goog-spanner-request-id'); + if (values && values.length > 0) { + const parts = String(values[0]).split('.'); + if (parts.length >= 4) { + return parts[3]; + } + } + return null; +} diff --git a/handwritten/spanner/test/index.ts b/handwritten/spanner/test/index.ts index e3aaee22a3d8..371bfdd06b43 100644 --- a/handwritten/spanner/test/index.ts +++ b/handwritten/spanner/test/index.ts @@ -53,12 +53,18 @@ assert.strictEqual(CLOUD_RESOURCE_HEADER, 'google-cloud-resource-prefix'); const apiConfig = require('../src/spanner_grpc_config.json'); async function disableMetrics(sandbox: sinon.SinonSandbox) { + if (!('SPANNER_DISABLE_BUILTIN_METRICS' in process.env)) { + process.env.SPANNER_DISABLE_BUILTIN_METRICS = ''; + } sandbox.stub(process.env, 'SPANNER_DISABLE_BUILTIN_METRICS').value('true'); await MetricsTracerFactory.resetInstance(); MetricsTracerFactory.enabled = false; } async function enableMetrics(sandbox: sinon.SinonSandbox) { + if (!('SPANNER_DISABLE_BUILTIN_METRICS' in process.env)) { + process.env.SPANNER_DISABLE_BUILTIN_METRICS = ''; + } sandbox.stub(process.env, 'SPANNER_DISABLE_BUILTIN_METRICS').value('false'); await MetricsTracerFactory.resetInstance(); } @@ -184,6 +190,7 @@ describe('Spanner', () => { const OPTIONS = { projectId: 'project-id', + channelPool: 'legacy' as const, }; before(() => { diff --git a/handwritten/spanner/test/metrics/metrics-tracer-factory.ts b/handwritten/spanner/test/metrics/metrics-tracer-factory.ts index 5a474286becb..2eaf23aa65c7 100644 --- a/handwritten/spanner/test/metrics/metrics-tracer-factory.ts +++ b/handwritten/spanner/test/metrics/metrics-tracer-factory.ts @@ -181,6 +181,25 @@ describe('MetricsTracerFactory', () => { assert.strictEqual((factory as any)._currentOperationLastUpdatedMs.size, 0); }); + it('should retrieve and clear a MetricsTracer when channel id was rewritten by channel pool', () => { + const factory = MetricsTracerFactory.getInstance('project-id'); + const createdTracer = factory!.createMetricsTracer( + 'some-method', + 'method-name', + '1.1a2bc3d4.1.1.1.1', + ); + + assert.strictEqual((factory as any)._currentOperationTracers.size, 1); + + // Channel pool assigned channel 5: requestId becomes 1.1a2bc3d4.1.5.1.1 + const retrievedTracer = factory!.getCurrentTracer('1.1a2bc3d4.1.5.1.1'); + assert.strictEqual(retrievedTracer, createdTracer); + + factory!.clearCurrentTracer('1.1a2bc3d4.1.5.1.1'); + assert.strictEqual((factory as any)._currentOperationTracers.size, 0); + assert.strictEqual((factory as any)._currentOperationLastUpdatedMs.size, 0); + }); + it('should correctly set default attributes', () => { const factory = MetricsTracerFactory.getInstance('project-id'); const tracer = factory!.createMetricsTracer( diff --git a/handwritten/spanner/test/spanner.ts b/handwritten/spanner/test/spanner.ts index d8fb862f383a..1577d37e8e48 100644 --- a/handwritten/spanner/test/spanner.ts +++ b/handwritten/spanner/test/spanner.ts @@ -1381,7 +1381,14 @@ describe('Spanner with mock server', () => { database .close() .then(() => { - const gotStreamingCalls = xGoogReqIDInterceptor.getStreamingCalls(); + const gotStreamingCalls = xGoogReqIDInterceptor + .getStreamingCalls() + .map(call => { + const parts = call.reqId.split('.'); + assert(parseInt(parts[3], 10) >= 1); + parts[3] = '1'; + return {...call, reqId: parts.join('.')}; + }); const wantStreamingCalls = [ { method: '/google.spanner.v1.Spanner/ExecuteStreamingSql', @@ -7377,6 +7384,12 @@ describe('Spanner with mock server', () => { reqId: `1.${randIdForProcess}.1.1.9.1`, }, ]; + const normalizeChannelId = (call: {method: string; reqId: string}) => { + const parts = call.reqId.split('.'); + assert(parseInt(parts[3], 10) >= 1); + parts[3] = '1'; + return {...call, reqId: parts.join('.')}; + }; const gotUnaryCalls = xGoogReqIDInterceptor.getUnaryCalls(); assert.deepStrictEqual( gotUnaryCalls[0].method, @@ -7385,11 +7398,13 @@ describe('Spanner with mock server', () => { // It is non-deterministic to try to get the exact clientId used to invoke .BatchCreateSessions // given that these tests run as a collective and sessions are pooled. assert.deepStrictEqual( - gotUnaryCalls.slice(1), + gotUnaryCalls.slice(1).map(normalizeChannelId), wantUnaryCallsWithoutBatchCreateSessions, ); - const gotStreamingCalls = xGoogReqIDInterceptor.getStreamingCalls(); + const gotStreamingCalls = xGoogReqIDInterceptor + .getStreamingCalls() + .map(normalizeChannelId); const wantStreamingCalls = [ { method: '/google.spanner.v1.Spanner/ExecuteStreamingSql', diff --git a/handwritten/spanner/test/transaction-runner.ts b/handwritten/spanner/test/transaction-runner.ts index 9a2bd818b313..86e6a95f9e44 100644 --- a/handwritten/spanner/test/transaction-runner.ts +++ b/handwritten/spanner/test/transaction-runner.ts @@ -299,6 +299,18 @@ describe('TransactionRunner', () => { }); }); + it('should reset transaction affinity on error', async () => { + const fakeError = new Error('err') as grpc.ServiceError; + fakeError.code = grpc.status.UNKNOWN; + const resetStub = sandbox.stub(); + (fakeTransaction as any).affinity = {reset: resetStub}; + + runFn.rejects(fakeError); + + await assert.rejects(runner.run(), fakeError); + assert.strictEqual(resetStub.callCount, 1); + }); + it('should retry on ABORTED errors', async () => { const fakeReturnValue = 11; const fakeError = new Error('err') as grpc.ServiceError; @@ -330,7 +342,7 @@ describe('TransactionRunner', () => { runner .run() .then(() => { - done(new Error('missing expected DEADLINE_EXCEEDED error')); + return done(new Error('missing expected DEADLINE_EXCEEDED error')); }) .catch(err => { assert.strictEqual(err.code, grpc.status.DEADLINE_EXCEEDED); diff --git a/handwritten/spanner/test/transaction.ts b/handwritten/spanner/test/transaction.ts index 647f6c763d8e..0a43c382e482 100644 --- a/handwritten/spanner/test/transaction.ts +++ b/handwritten/spanner/test/transaction.ts @@ -21,6 +21,7 @@ import {EventEmitter} from 'events'; import {common as p} from 'protobufjs'; import * as proxyquire from 'proxyquire'; import * as sinon from 'sinon'; +import * as extend from 'extend'; import {codec} from '../src/codec'; import {protos} from '@google-cloud/spanner-api'; @@ -250,7 +251,11 @@ describe('Transaction', () => { assert.deepStrictEqual(arg.gaxOpts, { timeout: 1000, otherArgs: { - options: {unbind: true, affinityKey: txn._affinityKey}, + options: { + unbind: true, + affinityKey: txn._affinityKey, + affinity: txn.affinity, + }, }, }); // The caller supplied gax options must not be modified. @@ -2268,6 +2273,53 @@ describe('Transaction', () => { it('should inherit from Dml', () => { assert(transaction instanceof Dml); }); + + it('should preserve this context and deep clone configuration in wrapped request and requestStream', () => { + let capturedConfig: any; + const fakeSession = Object.assign({}, SESSION, { + metadata: undefined, + request: sinon.spy((config: any, callback: Function) => { + capturedConfig = config; + callback(); + }), + requestStream: sinon.spy((config: any) => { + capturedConfig = config; + return {} as any; + }), + }); + const txn = new Transaction(fakeSession); + assert.ok((txn as any)._affinity); + assert.strictEqual((txn as any)._affinityKey, undefined); + + const originalConfig = { + gaxOpts: { + otherArgs: { + options: { + custom: 'value', + }, + }, + }, + }; + const configToPass = extend(true, {}, originalConfig); + + txn.request(configToPass, () => {}); + assert.strictEqual(fakeSession.request.calledOnce, true); + // Original config object must not be mutated + assert.deepStrictEqual(configToPass, originalConfig); + // The spy received a deep-cloned config with affinity injected + assert.strictEqual( + capturedConfig.gaxOpts.otherArgs.options.affinity, + (txn as any)._affinity, + ); + + txn.requestStream(configToPass); + assert.strictEqual(fakeSession.requestStream.calledOnce, true); + assert.deepStrictEqual(configToPass, originalConfig); + assert.strictEqual( + capturedConfig.gaxOpts.otherArgs.options.affinity, + (txn as any)._affinity, + ); + }); }); describe('batchUpdate', () => { @@ -3407,8 +3459,10 @@ describe('Transaction', () => { it('should not return an error if the `id` is not set', done => { delete transaction.id; + const endStub = sandbox.stub(transaction, 'end'); transaction.rollback(err => { assert.deepStrictEqual(err, null); + assert.strictEqual(endStub.callCount, 1); done(); }); });