From 2914b1fd464fd5e48475d42e86e06821fb4caf78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Thu, 30 Jul 2026 18:02:16 +0200 Subject: [PATCH] chore(spanner): implement server connection pool for location-aware routing Add `ServerConnection` and `ConnectionCache` manage server connections and inflight request metrics for location-aware routing. --- src/spanner/src/client.rs | 13 + src/spanner/src/lib.rs | 1 + src/spanner/src/routing/connection_cache.rs | 346 +++++++++++++++++++ src/spanner/src/routing/mod.rs | 18 + src/spanner/src/routing/server_connection.rs | 269 ++++++++++++++ 5 files changed, 647 insertions(+) create mode 100644 src/spanner/src/routing/connection_cache.rs create mode 100644 src/spanner/src/routing/mod.rs create mode 100644 src/spanner/src/routing/server_connection.rs diff --git a/src/spanner/src/client.rs b/src/spanner/src/client.rs index b5b6f5247c..11ef202719 100644 --- a/src/spanner/src/client.rs +++ b/src/spanner/src/client.rs @@ -481,6 +481,19 @@ impl Channel { } } +#[cfg(test)] +impl Channel { + pub(crate) fn new_for_test(stub: T) -> Self + where + T: crate::generated::gapic_dataplane::stub::Spanner + 'static, + { + Self { + inner: GapicSpanner::from_stub(stub), + grpc_client: None, + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/spanner/src/lib.rs b/src/spanner/src/lib.rs index 1a3bccd620..96f99918c5 100644 --- a/src/spanner/src/lib.rs +++ b/src/spanner/src/lib.rs @@ -88,6 +88,7 @@ pub(crate) mod read_only_transaction; pub(crate) mod read_write_transaction; pub(crate) mod result_set; pub(crate) mod result_set_metadata; +pub(crate) mod routing; pub(crate) mod row; pub(crate) mod server_streaming; pub(crate) mod session_maintainer; diff --git a/src/spanner/src/routing/connection_cache.rs b/src/spanner/src/routing/connection_cache.rs new file mode 100644 index 0000000000..f254c9790d --- /dev/null +++ b/src/spanner/src/routing/connection_cache.rs @@ -0,0 +1,346 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +//! Thread-safe cache of Spanner server node connections for location-aware routing. + +// TODO(location-aware-routing): Remove allow(dead_code) once location_router.rs integrates ConnectionCache. +#![allow(dead_code)] + +use crate::client::Channel; +use crate::routing::server_connection::ServerConnection; +use gaxi::options::ClientConfig; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +/// Cache for server connections used in location-aware routing. +/// +/// Stores and manages [`ServerConnection`] instances such that repeated calls with the same address +/// return the same connection wrapper. +/// +/// # Thread Safety +/// +/// This cache is thread-safe and allows concurrent lookups across multiple operations. +#[derive(Debug)] +pub(crate) struct ConnectionCache { + default_connection: ServerConnection, + servers: RwLock>>>, +} + +impl ConnectionCache { + /// Creates a new `ConnectionCache` with the specified default fallback connection. + pub(crate) fn new(default_connection: ServerConnection) -> Self { + let default_cell = Arc::new(tokio::sync::OnceCell::from(default_connection.clone())); + let mut map = HashMap::new(); + map.insert(default_connection.address().to_string(), default_cell); + Self { + default_connection, + servers: RwLock::new(map), + } + } + + /// Returns a reference to the default fallback server connection. + pub(crate) fn default_connection(&self) -> &ServerConnection { + &self.default_connection + } + + /// Returns a cached connection for the given address without creating it if missing. + /// + /// This method is used by location-aware routing to avoid foreground connection creation on the + /// hot RPC request path. + pub(crate) fn get_if_present(&self, address: &str) -> Option { + let guard = self + .servers + .read() + .expect("connection cache read lock poisoned"); + guard.get(address).and_then(|cell| cell.get().cloned()) + } + + /// Returns a cached connection for the given address, creating and caching a new connection + /// asynchronously if needed. + pub(crate) async fn get( + &self, + address: &str, + config: &ClientConfig, + ) -> crate::ClientBuilderResult { + let cell = { + let guard = self + .servers + .read() + .expect("connection cache read lock poisoned"); + if let Some(cell) = guard.get(address) { + if let Some(connection) = cell.get() { + return Ok(connection.clone()); + } + cell.clone() + } else { + drop(guard); + let mut guard = self + .servers + .write() + .expect("connection cache write lock poisoned"); + guard + .entry(address.to_string()) + .or_insert_with(|| Arc::new(tokio::sync::OnceCell::new())) + .clone() + } + }; + + cell.get_or_try_init(|| async { + let mut ep_config = config.clone(); + let addr = address.to_string(); + ep_config.endpoint = Some(addr.clone()); + let channel = Channel::create(&ep_config).await?; + Ok(ServerConnection::new(addr, channel)) + }) + .await + .cloned() + } + + /// Evicts a server connection from the cache. + /// + /// If `address` matches the default connection's address, this method does nothing and returns + /// `false`. Otherwise, returns `true` if a connection was removed from the cache. + pub(crate) fn evict(&self, address: &str) -> bool { + if self.default_connection.address() == address { + return false; + } + let mut guard = self + .servers + .write() + .expect("connection cache write lock poisoned"); + guard.remove(address).is_some() + } + + /// Returns the number of cached server connections (including the default connection). + pub(crate) fn len(&self) -> usize { + let guard = self + .servers + .read() + .expect("connection cache read lock poisoned"); + guard.values().filter(|cell| cell.get().is_some()).count() + } + + /// Returns whether the cache is empty. + pub(crate) fn is_empty(&self) -> bool { + let guard = self + .servers + .read() + .expect("connection cache read lock poisoned"); + !guard.values().any(|cell| cell.get().is_some()) + } + + /// Clears all cached server connections while preserving the default fallback connection. + pub(crate) fn clear(&self) { + let mut guard = self + .servers + .write() + .expect("connection cache write lock poisoned"); + let default_address = self.default_connection.address(); + guard.retain(|k, _| k == default_address); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Barrier; + use std::thread; + + #[derive(Debug)] + struct DummyStub; + impl crate::generated::gapic_dataplane::stub::Spanner for DummyStub {} + + fn create_test_connection(address: &str) -> ServerConnection { + let channel = Channel::new_for_test(DummyStub); + ServerConnection::new(address.to_string(), channel) + } + + #[test] + fn test_connection_cache_default_connection_and_get_if_present() { + let default_conn = create_test_connection("spanner.googleapis.com:443"); + let cache = ConnectionCache::new(default_conn.clone()); + + assert_eq!(cache.len(), 1); + assert!(!cache.is_empty()); + assert_eq!( + cache.default_connection().address(), + "spanner.googleapis.com:443" + ); + + let cached_default = cache + .get_if_present("spanner.googleapis.com:443") + .expect("default connection should be in cache"); + assert_eq!(cached_default.address(), "spanner.googleapis.com:443"); + + assert!(cache.get_if_present("10.0.0.1:15000").is_none()); + } + + #[test] + fn test_connection_cache_eviction_and_protection_of_default() { + let default_conn = create_test_connection("spanner.googleapis.com:443"); + let cache = ConnectionCache::new(default_conn); + + // Manually insert a tablet connection into the cache for testing eviction. + let tablet_conn = create_test_connection("10.0.0.1:15000"); + { + let mut guard = cache.servers.write().expect("write lock poisoned"); + let cell = Arc::new(tokio::sync::OnceCell::new()); + let _ = cell.set(tablet_conn.clone()); + guard.insert(tablet_conn.address().to_string(), cell); + } + assert_eq!(cache.len(), 2); + + // Evicting the default connection should be ignored. + assert!(!cache.evict("spanner.googleapis.com:443")); + assert_eq!(cache.len(), 2); + + // Evicting a normal tablet connection should succeed. + assert!(cache.evict("10.0.0.1:15000")); + assert_eq!(cache.len(), 1); + assert!(cache.get_if_present("10.0.0.1:15000").is_none()); + } + + #[test] + fn test_connection_cache_clear_preserves_default() { + let default_conn = create_test_connection("spanner.googleapis.com:443"); + let cache = ConnectionCache::new(default_conn); + + { + let mut guard = cache.servers.write().expect("write lock poisoned"); + let cell1 = Arc::new(tokio::sync::OnceCell::new()); + let _ = cell1.set(create_test_connection("10.0.0.1:15000")); + guard.insert("10.0.0.1:15000".to_string(), cell1); + + let cell2 = Arc::new(tokio::sync::OnceCell::new()); + let _ = cell2.set(create_test_connection("10.0.0.2:15000")); + guard.insert("10.0.0.2:15000".to_string(), cell2); + } + assert_eq!(cache.len(), 3); + + cache.clear(); + assert_eq!(cache.len(), 1); + assert_eq!( + cache.default_connection().address(), + "spanner.googleapis.com:443" + ); + assert!(cache.get_if_present("spanner.googleapis.com:443").is_some()); + assert!(cache.get_if_present("10.0.0.1:15000").is_none()); + assert!(cache.get_if_present("10.0.0.2:15000").is_none()); + } + + #[test] + fn test_connection_cache_concurrent_access() { + let default_conn = create_test_connection("spanner.googleapis.com:443"); + let cache = ConnectionCache::new(default_conn); + let worker_count = 10; + let iterations = 100; + let barrier = Barrier::new(worker_count); + + thread::scope(|scope| { + for _ in 0..worker_count { + scope.spawn(|| { + barrier.wait(); + for _ in 0..iterations { + assert!(cache.get_if_present("spanner.googleapis.com:443").is_some()); + assert!(!cache.evict("spanner.googleapis.com:443")); + } + }); + } + }); + + assert_eq!(cache.len(), 1); + } + + #[tokio::test] + async fn test_connection_cache_get_cached() { + let default_conn = create_test_connection("spanner.googleapis.com:443"); + let cache = ConnectionCache::new(default_conn); + let config = ClientConfig::default(); + + let ep = cache + .get("spanner.googleapis.com:443", &config) + .await + .expect("cached default connection"); + assert_eq!(ep.address(), "spanner.googleapis.com:443"); + assert_eq!(cache.len(), 1); + } + + #[tokio::test] + async fn test_connection_cache_concurrent_get_stampede_prevention() { + let default_conn = create_test_connection("spanner.googleapis.com:443"); + let cache = Arc::new(ConnectionCache::new(default_conn)); + let config = ClientConfig::default(); + + let mut handles = Vec::new(); + for _ in 0..10 { + let cache_clone = Arc::clone(&cache); + let config_clone = config.clone(); + handles.push(tokio::spawn(async move { + cache_clone + .get("http://10.0.0.1:15000", &config_clone) + .await + .expect("should obtain connection") + })); + } + + let mut results = Vec::new(); + for handle in handles { + results.push(handle.await.expect("task should complete")); + } + + for conn in &results { + assert_eq!(conn.address(), "http://10.0.0.1:15000"); + } + assert_eq!(cache.len(), 2); + } + + #[test] + fn test_connection_cache_uninitialized_cell_not_counted_in_len() { + let default_conn = create_test_connection("spanner.googleapis.com:443"); + let cache = ConnectionCache::new(default_conn); + assert_eq!(cache.len(), 1); + assert!(!cache.is_empty()); + + { + let mut guard = cache.servers.write().expect("write lock poisoned"); + guard.insert( + "10.0.0.1:15000".to_string(), + Arc::new(tokio::sync::OnceCell::new()), + ); + } + + assert_eq!(cache.len(), 1); + assert!(!cache.is_empty()); + assert!(cache.get_if_present("10.0.0.1:15000").is_none()); + } + + #[test] + fn test_connection_cache_evict_uninitialized_cell() { + let default_conn = create_test_connection("spanner.googleapis.com:443"); + let cache = ConnectionCache::new(default_conn); + + { + let mut guard = cache.servers.write().expect("write lock poisoned"); + guard.insert( + "10.0.0.1:15000".to_string(), + Arc::new(tokio::sync::OnceCell::new()), + ); + } + + assert_eq!(cache.len(), 1); + assert!(cache.evict("10.0.0.1:15000")); + assert_eq!(cache.len(), 1); + assert!(cache.get_if_present("10.0.0.1:15000").is_none()); + } +} diff --git a/src/spanner/src/routing/mod.rs b/src/spanner/src/routing/mod.rs new file mode 100644 index 0000000000..dcf6183852 --- /dev/null +++ b/src/spanner/src/routing/mod.rs @@ -0,0 +1,18 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +//! Location-aware routing and endpoint connection pooling for Spanner clients. + +pub(crate) mod connection_cache; +pub(crate) mod server_connection; diff --git a/src/spanner/src/routing/server_connection.rs b/src/spanner/src/routing/server_connection.rs new file mode 100644 index 0000000000..e7a8ffdff2 --- /dev/null +++ b/src/spanner/src/routing/server_connection.rs @@ -0,0 +1,269 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +//! Server connection wrapper and inflight request tracking for location-aware routing. + +// TODO(location-aware-routing): Remove allow(dead_code) once location_router.rs integrates ServerConnection. +#![allow(dead_code)] + +use crate::client::Channel; +use std::sync::Arc; +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; + +const STATE_READY: u8 = 0; +const STATE_TRANSIENT_FAILURE: u8 = 1; +const STATE_UNHEALTHY: u8 = 2; + +/// A Spanner server connection wrapper for location-aware routing. +/// +/// Wraps a gRPC [`Channel`] connected to a specific Spanner server node address and tracks +/// connectivity health state and active inflight requests. +/// +/// # Thread Safety +/// +/// All health transitions and request counter modifications are lock-free and thread-safe. +#[derive(Clone, Debug)] +pub(crate) struct ServerConnection { + inner: Arc, +} + +#[derive(Debug)] +struct Inner { + address: String, + channel: Channel, + state: AtomicU8, + active_requests: AtomicUsize, +} + +/// RAII guard that decrements the active request count of a [`ServerConnection`] when dropped. +#[must_use = "if unused the request count will decrement immediately"] +pub(crate) struct ActiveRequestGuard { + inner: Arc, +} + +impl Drop for ActiveRequestGuard { + fn drop(&mut self) { + let _ = + self.inner + .active_requests + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |val| { + Some(val.saturating_sub(1)) + }); + } +} + +impl ServerConnection { + /// Creates a new `ServerConnection` wrapping the given address and channel in the `READY` state. + pub(crate) fn new(address: String, channel: Channel) -> Self { + Self { + inner: Arc::new(Inner { + address, + channel, + state: AtomicU8::new(STATE_READY), + active_requests: AtomicUsize::new(0), + }), + } + } + + /// Returns the network address of this server in `"host:port"` format. + pub(crate) fn address(&self) -> &str { + &self.inner.address + } + + /// Returns a reference to the wrapped gRPC [`Channel`]. + pub(crate) fn channel(&self) -> &Channel { + &self.inner.channel + } + + /// Returns whether this connection is in the `READY` state and eligible for location-aware routing. + pub(crate) fn is_healthy(&self) -> bool { + self.inner.state.load(Ordering::Acquire) == STATE_READY + } + + /// Returns whether this connection is in the `TRANSIENT_FAILURE` state. + pub(crate) fn is_transient_failure(&self) -> bool { + self.inner.state.load(Ordering::Acquire) == STATE_TRANSIENT_FAILURE + } + + /// Marks this connection as `READY`. + pub(crate) fn set_ready(&self) { + self.inner.state.store(STATE_READY, Ordering::Release); + } + + /// Marks this connection as in `TRANSIENT_FAILURE`. + pub(crate) fn set_transient_failure(&self) { + self.inner + .state + .store(STATE_TRANSIENT_FAILURE, Ordering::Release); + } + + /// Marks this connection as `UNHEALTHY`. + pub(crate) fn set_unhealthy(&self) { + self.inner.state.store(STATE_UNHEALTHY, Ordering::Release); + } + + /// Increments the active inflight request count for this connection. + pub(crate) fn increment_active_requests(&self) { + self.inner.active_requests.fetch_add(1, Ordering::Relaxed); + } + + /// Decrements the active inflight request count for this connection without underflow. + pub(crate) fn decrement_active_requests(&self) { + let _ = + self.inner + .active_requests + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |val| { + Some(val.saturating_sub(1)) + }); + } + + /// Returns the current number of active inflight requests on this connection. + pub(crate) fn active_request_count(&self) -> usize { + self.inner.active_requests.load(Ordering::Relaxed) + } + + /// Increments the active request count and returns an RAII guard that automatically decrements + /// the count when dropped. + pub(crate) fn acquire_request_guard(&self) -> ActiveRequestGuard { + self.increment_active_requests(); + ActiveRequestGuard { + inner: Arc::clone(&self.inner), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Barrier; + use std::thread; + + #[derive(Debug)] + struct DummyStub; + impl crate::generated::gapic_dataplane::stub::Spanner for DummyStub {} + + fn create_test_connection(address: &str) -> ServerConnection { + let channel = Channel::new_for_test(DummyStub); + ServerConnection::new(address.to_string(), channel) + } + + #[test] + fn test_server_connection_health_state_transitions() { + let conn = create_test_connection("10.0.0.1:15000"); + assert_eq!(conn.address(), "10.0.0.1:15000"); + assert!(conn.is_healthy()); + assert!(!conn.is_transient_failure()); + + conn.set_transient_failure(); + assert!(!conn.is_healthy()); + assert!(conn.is_transient_failure()); + + conn.set_unhealthy(); + assert!(!conn.is_healthy()); + assert!(!conn.is_transient_failure()); + + conn.set_ready(); + assert!(conn.is_healthy()); + assert!(!conn.is_transient_failure()); + } + + #[test] + fn test_server_connection_active_requests_and_underflow_protection() { + let conn = create_test_connection("10.0.0.1:15000"); + assert_eq!(conn.active_request_count(), 0); + + conn.increment_active_requests(); + conn.increment_active_requests(); + conn.increment_active_requests(); + assert_eq!(conn.active_request_count(), 3); + + conn.decrement_active_requests(); + assert_eq!(conn.active_request_count(), 2); + + conn.decrement_active_requests(); + conn.decrement_active_requests(); + assert_eq!(conn.active_request_count(), 0); + + // Verify saturating subtraction prevents wrapping underflow. + conn.decrement_active_requests(); + conn.decrement_active_requests(); + assert_eq!(conn.active_request_count(), 0); + } + + #[test] + fn test_server_connection_channel_accessor() { + let conn = create_test_connection("10.0.0.1:15000"); + let channel = conn.channel(); + assert!(format!("{:?}", channel).contains("Channel")); + } + + #[test] + fn test_server_connection_acquire_request_guard() { + let conn = create_test_connection("10.0.0.1:15000"); + assert_eq!(conn.active_request_count(), 0); + + { + let _guard1 = conn.acquire_request_guard(); + assert_eq!(conn.active_request_count(), 1); + + { + let _guard2 = conn.acquire_request_guard(); + assert_eq!(conn.active_request_count(), 2); + } + + assert_eq!(conn.active_request_count(), 1); + } + + assert_eq!(conn.active_request_count(), 0); + } + + #[test] + fn test_server_connection_concurrent_guards() { + let conn = create_test_connection("10.0.0.1:15000"); + let worker_count = 10; + let iterations = 100; + let barrier = Barrier::new(worker_count); + + thread::scope(|scope| { + for _ in 0..worker_count { + scope.spawn(|| { + barrier.wait(); + for _ in 0..iterations { + let _guard = conn.acquire_request_guard(); + assert!(conn.active_request_count() > 0); + } + }); + } + }); + + assert_eq!(conn.active_request_count(), 0); + } + + #[tokio::test] + async fn test_server_connection_guard_static_lifetime() { + let conn = create_test_connection("10.0.0.1:15000"); + let guard = conn.acquire_request_guard(); + assert_eq!(conn.active_request_count(), 1); + + let conn_clone = conn.clone(); + tokio::spawn(async move { + let _moved_guard = guard; + assert_eq!(conn_clone.active_request_count(), 1); + }) + .await + .expect("spawned task should complete"); + + assert_eq!(conn.active_request_count(), 0); + } +}