From 51dff38519cca2b66387e5762ce20da93f5904ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Wed, 16 Sep 2026 13:10:18 +0200 Subject: [PATCH] feat(spanner): add support for static and dynamic channel pooling Adds support for configuring gRPC channel pooling when creating a Spanner client. Customers can now configure the channel pool on `Spanner::builder()` using `.with_channel_pool(...)`: - Static channel pool (`StaticChannelPoolConfig`): Configures a fixed number of gRPC channels (defaults to 4 channels, or 1 channel when targeting the Spanner emulator). - Dynamic channel pool (`DynamicChannelPoolConfig`): Automatically scales the number of channels up or down based on RPC concurrency and load. --- src/spanner/README.md | 25 + .../src/batch_read_only_transaction.rs | 108 +- src/spanner/src/batch_write_transaction.rs | 10 +- src/spanner/src/channel_pool/affinity.rs | 151 +- src/spanner/src/channel_pool/config.rs | 364 +++- src/spanner/src/channel_pool/entry.rs | 82 +- .../src/channel_pool/integration_tests.rs | 1706 ++++++++++++++++ src/spanner/src/channel_pool/mod.rs | 24 +- src/spanner/src/channel_pool/pool.rs | 95 +- src/spanner/src/channel_pool/scaler.rs | 183 +- src/spanner/src/client.rs | 605 ++++-- src/spanner/src/database_client.rs | 336 ++-- src/spanner/src/lib.rs | 3 +- .../src/partitioned_dml_transaction.rs | 10 +- src/spanner/src/read_only_transaction.rs | 57 +- src/spanner/src/read_write_transaction.rs | 692 ++++++- src/spanner/src/request_id.rs | 6 +- src/spanner/src/result_set.rs | 19 +- src/spanner/src/routing/mock_tests.rs | 147 +- src/spanner/src/server_streaming/builder.rs | 3 +- src/spanner/src/server_streaming/stream.rs | 15 +- src/spanner/src/session_maintainer.rs | 8 + src/spanner/src/transaction_runner.rs | 183 ++ src/spanner/src/write_only_transaction.rs | 15 +- src/spanner/tests/channel_pool_contention.rs | 1732 +++++++++++++++++ 25 files changed, 5934 insertions(+), 645 deletions(-) create mode 100644 src/spanner/src/channel_pool/integration_tests.rs create mode 100644 src/spanner/tests/channel_pool_contention.rs diff --git a/src/spanner/README.md b/src/spanner/README.md index 828db46eba..3ef3385306 100644 --- a/src/spanner/README.md +++ b/src/spanner/README.md @@ -69,6 +69,31 @@ export SPANNER_EMULATOR_HOST=localhost:9010 The client builder automatically detects this variable, connects to the emulator endpoint, and configures anonymous credentials. +### Configuring the Channel Pool + +By default, the client uses a static pool of 4 gRPC channels. You can configure +the channel pool using `with_channel_pool`: + +```rust +use google_cloud_spanner::client::{Spanner, SpannerBuilderExt}; +use google_cloud_spanner::channel_pool::{DynamicChannelPoolConfig, StaticChannelPoolConfig}; + +# async fn sample() -> Result<(), google_cloud_spanner::Error> { +// Custom static pool: +let spanner = Spanner::builder() + .with_channel_pool(StaticChannelPoolConfig::new(8)) + .build() + .await?; + +// Or dynamic load-based channel pool: +let spanner = Spanner::builder() + .with_channel_pool(DynamicChannelPoolConfig::new()) + .build() + .await?; +# Ok(()) +# } +``` + ## Session Management and Client Lifecycle The Spanner Rust client manages a long-lived multiplexed session under the hood. diff --git a/src/spanner/src/batch_read_only_transaction.rs b/src/spanner/src/batch_read_only_transaction.rs index 43028e4a18..b0fad118ff 100644 --- a/src/spanner/src/batch_read_only_transaction.rs +++ b/src/spanner/src/batch_read_only_transaction.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::channel_pool::ChannelTarget; use crate::database_client::DatabaseClient; use crate::model::{ExecuteSqlRequest, PartitionOptions, ReadRequest}; use crate::precommit::PrecommitTokenTracker; @@ -161,15 +162,12 @@ impl BatchReadOnlyTransaction { .set_transaction(selector.clone()) .set_partition_options(options); + let target = ChannelTarget::from(self.inner.context.affinity()); let response = self .inner .context .client - .partition_query( - request, - crate::RequestOptions::default(), - self.inner.context.channel_hint, - ) + .partition_query(request, crate::RequestOptions::default(), target) .await?; Ok(response @@ -224,15 +222,12 @@ impl BatchReadOnlyTransaction { .set_transaction(selector.clone()) .set_partition_options(options); + let target = ChannelTarget::from(self.inner.context.affinity()); let response = self .inner .context .client - .partition_read( - request, - crate::RequestOptions::default(), - self.inner.context.channel_hint, - ) + .partition_read(request, crate::RequestOptions::default(), target) .await?; Ok(response @@ -405,12 +400,10 @@ impl Partition { req: &ExecuteSqlRequest, gax_options: GaxRequestOptions, ) -> crate::Result { - let channel_hint = client.next_channel_hint(); - let gax_options = client.attach_request_id(gax_options, channel_hint); let (stream, attempt_start_time) = Self::execute_partition_stream(client, "ExecuteStreamingSql", || { client - .execute_streaming_sql(req.clone(), gax_options.clone(), channel_hint) + .execute_streaming_sql(req.clone(), gax_options.clone(), ChannelTarget::Any) .send() }) .await?; @@ -428,7 +421,6 @@ impl Partition { session_name: req.session.clone(), transaction_tag: None, operation: StreamOperation::Query(req.clone()), - channel_hint, gax_options, method_name: "ExecuteStreamingSql", attempt_start_time: Some(attempt_start_time), @@ -443,12 +435,10 @@ impl Partition { req: &ReadRequest, gax_options: GaxRequestOptions, ) -> crate::Result { - let channel_hint = client.next_channel_hint(); - let gax_options = client.attach_request_id(gax_options, channel_hint); let (stream, attempt_start_time) = Self::execute_partition_stream(client, "StreamingRead", || { client - .streaming_read(req.clone(), gax_options.clone(), channel_hint) + .streaming_read(req.clone(), gax_options.clone(), ChannelTarget::Any) .send() }) .await?; @@ -466,7 +456,6 @@ impl Partition { session_name: req.session.clone(), transaction_tag: None, operation: StreamOperation::Read(req.clone()), - channel_hint, gax_options, method_name: "StreamingRead", attempt_start_time: Some(attempt_start_time), @@ -493,7 +482,9 @@ pub(crate) mod tests { use crate::read_only_transaction::tests::{create_session_mock, setup_db_client}; use crate::statement::Statement; use crate::transaction::TimestampBound; - use gaxi::grpc::tonic::Response; + use gaxi::grpc::tonic::{Response, Status}; + use google_cloud_gax::exponential_backoff::ExponentialBackoff; + use google_cloud_gax::retry_policy::NeverRetry; use google_cloud_test_macros::tokio_test_no_panics; use prost_types::Timestamp; use spanner_grpc_mock::google::spanner::v1::{ @@ -502,6 +493,7 @@ pub(crate) mod tests { }; use static_assertions::assert_impl_all; use std::fmt::Debug; + use std::time::Duration; #[test] fn auto_traits() { @@ -512,8 +504,6 @@ pub(crate) mod tests { #[test] fn serialize_partition_skips_gax_options() -> anyhow::Result<()> { - use std::time::Duration; - let req = crate::model::ExecuteSqlRequest::new() .set_sql("SELECT 1") .set_partition_token(b"token".to_vec()); @@ -937,4 +927,80 @@ pub(crate) mod tests { Ok(()) } + + #[tokio_test_no_panics] + async fn execute_query_with_retry_and_backoff_policy() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + + mock.expect_execute_streaming_sql().once().returning(|_| { + Ok(Response::from(crate::result_set::tests::adapt([Ok( + crate::read_only_transaction::tests::setup_select1(), + )]))) + }); + + let (db_client, _server) = setup_db_client(mock).await; + + let req = crate::model::ExecuteSqlRequest::new() + .set_session("projects/p/instances/i/databases/d/sessions/123") + .set_transaction(crate::model::TransactionSelector { + selector: Some(crate::model::transaction_selector::Selector::Id( + b"tx_id_1".to_vec().into(), + )), + ..Default::default() + }) + .set_sql("SELECT * FROM Users") + .set_partition_token(b"partition_token_123".to_vec()); + + let partition = Partition { + inner: PartitionedOperation::Query(req), + gax_options: GaxRequestOptions::default(), + }; + + let _result_set = partition + .with_retry_policy(NeverRetry) + .with_backoff_policy(ExponentialBackoff::default()) + .execute(&db_client) + .await?; + + Ok(()) + } + + #[tokio_test_no_panics] + async fn execute_query_error_records_telemetry() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + + mock.expect_execute_streaming_sql() + .once() + .returning(|_| Err(Status::internal("rpc failed"))); + + let (db_client, _server) = setup_db_client(mock).await; + + let req = crate::model::ExecuteSqlRequest::new() + .set_session("projects/p/instances/i/databases/d/sessions/123") + .set_transaction(crate::model::TransactionSelector { + selector: Some(crate::model::transaction_selector::Selector::Id( + b"tx_id_1".to_vec().into(), + )), + ..Default::default() + }) + .set_sql("SELECT * FROM Users") + .set_partition_token(b"partition_token_123".to_vec()); + + let partition = Partition { + inner: PartitionedOperation::Query(req), + gax_options: GaxRequestOptions::default(), + }; + + let result = partition + .with_retry_policy(NeverRetry) + .execute(&db_client) + .await; + + assert!( + result.is_err(), + "partition execution must fail on rpc error" + ); + + Ok(()) + } } diff --git a/src/spanner/src/batch_write_transaction.rs b/src/spanner/src/batch_write_transaction.rs index 3ad9a95264..631bc0b7c3 100644 --- a/src/spanner/src/batch_write_transaction.rs +++ b/src/spanner/src/batch_write_transaction.rs @@ -14,6 +14,7 @@ use crate::Error; use crate::Result; +use crate::channel_pool::ChannelTarget; use crate::client::DatabaseClient; use crate::error::internal_error; use crate::google::spanner::v1::BatchWriteResponse as ProtoBatchWriteResponse; @@ -209,12 +210,10 @@ impl BatchWriteTransactionBuilder { /// ``` pub fn build(self) -> BatchWriteTransaction { let session_name = self.client.session_name(); - let channel_hint = self.client.next_channel_hint(); let gax_options = apply_defaults(self.gax_options); BatchWriteTransaction { session_name, client: self.client, - channel_hint, transaction_tag: self.transaction_tag, priority: self.priority, exclude_txn_from_change_streams: self.exclude_txn_from_change_streams, @@ -230,7 +229,6 @@ impl BatchWriteTransactionBuilder { pub struct BatchWriteTransaction { session_name: String, client: DatabaseClient, - channel_hint: usize, transaction_tag: Option, priority: Priority, exclude_txn_from_change_streams: bool, @@ -291,7 +289,6 @@ impl BatchWriteTransaction { Ok(BatchWriteResponseStream { client: self.client, session_name: self.session_name, - channel_hint: self.channel_hint, transaction_tag: self.transaction_tag, priority: self.priority, exclude_txn_from_change_streams: self.exclude_txn_from_change_streams, @@ -317,7 +314,6 @@ impl BatchWriteTransaction { pub struct BatchWriteResponseStream { client: DatabaseClient, session_name: String, - channel_hint: usize, transaction_tag: Option, priority: Priority, exclude_txn_from_change_streams: bool, @@ -436,7 +432,7 @@ impl BatchWriteResponseStream { let stream_result = self .client - .batch_write(request, self.gax_options.clone(), self.channel_hint) + .batch_write(request, self.gax_options.clone(), ChannelTarget::Any) .send() .await; @@ -540,8 +536,6 @@ impl BatchWriteResponseStream { match self.check_retry(error) { Ok(()) => { self.retry_count += 1; - // Rotate channel hint only when a retry is confirmed to distribute load across healthy connections. - self.channel_hint = self.client.next_channel_hint(); if let Some(policy) = self.gax_options.backoff_policy() { let state = RetryState::new(true).set_attempt_count(self.retry_count as u32); let delay = policy.on_failure(&state); diff --git a/src/spanner/src/channel_pool/affinity.rs b/src/spanner/src/channel_pool/affinity.rs index ac7569b31e..c37591f4d3 100644 --- a/src/spanner/src/channel_pool/affinity.rs +++ b/src/spanner/src/channel_pool/affinity.rs @@ -55,16 +55,6 @@ impl TransactionAffinity { } } - /// Returns the provided affinity handle, or creates a new default `ReadOnly` affinity if `None`. - pub(crate) fn default_read_only(existing: Option>) -> Arc { - existing.unwrap_or_else(|| Arc::new(Self::new_read_only())) - } - - /// Returns the provided affinity handle, or creates a new default `ReadWrite` affinity if `None`. - pub(crate) fn default_read_write(existing: Option>) -> Arc { - existing.unwrap_or_else(|| Arc::new(Self::new_read_write())) - } - /// Returns `true` if this handle requires hard stickiness (Read/Write transactions). pub(crate) fn is_read_write(&self) -> bool { self.kind == AffinityKind::ReadWrite @@ -102,6 +92,16 @@ impl TransactionAffinity { } *slot = Some(lease.rw_affinity_guard()); } + + /// Releases the active Read/Write transaction guard on the channel entry, + /// allowing draining channels to close once the transaction completes. + pub(crate) fn release_rw_guard(&self) { + let mut slot = self + .rw_guard + .lock() + .expect("affinity rw_guard lock poisoned"); + *slot = None; + } } /// Stickiness kind for transaction channel affinity. @@ -116,6 +116,50 @@ pub(crate) enum AffinityKind { ReadOnly, } +/// Routing target for channel selection: either a round-robin hint or a transaction affinity handle. +#[derive(Clone, Copy, Debug, Default)] +pub(crate) enum ChannelTarget<'a> { + #[default] + Any, + Affinity(&'a TransactionAffinity), +} + +impl<'a> From<&'a TransactionAffinity> for ChannelTarget<'a> { + fn from(affinity: &'a TransactionAffinity) -> Self { + Self::Affinity(affinity) + } +} + +impl<'a> From<&'a Arc> for ChannelTarget<'a> { + fn from(affinity: &'a Arc) -> Self { + Self::Affinity(affinity.as_ref()) + } +} + +impl<'a> From> for ChannelTarget<'a> { + fn from(affinity: Option<&'a TransactionAffinity>) -> Self { + match affinity { + Some(affinity) => Self::Affinity(affinity), + None => Self::Any, + } + } +} + +impl<'a> From<&'a Option>> for ChannelTarget<'a> { + fn from(affinity: &'a Option>) -> Self { + match affinity { + Some(affinity) => Self::Affinity(affinity.as_ref()), + None => Self::Any, + } + } +} + +impl From<()> for ChannelTarget<'_> { + fn from(_: ()) -> Self { + Self::Any + } +} + #[cfg(test)] impl TransactionAffinity { pub(crate) fn set_entry_id(&self, entry_id: u64) { @@ -155,6 +199,8 @@ mod tests { use crate::client::Channel; use crate::generated::gapic_dataplane::stub::Spanner as SpannerStub; use std::fmt::Debug; + use std::ptr; + use std::sync::Arc; use std::time::Duration; #[derive(Debug)] @@ -255,37 +301,6 @@ mod tests { ); } - #[test] - fn transaction_affinity_defaults() { - let default_read_only = TransactionAffinity::default_read_only(None); - assert!( - default_read_only.is_read_only(), - "default_read_only(None) must return ReadOnly affinity" - ); - - let custom_read_only = Arc::new(TransactionAffinity::new_read_only()); - let passed_read_only = - TransactionAffinity::default_read_only(Some(Arc::clone(&custom_read_only))); - assert!( - Arc::ptr_eq(&custom_read_only, &passed_read_only), - "default_read_only(Some(handle)) must return existing handle without recreating" - ); - - let default_read_write = TransactionAffinity::default_read_write(None); - assert!( - default_read_write.is_read_write(), - "default_read_write(None) must return ReadWrite affinity" - ); - - let custom_read_write = Arc::new(TransactionAffinity::new_read_write()); - let passed_read_write = - TransactionAffinity::default_read_write(Some(Arc::clone(&custom_read_write))); - assert!( - Arc::ptr_eq(&custom_read_write, &passed_read_write), - "default_read_write(Some(handle)) must return existing handle without recreating" - ); - } - #[test] fn attach_rw_guard_repinning_replaces_old_guard() { let channel1 = Channel::new_for_test(DummyStub); @@ -391,4 +406,58 @@ mod tests { assert_eq!(entry1.active_rw_count(), 0, "entry1 count must drop to 0"); assert_eq!(entry2.active_rw_count(), 1, "entry2 count must be 1"); } + + #[test] + fn channel_target_from_conversions() { + let affinity = TransactionAffinity::new_read_write(); + let target_from_ref = ChannelTarget::from(&affinity); + assert!( + matches!(target_from_ref, ChannelTarget::Affinity(a) if ptr::eq(a, &affinity)), + "target_from_ref affinity must match" + ); + + let arc_affinity = Arc::new(TransactionAffinity::new_read_only()); + let target_from_arc = ChannelTarget::from(&arc_affinity); + assert!( + matches!(target_from_arc, ChannelTarget::Affinity(a) if ptr::eq(a, arc_affinity.as_ref())), + "target_from_arc affinity must match" + ); + + let target_from_some = ChannelTarget::from(Some(&affinity)); + assert!( + matches!(target_from_some, ChannelTarget::Affinity(a) if ptr::eq(a, &affinity)), + "target_from_some affinity must match" + ); + + let target_from_none = ChannelTarget::from(None); + assert!( + matches!(target_from_none, ChannelTarget::Any), + "target_from_none must be ChannelTarget::Any" + ); + + let opt_arc: Option> = Some(Arc::clone(&arc_affinity)); + let target_from_opt_arc = ChannelTarget::from(&opt_arc); + assert!( + matches!(target_from_opt_arc, ChannelTarget::Affinity(a) if ptr::eq(a, arc_affinity.as_ref())), + "target_from_opt_arc must be ChannelTarget::Affinity" + ); + + let opt_arc_none: Option> = None; + let target_from_opt_arc_none = ChannelTarget::from(&opt_arc_none); + assert!( + matches!(target_from_opt_arc_none, ChannelTarget::Any), + "target_from_opt_arc_none must be ChannelTarget::Any" + ); + + let target_from_unit = ChannelTarget::from(()); + assert!( + matches!(target_from_unit, ChannelTarget::Any), + "target_from_unit must be ChannelTarget::Any" + ); + + assert!( + matches!(ChannelTarget::default(), ChannelTarget::Any), + "ChannelTarget default must be ChannelTarget::Any" + ); + } } diff --git a/src/spanner/src/channel_pool/config.rs b/src/spanner/src/channel_pool/config.rs index bf2346bde0..d15656498e 100644 --- a/src/spanner/src/channel_pool/config.rs +++ b/src/spanner/src/channel_pool/config.rs @@ -23,7 +23,7 @@ pub(crate) const MAX_SUPPORTED_CHANNELS: usize = 256; /// Strategy used to select channels from the active pool. /// /// # Example -/// ```no_rust +/// ``` /// use google_cloud_spanner::channel_pool::{ChannelSelectionStrategy, DynamicChannelPoolConfig}; /// /// let config = DynamicChannelPoolConfig::new() @@ -31,8 +31,7 @@ pub(crate) const MAX_SUPPORTED_CHANNELS: usize = 256; /// ``` #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] #[non_exhaustive] -#[allow(dead_code)] -pub(crate) enum ChannelSelectionStrategy { +pub enum ChannelSelectionStrategy { /// Power of Two Least Busy (samples 2 candidates, picks lower effective load, breaks ties with warmer channel). #[default] PowerOfTwoLeastBusy, @@ -41,21 +40,22 @@ pub(crate) enum ChannelSelectionStrategy { /// Configuration for the Spanner client channel pool. /// /// # Example -/// ```no_rust -/// use google_cloud_spanner::client::{Spanner, SpannerPoolBuilderExt}; -/// use google_cloud_spanner::channel_pool::{ChannelPoolConfig, StaticChannelPoolConfig}; -/// +/// ``` +/// # use google_cloud_spanner::client::{Spanner, SpannerBuilderExt}; +/// # use google_cloud_spanner::channel_pool::{ChannelPoolConfig, StaticChannelPoolConfig}; +/// # async fn sample() -> anyhow::Result<()> { /// let config = ChannelPoolConfig::from(StaticChannelPoolConfig::new(8)); /// let client = Spanner::builder() /// .with_channel_pool(config) /// .build() /// .await?; +/// # Ok(()) } /// ``` /// /// Supports either static fixed-size channel pooling or autonomous dynamic load-based channel scaling. #[derive(Clone, Debug, PartialEq)] #[non_exhaustive] -pub(crate) enum ChannelPoolConfig { +pub enum ChannelPoolConfig { /// Fixed-size static channel pool (default: 4 channels). Static(StaticChannelPoolConfig), /// Dynamic load-based channel pool. @@ -77,18 +77,16 @@ impl ChannelPoolConfig { } } - /// Returns a reference to the `DynamicChannelPoolConfig` if dynamic. - pub(crate) fn dynamic_config(&self) -> Option<&DynamicChannelPoolConfig> { + /// Returns a reference to the [`DynamicChannelPoolConfig`] if dynamic. + pub fn dynamic_config(&self) -> Option<&DynamicChannelPoolConfig> { match self { Self::Dynamic(config) => Some(config), Self::Static(_) => None, } } -} -#[cfg(test)] -impl ChannelPoolConfig { - pub(crate) fn static_config(&self) -> Option<&StaticChannelPoolConfig> { + /// Returns a reference to the [`StaticChannelPoolConfig`] if static. + pub fn static_config(&self) -> Option<&StaticChannelPoolConfig> { match self { Self::Static(config) => Some(config), Self::Dynamic(_) => None, @@ -99,19 +97,20 @@ impl ChannelPoolConfig { /// Configuration for a static (fixed-size) channel pool. /// /// # Example -/// ```no_rust -/// use google_cloud_spanner::client::{Spanner, SpannerPoolBuilderExt}; -/// use google_cloud_spanner::channel_pool::StaticChannelPoolConfig; -/// +/// ``` +/// # use google_cloud_spanner::client::{Spanner, SpannerBuilderExt}; +/// # use google_cloud_spanner::channel_pool::StaticChannelPoolConfig; +/// # async fn sample() -> anyhow::Result<()> { /// let config = StaticChannelPoolConfig::new(8); /// let client = Spanner::builder() /// .with_channel_pool(config) /// .build() /// .await?; +/// # Ok(()) } /// ``` #[derive(Clone, Debug, PartialEq, Eq)] #[non_exhaustive] -pub(crate) struct StaticChannelPoolConfig { +pub struct StaticChannelPoolConfig { /// Number of channels in the static pool (default: 4). pub(crate) num_channels: usize, } @@ -122,24 +121,29 @@ impl Default for StaticChannelPoolConfig { } } -#[allow(dead_code)] impl StaticChannelPoolConfig { /// Creates a new static channel pool configuration with the specified number of channels. /// /// # Example - /// ```no_rust - /// use google_cloud_spanner::client::{Spanner, SpannerPoolBuilderExt}; - /// use google_cloud_spanner::channel_pool::StaticChannelPoolConfig; - /// + /// ``` + /// # use google_cloud_spanner::client::{Spanner, SpannerBuilderExt}; + /// # use google_cloud_spanner::channel_pool::StaticChannelPoolConfig; + /// # async fn sample() -> anyhow::Result<()> { /// let client = Spanner::builder() /// .with_channel_pool(StaticChannelPoolConfig::new(4)) /// .build() /// .await?; + /// # Ok(()) } /// ``` - pub(crate) fn new(num_channels: usize) -> Self { + pub fn new(num_channels: usize) -> Self { Self { num_channels } } + /// Returns the number of channels configured for the static pool. + pub fn num_channels(&self) -> usize { + self.num_channels + } + /// Validates the static pool configuration. pub(crate) fn validate(&self) -> Result<(), GaxError> { if self.num_channels == 0 { @@ -163,10 +167,10 @@ impl From for ChannelPoolConfig { /// Configuration for a dynamically scaling channel pool. /// /// # Example -/// ```no_rust -/// use google_cloud_spanner::client::{Spanner, SpannerPoolBuilderExt}; -/// use google_cloud_spanner::channel_pool::DynamicChannelPoolConfig; -/// +/// ``` +/// # use google_cloud_spanner::client::{Spanner, SpannerBuilderExt}; +/// # use google_cloud_spanner::channel_pool::DynamicChannelPoolConfig; +/// # async fn sample() -> anyhow::Result<()> { /// let config = DynamicChannelPoolConfig::new() /// .with_initial_channels(4) /// .with_min_channels(2) @@ -176,39 +180,40 @@ impl From for ChannelPoolConfig { /// .with_channel_pool(config) /// .build() /// .await?; +/// # Ok(()) } /// ``` /// /// Manages autonomous elastic scaling of gRPC channels based on in-flight RPC load and error feedback. #[derive(Clone, Debug, PartialEq)] #[non_exhaustive] -pub(crate) struct DynamicChannelPoolConfig { +pub struct DynamicChannelPoolConfig { /// Number of channels created eagerly at startup (default: 4). pub(crate) initial_channels: usize, /// Minimum number of channels retained during scale-down (default: 4). pub(crate) min_channels: usize, - /// Maximum number of channels allowed during scale-up (default: 10, configurable up to 256). + /// Maximum number of channels allowed during scale-up (default: 256, configurable up to 256). pub(crate) max_channels: usize, - /// Low-load threshold (per channel) triggering scale-down evaluation (default: 15.0). + /// Low-load threshold (per channel) triggering scale-down evaluation (default: 2.0). pub(crate) min_rpc_per_channel: f64, - /// High-load threshold (per channel) triggering scale-up (default: 25.0). + /// High-load threshold (per channel) triggering scale-up (default: 8.0). pub(crate) max_rpc_per_channel: f64, - /// Synthetic picker load added per qualifying error (default: 5). + /// Synthetic picker load added per qualifying error (default: 2). pub(crate) error_penalty_step: u32, /// Sliding window duration for active error penalties (default: 5 seconds). pub(crate) error_penalty_duration: Duration, - /// Interval between periodic scale-down evaluations (default: 3 minutes). + /// Interval between periodic scale-down evaluations (default: 1 minute). pub(crate) scale_down_check_interval: Duration, - /// Cooldown period between consecutive scale-up bursts (default: 10 seconds). + /// Cooldown period between consecutive scale-up bursts (default: 1 second). pub(crate) scale_up_cooldown: Duration, /// Number of consecutive low-load checks required before scale-down (default: 3). pub(crate) consecutive_low_load_checks: usize, - /// Maximum percentage of current pool size added per scale-up event (default: 30%, min 2). + /// Maximum percentage of current pool size added per scale-up event (default: 100%, min 2). pub(crate) max_scale_up_percent: u32, - /// Maximum number of channels marked draining per scale-down cycle (default: 2). + /// Maximum number of channels marked draining per scale-down cycle (default: 4). pub(crate) max_remove_channels: usize, /// Idle grace duration a draining channel is kept alive after load drops to 0 (default: 1 minute). pub(crate) drain_idle_grace: Duration, - /// Timeout for executing SELECT 1 priming on a new scaled-up channel (default: 10 seconds). + /// Timeout for executing SELECT 1 priming on a new scaled-up channel (default: 5 seconds). pub(crate) prime_timeout: Duration, /// Maximum retry attempts for SELECT 1 priming (default: 3). pub(crate) prime_max_attempts: usize, @@ -221,138 +226,218 @@ impl Default for DynamicChannelPoolConfig { Self { initial_channels: 4, min_channels: 4, - max_channels: 10, - min_rpc_per_channel: 15.0, - max_rpc_per_channel: 25.0, - error_penalty_step: 5, + max_channels: 256, + min_rpc_per_channel: 2.0, + max_rpc_per_channel: 8.0, + error_penalty_step: 2, error_penalty_duration: Duration::from_secs(5), - scale_down_check_interval: Duration::from_secs(180), - scale_up_cooldown: Duration::from_secs(10), + scale_down_check_interval: Duration::from_secs(60), + scale_up_cooldown: Duration::from_secs(1), consecutive_low_load_checks: 3, - max_scale_up_percent: 30, - max_remove_channels: 2, + max_scale_up_percent: 100, + max_remove_channels: 4, drain_idle_grace: Duration::from_secs(60), - prime_timeout: Duration::from_secs(10), + prime_timeout: Duration::from_secs(5), prime_max_attempts: 3, selection_strategy: ChannelSelectionStrategy::PowerOfTwoLeastBusy, } } } -#[allow(dead_code)] impl DynamicChannelPoolConfig { /// Creates a new default dynamic channel pool configuration. /// /// # Example - /// ```no_rust - /// use google_cloud_spanner::client::{Spanner, SpannerPoolBuilderExt}; - /// use google_cloud_spanner::channel_pool::DynamicChannelPoolConfig; - /// + /// ``` + /// # use google_cloud_spanner::client::{Spanner, SpannerBuilderExt}; + /// # use google_cloud_spanner::channel_pool::DynamicChannelPoolConfig; + /// # async fn sample() -> anyhow::Result<()> { /// let client = Spanner::builder() /// .with_channel_pool(DynamicChannelPoolConfig::new()) /// .build() /// .await?; + /// # Ok(()) } /// ``` - pub(crate) fn new() -> Self { + pub fn new() -> Self { Self::default() } /// Sets the number of channels created eagerly at startup. - pub(crate) fn with_initial_channels(mut self, channels: usize) -> Self { + pub fn with_initial_channels(mut self, channels: usize) -> Self { self.initial_channels = channels; self } /// Sets the minimum number of channels retained during scale-down. - pub(crate) fn with_min_channels(mut self, channels: usize) -> Self { + pub fn with_min_channels(mut self, channels: usize) -> Self { self.min_channels = channels; self } /// Sets the maximum number of channels allowed during scale-up. - pub(crate) fn with_max_channels(mut self, channels: usize) -> Self { + pub fn with_max_channels(mut self, channels: usize) -> Self { self.max_channels = channels; self } /// Sets the low-load threshold (per channel) triggering scale-down evaluation. - pub(crate) fn with_min_rpc_per_channel(mut self, min_rpc: f64) -> Self { + pub fn with_min_rpc_per_channel(mut self, min_rpc: f64) -> Self { self.min_rpc_per_channel = min_rpc; self } /// Sets the high-load threshold (per channel) triggering scale-up. - pub(crate) fn with_max_rpc_per_channel(mut self, max_rpc: f64) -> Self { + pub fn with_max_rpc_per_channel(mut self, max_rpc: f64) -> Self { self.max_rpc_per_channel = max_rpc; self } /// Sets the synthetic picker load added per qualifying transport error. - pub(crate) fn with_error_penalty_step(mut self, step: u32) -> Self { + pub fn with_error_penalty_step(mut self, step: u32) -> Self { self.error_penalty_step = step; self } /// Sets the sliding window duration for active error penalties. - pub(crate) fn with_error_penalty_duration(mut self, duration: Duration) -> Self { + pub fn with_error_penalty_duration(mut self, duration: Duration) -> Self { self.error_penalty_duration = duration; self } /// Sets the interval between periodic scale-down evaluations. - pub(crate) fn with_scale_down_check_interval(mut self, interval: Duration) -> Self { + pub fn with_scale_down_check_interval(mut self, interval: Duration) -> Self { self.scale_down_check_interval = interval; self } /// Sets the cooldown period between consecutive scale-up bursts. - pub(crate) fn with_scale_up_cooldown(mut self, cooldown: Duration) -> Self { + pub fn with_scale_up_cooldown(mut self, cooldown: Duration) -> Self { self.scale_up_cooldown = cooldown; self } /// Sets the number of consecutive low-load checks required before scale-down. - pub(crate) fn with_consecutive_low_load_checks(mut self, checks: usize) -> Self { + pub fn with_consecutive_low_load_checks(mut self, checks: usize) -> Self { self.consecutive_low_load_checks = checks; self } /// Sets the maximum percentage of current pool size added per scale-up event. - pub(crate) fn with_max_scale_up_percent(mut self, percent: u32) -> Self { + pub fn with_max_scale_up_percent(mut self, percent: u32) -> Self { self.max_scale_up_percent = percent; self } /// Sets the maximum number of channels marked draining per scale-down cycle. - pub(crate) fn with_max_remove_channels(mut self, max_channels: usize) -> Self { + pub fn with_max_remove_channels(mut self, max_channels: usize) -> Self { self.max_remove_channels = max_channels; self } /// Sets the idle grace duration a draining channel is kept alive after load drops to zero. - pub(crate) fn with_drain_idle_grace(mut self, grace: Duration) -> Self { + pub fn with_drain_idle_grace(mut self, grace: Duration) -> Self { self.drain_idle_grace = grace; self } /// Sets the timeout for executing `SELECT 1` priming on a new scaled-up channel. - pub(crate) fn with_prime_timeout(mut self, timeout: Duration) -> Self { + pub fn with_prime_timeout(mut self, timeout: Duration) -> Self { self.prime_timeout = timeout; self } /// Sets the maximum retry attempts for `SELECT 1` priming. - pub(crate) fn with_prime_max_attempts(mut self, attempts: usize) -> Self { + pub fn with_prime_max_attempts(mut self, attempts: usize) -> Self { self.prime_max_attempts = attempts; self } /// Sets the channel selection strategy. - pub(crate) fn with_selection_strategy(mut self, strategy: ChannelSelectionStrategy) -> Self { + pub fn with_selection_strategy(mut self, strategy: ChannelSelectionStrategy) -> Self { self.selection_strategy = strategy; self } + /// Returns the number of channels created eagerly at startup. + pub fn initial_channels(&self) -> usize { + self.initial_channels + } + + /// Returns the minimum number of channels retained during scale-down. + pub fn min_channels(&self) -> usize { + self.min_channels + } + + /// Returns the maximum number of channels allowed during scale-up. + pub fn max_channels(&self) -> usize { + self.max_channels + } + + /// Returns the low-load threshold (per channel) triggering scale-down evaluation. + pub fn min_rpc_per_channel(&self) -> f64 { + self.min_rpc_per_channel + } + + /// Returns the high-load threshold (per channel) triggering scale-up. + pub fn max_rpc_per_channel(&self) -> f64 { + self.max_rpc_per_channel + } + + /// Returns the synthetic picker load added per qualifying transport error. + pub fn error_penalty_step(&self) -> u32 { + self.error_penalty_step + } + + /// Returns the sliding window duration for active error penalties. + pub fn error_penalty_duration(&self) -> Duration { + self.error_penalty_duration + } + + /// Returns the interval between periodic scale-down evaluations. + pub fn scale_down_check_interval(&self) -> Duration { + self.scale_down_check_interval + } + + /// Returns the cooldown period between consecutive scale-up bursts. + pub fn scale_up_cooldown(&self) -> Duration { + self.scale_up_cooldown + } + + /// Returns the number of consecutive low-load checks required before scale-down. + pub fn consecutive_low_load_checks(&self) -> usize { + self.consecutive_low_load_checks + } + + /// Returns the maximum percentage of current pool size added per scale-up event. + pub fn max_scale_up_percent(&self) -> u32 { + self.max_scale_up_percent + } + + /// Returns the maximum number of channels marked draining per scale-down cycle. + pub fn max_remove_channels(&self) -> usize { + self.max_remove_channels + } + + /// Returns the idle grace duration a draining channel is kept alive after load drops to zero. + pub fn drain_idle_grace(&self) -> Duration { + self.drain_idle_grace + } + + /// Returns the timeout for executing `SELECT 1` priming on a new scaled-up channel. + pub fn prime_timeout(&self) -> Duration { + self.prime_timeout + } + + /// Returns the maximum retry attempts for `SELECT 1` priming. + pub fn prime_max_attempts(&self) -> usize { + self.prime_max_attempts + } + + /// Returns the channel selection strategy. + pub fn selection_strategy(&self) -> ChannelSelectionStrategy { + self.selection_strategy + } + /// Validates dynamic channel pool configuration boundaries and invariant relationships. pub(crate) fn validate(&self) -> Result<(), GaxError> { if self.min_channels == 0 { @@ -423,7 +508,7 @@ impl DynamicChannelPoolConfig { self.max_rpc_per_channel.ceil() as u32 } - /// Computes the midpoint target RPC capacity per channel (e.g. (15 + 25) / 2 = 20). + /// Computes the midpoint target RPC capacity per channel (e.g. (2 + 8) / 2 = 5). pub(crate) fn target_rpc_per_channel(&self) -> u32 { let midpoint = ((self.min_rpc_per_channel + self.max_rpc_per_channel) / 2.0).floor() as u32; midpoint.max(1) @@ -504,20 +589,20 @@ mod tests { "DynamicChannelPoolConfig default min_channels must be 4" ); assert_eq!( - dynamic_config.max_channels, 10, - "DynamicChannelPoolConfig default max_channels must be 10" + dynamic_config.max_channels, 256, + "DynamicChannelPoolConfig default max_channels must be 256" ); assert_eq!( - dynamic_config.min_rpc_per_channel, 15.0, - "DynamicChannelPoolConfig default min_rpc_per_channel must be 15.0" + dynamic_config.min_rpc_per_channel, 2.0, + "DynamicChannelPoolConfig default min_rpc_per_channel must be 2.0" ); assert_eq!( - dynamic_config.max_rpc_per_channel, 25.0, - "DynamicChannelPoolConfig default max_rpc_per_channel must be 25.0" + dynamic_config.max_rpc_per_channel, 8.0, + "DynamicChannelPoolConfig default max_rpc_per_channel must be 8.0" ); assert_eq!( - dynamic_config.error_penalty_step, 5, - "DynamicChannelPoolConfig default error_penalty_step must be 5" + dynamic_config.error_penalty_step, 2, + "DynamicChannelPoolConfig default error_penalty_step must be 2" ); assert_eq!( dynamic_config.error_penalty_duration, @@ -526,22 +611,50 @@ mod tests { ); assert_eq!( dynamic_config.error_penalty_max(), - 25, - "DynamicChannelPoolConfig default error_penalty_max must be 25" + 8, + "DynamicChannelPoolConfig default error_penalty_max must be 8" + ); + assert_eq!( + dynamic_config.scale_down_check_interval, + Duration::from_secs(60), + "DynamicChannelPoolConfig default scale_down_check_interval must be 60s" + ); + assert_eq!( + dynamic_config.scale_up_cooldown, + Duration::from_secs(1), + "DynamicChannelPoolConfig default scale_up_cooldown must be 1s" ); assert_eq!( dynamic_config.consecutive_low_load_checks, 3, "DynamicChannelPoolConfig default consecutive_low_load_checks must be 3" ); assert_eq!( - dynamic_config.max_remove_channels, 2, - "DynamicChannelPoolConfig default max_remove_channels must be 2" + dynamic_config.max_scale_up_percent, 100, + "DynamicChannelPoolConfig default max_scale_up_percent must be 100" + ); + assert_eq!( + dynamic_config.max_remove_channels, 4, + "DynamicChannelPoolConfig default max_remove_channels must be 4" ); assert_eq!( dynamic_config.drain_idle_grace, Duration::from_secs(60), "DynamicChannelPoolConfig default drain_idle_grace must be 60s" ); + assert_eq!( + dynamic_config.prime_timeout, + Duration::from_secs(5), + "DynamicChannelPoolConfig default prime_timeout must be 5s" + ); + assert_eq!( + dynamic_config.prime_max_attempts, 3, + "DynamicChannelPoolConfig default prime_max_attempts must be 3" + ); + assert_eq!( + dynamic_config.selection_strategy, + ChannelSelectionStrategy::PowerOfTwoLeastBusy, + "DynamicChannelPoolConfig default selection_strategy must be PowerOfTwoLeastBusy" + ); assert!( dynamic_config.validate().is_ok(), "Default DynamicChannelPoolConfig must pass validation" @@ -928,6 +1041,11 @@ mod tests { fn static_channel_pool_config_new() { let config = StaticChannelPoolConfig::new(8); assert_eq!(config.num_channels, 8, "num_channels must be 8"); + assert_eq!( + config.num_channels(), + 8, + "num_channels() getter must return 8" + ); assert!(config.validate().is_ok(), "validation must succeed"); } @@ -952,66 +1070,146 @@ mod tests { .with_selection_strategy(ChannelSelectionStrategy::PowerOfTwoLeastBusy); assert_eq!(config.initial_channels, 5, "initial_channels must match"); + assert_eq!( + config.initial_channels(), + 5, + "initial_channels() getter must return 5" + ); assert_eq!(config.min_channels, 3, "min_channels must match"); + assert_eq!( + config.min_channels(), + 3, + "min_channels() getter must return 3" + ); assert_eq!(config.max_channels, 12, "max_channels must match"); + assert_eq!( + config.max_channels(), + 12, + "max_channels() getter must return 12" + ); assert_eq!( config.min_rpc_per_channel, 10.0, "min_rpc_per_channel must match" ); + assert_eq!( + config.min_rpc_per_channel(), + 10.0, + "min_rpc_per_channel() getter must return 10.0" + ); assert_eq!( config.max_rpc_per_channel, 20.0, "max_rpc_per_channel must match" ); + assert_eq!( + config.max_rpc_per_channel(), + 20.0, + "max_rpc_per_channel() getter must return 20.0" + ); assert_eq!( config.error_penalty_step, 8, "error_penalty_step must match" ); + assert_eq!( + config.error_penalty_step(), + 8, + "error_penalty_step() getter must return 8" + ); assert_eq!( config.error_penalty_duration, Duration::from_secs(15), "error_penalty_duration must match" ); + assert_eq!( + config.error_penalty_duration(), + Duration::from_secs(15), + "error_penalty_duration() getter must return 15s" + ); assert_eq!( config.scale_down_check_interval, Duration::from_secs(120), "scale_down_check_interval must match" ); + assert_eq!( + config.scale_down_check_interval(), + Duration::from_secs(120), + "scale_down_check_interval() getter must return 120s" + ); assert_eq!( config.scale_up_cooldown, Duration::from_secs(30), "scale_up_cooldown must match" ); + assert_eq!( + config.scale_up_cooldown(), + Duration::from_secs(30), + "scale_up_cooldown() getter must return 30s" + ); assert_eq!( config.consecutive_low_load_checks, 5, "consecutive_low_load_checks must match" ); + assert_eq!( + config.consecutive_low_load_checks(), + 5, + "consecutive_low_load_checks() getter must return 5" + ); assert_eq!( config.max_scale_up_percent, 50, "max_scale_up_percent must match" ); + assert_eq!( + config.max_scale_up_percent(), + 50, + "max_scale_up_percent() getter must return 50" + ); assert_eq!( config.max_remove_channels, 3, "max_remove_channels must match" ); + assert_eq!( + config.max_remove_channels(), + 3, + "max_remove_channels() getter must return 3" + ); assert_eq!( config.drain_idle_grace, Duration::from_secs(90), "drain_idle_grace must match" ); + assert_eq!( + config.drain_idle_grace(), + Duration::from_secs(90), + "drain_idle_grace() getter must return 90s" + ); assert_eq!( config.prime_timeout, Duration::from_secs(20), "prime_timeout must match" ); + assert_eq!( + config.prime_timeout(), + Duration::from_secs(20), + "prime_timeout() getter must return 20s" + ); assert_eq!( config.prime_max_attempts, 5, "prime_max_attempts must match" ); + assert_eq!( + config.prime_max_attempts(), + 5, + "prime_max_attempts() getter must return 5" + ); assert_eq!( config.selection_strategy, ChannelSelectionStrategy::PowerOfTwoLeastBusy, "selection_strategy must match" ); + assert_eq!( + config.selection_strategy(), + ChannelSelectionStrategy::PowerOfTwoLeastBusy, + "selection_strategy() getter must return PowerOfTwoLeastBusy" + ); assert!(config.validate().is_ok(), "validation must succeed"); } } diff --git a/src/spanner/src/channel_pool/entry.rs b/src/spanner/src/channel_pool/entry.rs index 707a0bebd9..2c6fe780ac 100644 --- a/src/spanner/src/channel_pool/entry.rs +++ b/src/spanner/src/channel_pool/entry.rs @@ -327,31 +327,12 @@ impl ChannelLease { Self { guard } } - /// Consumes the lease, returning the underlying active RPC guard. - pub(crate) fn into_guard(self) -> ActiveRpcGuard { - self.guard - } - - /// Records the result of an RPC call and applies an error penalty if a qualifying error occurred. - pub(crate) fn record_result( - &self, - result: &Result, - extract_code: impl Fn(&E) -> Option, - ) { - self.guard.record_result(result, extract_code); - } - /// Records the result of a standard GAX RPC call, extracting the gRPC status code if present. pub(crate) fn record_call_result(&self, result: &crate::Result) { self.guard .record_result(result, |error| error.status().map(|status| status.code)); } - /// Returns a reference to the physical `Channel`. - pub(crate) fn channel(&self) -> &Channel { - &self.guard.entry.channel - } - /// Returns the unique monotonic internal entry ID. pub(crate) fn entry_id(&self) -> u64 { self.guard.entry.id @@ -420,6 +401,13 @@ mod tests { ); } + #[tokio::test] + async fn mock_stub_create_session() { + let channel = create_mock_channel(); + let result = channel.inner.create_session().send().await; + assert!(result.is_ok(), "mock session create must succeed"); + } + #[test] fn error_penalty_allowlist_and_sliding_expiry() { let channel = create_mock_channel(); @@ -716,11 +704,10 @@ mod tests { "entry_id() must return entry's internal id 42" ); assert_eq!( - lease.channel().channel_id, - 3, - "channel.channel_id must match entry's logical id 3" + lease.channel_id, 3, + "channel_id via Deref must match entry's logical id 3" ); - let _channel = lease.channel(); + let _channel: &Channel = &lease; // rw_affinity_guard helper creates an RAII guard incrementing active_rw_transactions let rw_guard = lease.rw_affinity_guard(); @@ -743,61 +730,48 @@ mod tests { } #[test] - fn channel_lease_into_guard() { + fn channel_lease_lifecycle_and_deref() { let channel = create_mock_channel(); let entry = Arc::new(ChannelEntry::new(42, 3, channel)); assert_eq!(entry.in_flight(), 0, "initial in-flight count must be 0"); - let guard = ActiveRpcGuard::new(Arc::clone(&entry), 0, Duration::ZERO, 0); + let guard = ActiveRpcGuard::new(Arc::clone(&entry), 5, Duration::from_secs(10), 10); let lease = ChannelLease::new(guard); assert_eq!( entry.in_flight(), 1, - "creating guard must increment in-flight count" - ); - - let guard = lease.into_guard(); - assert_eq!( - entry.in_flight(), - 1, - "into_guard must preserve in-flight count" + "creating lease must increment in-flight count" ); - drop(guard); - assert_eq!( - entry.in_flight(), - 0, - "dropping ActiveRpcGuard must decrement in-flight count" - ); - } - - #[test] - fn channel_lease_record_result_and_deref() { - let channel = create_mock_channel(); - let entry = Arc::new(ChannelEntry::new(42, 3, channel)); - let guard = ActiveRpcGuard::new(Arc::clone(&entry), 5, Duration::from_secs(10), 10); - let lease = ChannelLease::new(guard); - // Verify Deref to Channel assert_eq!( lease.channel_id, 3, "Deref must allow accessing underlying channel fields" ); - let ok_result: Result<&str, Status> = Ok("success"); - lease.record_result(&ok_result, |status| Some(status.code)); + let ok_result: crate::Result<&str> = Ok("success"); + lease.record_call_result(&ok_result); assert_eq!( entry.current_penalty(), 0, - "record_result on Ok must not add penalty load" + "record_call_result on Ok must not add penalty load" ); - let err_result: Result<&str, Status> = Err(Status::default().set_code(Code::Unavailable)); - lease.record_result(&err_result, |status| Some(status.code)); + let err_result: crate::Result<&str> = Err(google_cloud_gax::error::Error::service( + Status::default().set_code(Code::Unavailable), + )); + lease.record_call_result(&err_result); assert_eq!( entry.current_penalty(), 5, - "record_result on qualifying error must add penalty load" + "record_call_result on qualifying error must add penalty load" + ); + + drop(lease); + assert_eq!( + entry.in_flight(), + 0, + "dropping ChannelLease must decrement in-flight count" ); } diff --git a/src/spanner/src/channel_pool/integration_tests.rs b/src/spanner/src/channel_pool/integration_tests.rs new file mode 100644 index 0000000000..51b48ace05 --- /dev/null +++ b/src/spanner/src/channel_pool/integration_tests.rs @@ -0,0 +1,1706 @@ +// 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. + +//! Semantic integration verification for Spanner channel pools. +//! +//! Validates that both dynamic and static channel pools are fully wired into +//! [`DatabaseClient`] and all transaction types, verifying: +//! 1. In-flight RPC permits increment during active execution and decrement to zero on completion. +//! 2. Early stream drops and errors release permits immediately without leaks. +//! 3. Read/Write transaction affinity hard-pins all operations to the exact same channel entry. +//! 4. Multi-use Read-Only transaction affinity pins queries across statements. +//! 5. Write-only and Partitioned DML transactions preserve channel affinity. +//! 6. Dynamic channel pools scale up under concurrent load from `DatabaseClient`. +//! 7. Static channel pools maintain fixed channel capacity under concurrent load. +//! 8. Transport error penalties deprioritize degraded channels in P2C selection. +//! 9. Mixed workloads leave zero leaked permits or active Read/Write transaction guards. + +use crate::channel_pool::entry::ChannelEntry; +use crate::channel_pool::{DynamicChannelPoolConfig, StaticChannelPoolConfig}; +use crate::client::{Spanner, SpannerBuilderExt}; +use crate::database_client::DatabaseClient; +use crate::key::KeySet; +use crate::model::PartitionOptions; +use crate::mutation::Mutation; +use crate::read::ReadRequest; +use crate::read_only_transaction::BeginTransactionOption; +use crate::read_only_transaction::tests::{create_session_mock, setup_select1}; +use crate::read_write_transaction::ReadWriteTransactionBuilder; +use crate::result_set::tests::adapt; +use crate::statement::Statement; +use gaxi::grpc::tonic::{Response, Status}; +use google_cloud_auth::credentials::anonymous::Builder as Anonymous; +use google_cloud_test_macros::tokio_test_no_panics; +use spanner_grpc_mock::google::spanner::v1 as mock_v1; +use spanner_grpc_mock::{MockSpanner, start}; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; +use tokio::time::sleep; + +async fn setup_client_with_static_pool( + mock: MockSpanner, + channel_count: usize, +) -> (DatabaseClient, Spanner, tokio::task::JoinHandle<()>) { + let (address, server) = start("127.0.0.1:0", mock) + .await + .expect("mock server should start"); + + let static_config = StaticChannelPoolConfig::new(channel_count); + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(static_config) + .build() + .await + .expect("spanner client should build"); + + let database_client = spanner + .database_client("projects/p/instances/i/databases/d") + .build() + .await + .expect("database client should build"); + + (database_client, spanner, server) +} + +async fn setup_client_with_dynamic_pool( + mock: MockSpanner, + config: DynamicChannelPoolConfig, +) -> (DatabaseClient, Spanner, tokio::task::JoinHandle<()>) { + let (address, server) = start("127.0.0.1:0", mock) + .await + .expect("mock server should start"); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(config) + .build() + .await + .expect("spanner client should build"); + + let database_client = spanner + .database_client("projects/p/instances/i/databases/d") + .build() + .await + .expect("database client should build"); + + (database_client, spanner, server) +} + +fn total_in_flight(spanner: &Spanner) -> u32 { + let active_guard = spanner + .channel_pool + .inner + .active_entries + .read() + .expect("lock poisoned"); + active_guard.iter().map(|entry| entry.in_flight()).sum() +} + +fn total_active_rw(spanner: &Spanner) -> u32 { + let active_guard = spanner + .channel_pool + .inner + .active_entries + .read() + .expect("lock poisoned"); + active_guard + .iter() + .map(|entry| entry.active_rw_count()) + .sum() +} + +fn active_channel_entries(spanner: &Spanner) -> Vec> { + let active_guard = spanner + .channel_pool + .inner + .active_entries + .read() + .expect("lock poisoned"); + active_guard.clone() +} + +async fn wait_for_in_flight(spanner: &Spanner, expected: u32, timeout_duration: Duration) { + let start = tokio::time::Instant::now(); + while start.elapsed() < timeout_duration { + if total_in_flight(spanner) == expected { + return; + } + sleep(Duration::from_millis(5)).await; + } + assert_eq!( + total_in_flight(spanner), + expected, + "Timed out waiting for in-flight permits to reach {expected}" + ); +} + +async fn wait_for_channels_greater_than( + spanner: &Spanner, + threshold: usize, + timeout_duration: Duration, +) { + let start = tokio::time::Instant::now(); + while start.elapsed() < timeout_duration { + if spanner.channel_pool.active_channel_count() > threshold { + return; + } + sleep(Duration::from_millis(5)).await; + } + let count = spanner.channel_pool.active_channel_count(); + assert!( + count > threshold, + "Timed out waiting for channel count ({count}) to exceed {threshold}" + ); +} + +fn setup_chunk(last: bool) -> mock_v1::PartialResultSet { + let mut chunk = setup_select1(); + chunk.last = last; + chunk +} + +#[tokio_test_no_panics] +async fn single_use_query_in_flight_lifecycle() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + mock.expect_execute_streaming_sql() + .once() + .returning(|_| Ok(Response::from(adapt([Ok(setup_select1())])))); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 2).await; + + assert_eq!( + total_in_flight(&spanner), + 0, + "Pool must have 0 in-flight before query" + ); + + let mut result_set = database_client + .single_use() + .build() + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + + assert_eq!( + total_in_flight(&spanner), + 1, + "Active result set stream must hold exactly 1 in-flight permit" + ); + + let row = result_set.next().await; + assert!(row.is_some(), "Stream must return first row"); + + let end_of_stream = result_set.next().await; + assert!(end_of_stream.is_none(), "Stream must reach EOF"); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn single_use_read_in_flight_lifecycle() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + mock.expect_streaming_read() + .once() + .returning(|_| Ok(Response::from(adapt([Ok(setup_select1())])))); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 2).await; + + assert_eq!( + total_in_flight(&spanner), + 0, + "Pool must have 0 in-flight before read" + ); + + let read_request = ReadRequest::builder("Users", vec!["id".to_string()]) + .with_keys(KeySet::all()) + .build(); + + let mut result_set = database_client + .single_use() + .build() + .execute_read(read_request) + .await?; + + assert_eq!( + total_in_flight(&spanner), + 1, + "Active read stream must hold exactly 1 in-flight permit" + ); + + let row = result_set.next().await; + assert!(row.is_some(), "Read stream must return first row"); + + let end_of_stream = result_set.next().await; + assert!(end_of_stream.is_none(), "Read stream must reach EOF"); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn streaming_query_early_drop_releases_in_flight_immediately() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + mock.expect_execute_streaming_sql().once().returning(|_| { + Ok(Response::from(adapt([ + Ok(setup_chunk(false)), + Ok(setup_chunk(false)), + Ok(setup_chunk(true)), + ]))) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 2).await; + + let mut result_set = database_client + .single_use() + .build() + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + + assert_eq!( + total_in_flight(&spanner), + 1, + "In-flight permit must be held by open stream" + ); + + let first_row = result_set.next().await; + assert!(first_row.is_some(), "First row should be available"); + + // Drop the ResultSet before stream finishes + drop(result_set); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(100)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn streaming_query_initial_error_releases_in_flight_and_applies_penalty() -> anyhow::Result<()> +{ + let mut mock = create_session_mock(); + mock.expect_execute_streaming_sql() + .once() + .returning(|_| Err(Status::unavailable("simulated network partition"))); + + let config = DynamicChannelPoolConfig::new() + .with_initial_channels(2) + .with_min_channels(2) + .with_max_channels(4) + .with_min_rpc_per_channel(2.0) + .with_max_rpc_per_channel(10.0) + .with_error_penalty_step(3) + .with_scale_up_cooldown(Duration::from_millis(500)); + let (database_client, spanner, _server) = setup_client_with_dynamic_pool(mock, config).await; + + let query_result = database_client + .single_use() + .build() + .execute_query(Statement::builder("SELECT 1").build()) + .await; + + assert!( + query_result.is_err(), + "Query must propagate unavailable error" + ); + + assert_eq!( + total_in_flight(&spanner), + 0, + "Error on stream creation must decrement in-flight permit immediately" + ); + + let entries = active_channel_entries(&spanner); + let penalized_channel_count = entries + .iter() + .filter(|entry| entry.current_penalty() > 0) + .count(); + assert_eq!( + penalized_channel_count, 1, + "The failed channel entry must record a synthetic error penalty" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn multi_use_read_only_transaction_pins_same_channel_inline_begin() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + for _ in 0..2 { + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + let mut partial_result_set = setup_select1(); + let metadata = partial_result_set + .metadata + .as_mut() + .expect("metadata present"); + metadata.transaction = Some(mock_v1::Transaction { + id: vec![10, 20, 30], + ..Default::default() + }); + Ok(Response::from(adapt([Ok(partial_result_set)]))) + }); + } + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = database_client + .read_only_transaction() + .with_begin_transaction_option(BeginTransactionOption::InlineBegin) + .build() + .await?; + + let mut result_set1 = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set1.next().await).is_some() {} + + let mut result_set2 = transaction + .execute_query(Statement::builder("SELECT 2").build()) + .await?; + while (result_set2.next().await).is_some() {} + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!(addresses.len(), 2, "Both queries must execute"); + assert_eq!( + addresses[0], addresses[1], + "Both queries in multi-use read-only transaction must pin to the exact same channel" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn multi_use_read_only_transaction_pins_same_channel_explicit_begin() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_begin_transaction() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::Transaction { + id: vec![10, 20, 30], + ..Default::default() + })) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::from(adapt([Ok(setup_select1())]))) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = database_client + .read_only_transaction() + .with_begin_transaction_option(BeginTransactionOption::ExplicitBegin) + .build() + .await?; + + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set.next().await).is_some() {} + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!( + addresses.len(), + 2, + "BeginTransaction and query must execute" + ); + assert_eq!( + addresses[0], addresses[1], + "BeginTransaction and query in multi-use transaction must pin to the exact same channel" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn multi_use_read_only_transaction_parallel_initial_queries_inline_begin() +-> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + // One query will be leader (has begin), the second will wait for transaction ID + mock.expect_execute_streaming_sql().times(2).returning({ + let captured_addresses = Arc::clone(&captured_addresses); + move |request| { + let address = request.remote_addr().expect("remote_addr"); + captured_addresses.lock().expect("lock").push(address); + + let message = request.into_inner(); + let has_begin = message + .transaction + .as_ref() + .and_then(|selector| selector.selector.as_ref()) + .is_some_and(|mode| { + matches!(mode, mock_v1::transaction_selector::Selector::Begin(_)) + }); + + let (sender, receiver) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + if has_begin { + // Leader simulates small network delay then returns transaction ID + sleep(Duration::from_millis(50)).await; + let mut partial_result_set = setup_select1(); + let metadata = partial_result_set + .metadata + .as_mut() + .expect("metadata present"); + metadata.transaction = Some(mock_v1::Transaction { + id: vec![77, 88, 99], + ..Default::default() + }); + let _ = sender.send(Ok(partial_result_set)).await; + } else { + // Follower query runs after receiving transaction ID + let _ = sender.send(Ok(setup_select1())).await; + } + }); + + Ok(Response::from(receiver)) + } + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = database_client + .read_only_transaction() + .with_begin_transaction_option(BeginTransactionOption::InlineBegin) + .build() + .await?; + + let query1_future = async { + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set.next().await).is_some() {} + Ok::<(), anyhow::Error>(()) + }; + + let query2_future = async { + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 2").build()) + .await?; + while (result_set.next().await).is_some() {} + Ok::<(), anyhow::Error>(()) + }; + + let (result1, result2) = tokio::join!(query1_future, query2_future); + result1?; + result2?; + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!(addresses.len(), 2, "Both parallel queries must execute"); + assert_eq!( + addresses[0], addresses[1], + "Parallel initial queries in inline-begin read-only transaction must pin to identical channel" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn multi_use_read_only_transaction_parallel_initial_queries_explicit_begin() +-> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + // 1. BeginTransaction called during build() + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_begin_transaction() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::Transaction { + id: vec![77, 88, 99], + ..Default::default() + })) + }); + + // 2. Both parallel queries execute concurrently on the wire + mock.expect_execute_streaming_sql().times(2).returning({ + let captured_addresses = Arc::clone(&captured_addresses); + move |request| { + let address = request.remote_addr().expect("remote_addr"); + captured_addresses.lock().expect("lock").push(address); + + let (sender, receiver) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + // Hold stream open to allow concurrent in-flight overlap + sleep(Duration::from_millis(60)).await; + let _ = sender.send(Ok(setup_select1())).await; + }); + + Ok(Response::from(receiver)) + } + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = database_client + .read_only_transaction() + .with_begin_transaction_option(BeginTransactionOption::ExplicitBegin) + .build() + .await?; + + let query1_future = async { + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set.next().await).is_some() {} + Ok::<(), anyhow::Error>(()) + }; + + let query2_future = async { + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 2").build()) + .await?; + while (result_set.next().await).is_some() {} + Ok::<(), anyhow::Error>(()) + }; + + let (result1, result2) = tokio::join!(query1_future, query2_future); + result1?; + result2?; + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!( + addresses.len(), + 3, + "BeginTransaction and both queries must execute" + ); + assert_eq!( + addresses[0], addresses[1], + "Query 1 must route to the same channel as BeginTransaction" + ); + assert_eq!( + addresses[1], addresses[2], + "Query 2 must route to the same channel as Query 1" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn read_write_transaction_full_lifecycle_hard_affinity() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + // 1. Query: execute_streaming_sql (starts transaction via inline begin) + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + let mut partial_result_set = setup_select1(); + let metadata = partial_result_set + .metadata + .as_mut() + .expect("metadata present"); + metadata.transaction = Some(mock_v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + }); + Ok(Response::from(adapt([Ok(partial_result_set)]))) + }); + + // 2. Update: execute_sql + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_sql().once().returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::ResultSet { + metadata: Some(mock_v1::ResultSetMetadata { + row_type: Some(mock_v1::StructType { fields: vec![] }), + ..Default::default() + }), + stats: Some(mock_v1::ResultSetStats { + row_count: Some(mock_v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + })) + }); + + // 3. Batch DML: execute_batch_dml + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_batch_dml() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::ExecuteBatchDmlResponse { + result_sets: vec![mock_v1::ResultSet { + stats: Some(mock_v1::ResultSetStats { + row_count: Some(mock_v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + }], + ..Default::default() + })) + }); + + // 4. Commit: commit + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_commit().once().returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 5000, + nanos: 0, + }), + ..Default::default() + })) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + assert_eq!( + total_active_rw(&spanner), + 0, + "Active RW count must start at 0" + ); + + let transaction = ReadWriteTransactionBuilder::new(database_client) + .with_begin_transaction_option(BeginTransactionOption::InlineBegin) + .build(None) + .await?; + + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set.next().await).is_some() {} + + assert_eq!( + total_active_rw(&spanner), + 1, + "Active RW guard must be attached to pinned channel during transaction" + ); + + let updated_rows = transaction + .execute_update("UPDATE Users SET Name = 'Alice' WHERE Id = 1") + .await?; + assert_eq!(updated_rows, 1, "Must update 1 row"); + + let batch_result = transaction + .execute_batch_update(vec![Statement::from( + "UPDATE Users SET Name = 'Bob' WHERE Id = 2", + )]) + .await?; + assert_eq!(batch_result.len(), 1); + + transaction.commit().await?; + + assert_eq!( + total_active_rw(&spanner), + 0, + "Active RW guard must be released immediately upon transaction commit" + ); + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!( + addresses.len(), + 4, + "Query, Update, Batch DML, and Commit must all be captured" + ); + assert_eq!( + addresses[0], addresses[1], + "Update must route to identical channel as initial query" + ); + assert_eq!( + addresses[1], addresses[2], + "Batch DML must route to identical channel as update" + ); + assert_eq!( + addresses[2], addresses[3], + "Commit must route to identical channel as prior operations" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn read_write_transaction_parallel_initial_queries_inline_begin() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + // One statement initiates the inline begin, second waits for transaction ID + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + + let (sender, receiver) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + sleep(Duration::from_millis(50)).await; + let mut partial_result_set = setup_select1(); + let metadata = partial_result_set + .metadata + .as_mut() + .expect("metadata present"); + metadata.transaction = Some(mock_v1::Transaction { + id: vec![55, 66, 77], + ..Default::default() + }); + let _ = sender.send(Ok(partial_result_set)).await; + }); + Ok(Response::from(receiver)) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_sql().once().returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::ResultSet { + metadata: Some(mock_v1::ResultSetMetadata { + row_type: Some(mock_v1::StructType { fields: vec![] }), + ..Default::default() + }), + stats: Some(mock_v1::ResultSetStats { + row_count: Some(mock_v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + })) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_commit().once().returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 5000, + nanos: 0, + }), + ..Default::default() + })) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = ReadWriteTransactionBuilder::new(database_client) + .with_begin_transaction_option(BeginTransactionOption::InlineBegin) + .build(None) + .await?; + + let op1 = async { + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set.next().await).is_some() {} + Ok::<(), anyhow::Error>(()) + }; + + let op2 = async { + let _ = transaction + .execute_update("UPDATE Users SET Name = 'Alice' WHERE Id = 1") + .await?; + Ok::<(), anyhow::Error>(()) + }; + + let (res1, res2) = tokio::join!(op1, op2); + res1?; + res2?; + + assert_eq!( + total_active_rw(&spanner), + 1, + "Active RW guard must be held during transaction" + ); + + transaction.commit().await?; + + assert_eq!( + total_active_rw(&spanner), + 0, + "Active RW guard released after commit" + ); + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!(addresses.len(), 3, "Query, Update, and Commit executed"); + assert_eq!( + addresses[0], addresses[1], + "Query and Update must route to identical channel" + ); + assert_eq!( + addresses[1], addresses[2], + "Commit must route to identical channel" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn read_write_transaction_parallel_initial_queries_explicit_begin() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_begin_transaction() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::Transaction { + id: vec![55, 66, 77], + ..Default::default() + })) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + + let (sender, receiver) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + sleep(Duration::from_millis(50)).await; + let _ = sender.send(Ok(setup_select1())).await; + }); + Ok(Response::from(receiver)) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_sql().once().returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::ResultSet { + metadata: Some(mock_v1::ResultSetMetadata { + row_type: Some(mock_v1::StructType { fields: vec![] }), + ..Default::default() + }), + stats: Some(mock_v1::ResultSetStats { + row_count: Some(mock_v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + })) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_commit().once().returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 5000, + nanos: 0, + }), + ..Default::default() + })) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = ReadWriteTransactionBuilder::new(database_client) + .with_begin_transaction_option(BeginTransactionOption::ExplicitBegin) + .build(None) + .await?; + + assert_eq!( + total_active_rw(&spanner), + 1, + "Active RW guard attached upon explicit begin" + ); + + let op1 = async { + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set.next().await).is_some() {} + Ok::<(), anyhow::Error>(()) + }; + + let op2 = async { + let _ = transaction + .execute_update("UPDATE Users SET Name = 'Alice' WHERE Id = 1") + .await?; + Ok::<(), anyhow::Error>(()) + }; + + let (res1, res2) = tokio::join!(op1, op2); + res1?; + res2?; + + transaction.commit().await?; + + assert_eq!( + total_active_rw(&spanner), + 0, + "Active RW guard released after commit" + ); + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!( + addresses.len(), + 4, + "BeginTransaction, Query, Update, and Commit executed" + ); + assert_eq!( + addresses[0], addresses[1], + "Query must route to identical channel as BeginTransaction" + ); + assert_eq!( + addresses[1], addresses[2], + "Update must route to identical channel as Query" + ); + assert_eq!( + addresses[2], addresses[3], + "Commit must route to identical channel as Update" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn read_write_transaction_rollback_releases_rw_guard() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + let mut partial_result_set = setup_select1(); + let metadata = partial_result_set + .metadata + .as_mut() + .expect("metadata present"); + metadata.transaction = Some(mock_v1::Transaction { + id: vec![7, 8, 9], + ..Default::default() + }); + Ok(Response::from(adapt([Ok(partial_result_set)]))) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_rollback().once().returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(())) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = ReadWriteTransactionBuilder::new(database_client) + .build(None) + .await?; + + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set.next().await).is_some() {} + + assert_eq!( + total_active_rw(&spanner), + 1, + "Active RW guard held during active transaction" + ); + + transaction.rollback().await?; + + assert_eq!( + total_active_rw(&spanner), + 0, + "Active RW guard released upon rollback" + ); + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!(addresses.len(), 2, "Query and Rollback executed"); + assert_eq!( + addresses[0], addresses[1], + "Rollback must route to the identical pinned channel" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn read_write_transaction_drop_releases_rw_guard() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + + mock.expect_execute_streaming_sql().once().returning(|_| { + let mut partial_result_set = setup_select1(); + let metadata = partial_result_set + .metadata + .as_mut() + .expect("metadata present"); + metadata.transaction = Some(mock_v1::Transaction { + id: vec![99], + ..Default::default() + }); + Ok(Response::from(adapt([Ok(partial_result_set)]))) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = ReadWriteTransactionBuilder::new(database_client) + .build(None) + .await?; + + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set.next().await).is_some() {} + + assert_eq!( + total_active_rw(&spanner), + 1, + "Active RW guard held by live transaction" + ); + + drop(result_set); + drop(transaction); + + assert_eq!( + total_active_rw(&spanner), + 0, + "Dropping active transaction and its result sets must decrement active RW guard immediately" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn write_only_transaction_pins_begin_and_commit_and_releases_guards() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_begin_transaction() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::Transaction { + id: vec![55, 66], + ..Default::default() + })) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_commit().once().returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 100, + nanos: 0, + }), + ..Default::default() + })) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = database_client.write_only_transaction().build(); + let mutations = vec![Mutation::delete("Users", KeySet::all())]; + let response = transaction.write(mutations).await?; + + assert!( + response.commit_timestamp.is_some(), + "Write transaction must return commit timestamp" + ); + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!(addresses.len(), 2, "Begin and Commit executed"); + assert_eq!( + addresses[0], addresses[1], + "Write-only transaction Begin and Commit must pin to identical channel" + ); + + assert_eq!( + total_active_rw(&spanner), + 0, + "Write-only transaction must not leak RW guards" + ); + assert_eq!( + total_in_flight(&spanner), + 0, + "Write-only transaction must leave 0 in-flight permits" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn write_only_transaction_write_at_least_once_lifecycle() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + + mock.expect_commit().once().returning(|_| { + Ok(Response::new(mock_v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 200, + nanos: 0, + }), + ..Default::default() + })) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = database_client.write_only_transaction().build(); + let mutations = vec![Mutation::delete("Users", KeySet::all())]; + let response = transaction.write_at_least_once(mutations).await?; + + assert!(response.commit_timestamp.is_some()); + + assert_eq!( + total_active_rw(&spanner), + 0, + "write_at_least_once must have 0 active RW guards" + ); + assert_eq!( + total_in_flight(&spanner), + 0, + "write_at_least_once must leave 0 in-flight permits" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn partitioned_dml_pins_begin_and_execute_and_releases() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_begin_transaction() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::Transaction { + id: vec![77, 88], + ..Default::default() + })) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + let mut partial_result_set = setup_select1(); + partial_result_set.stats = Some(mock_v1::ResultSetStats { + row_count: Some(mock_v1::result_set_stats::RowCount::RowCountLowerBound(100)), + ..Default::default() + }); + Ok(Response::from(adapt([Ok(partial_result_set)]))) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = database_client + .partitioned_dml_transaction() + .build() + .await?; + + let modified = transaction + .execute_update("UPDATE Users SET active = true WHERE true") + .await?; + assert_eq!(modified, 100, "Must return modified row count"); + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!(addresses.len(), 2, "Begin and Streaming SQL executed"); + assert_eq!( + addresses[0], addresses[1], + "Partitioned DML Begin and Execute must route to the same pinned channel" + ); + + assert_eq!( + total_active_rw(&spanner), + 0, + "Active RW guard must be released after Partitioned DML" + ); + assert_eq!( + total_in_flight(&spanner), + 0, + "In-flight permits must be 0 after Partitioned DML" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn batch_read_only_transaction_routes_through_pool() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_begin_transaction() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::Transaction { + id: vec![44, 55, 66], + ..Default::default() + })) + }); + + let addresses_clone = Arc::clone(&captured_addresses); + mock.expect_partition_query() + .once() + .returning(move |request| { + addresses_clone + .lock() + .expect("lock") + .push(request.remote_addr().expect("remote_addr")); + Ok(Response::new(mock_v1::PartitionResponse { + partitions: vec![mock_v1::Partition { + partition_token: vec![1, 2, 3], + }], + transaction: Some(mock_v1::Transaction { + id: vec![44, 55, 66], + ..Default::default() + }), + })) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + let transaction = database_client + .batch_read_only_transaction() + .build() + .await?; + + let partitions = transaction + .partition_query( + Statement::builder("SELECT * FROM Users").build(), + PartitionOptions::default(), + ) + .await?; + + assert_eq!(partitions.len(), 1, "Must return 1 partition"); + + let addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!( + addresses.len(), + 2, + "BeginTransaction and partition_query must execute" + ); + assert_eq!( + addresses[0], addresses[1], + "BeginTransaction and partition_query must route to the same pinned channel" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn dynamic_channel_pool_scales_up_under_concurrent_queries() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + + // Scale-up worker primes newly opened channels using unary execute_sql("SELECT 1") + mock.expect_execute_sql().returning(|_| { + Ok(Response::new(mock_v1::ResultSet { + metadata: Some(mock_v1::ResultSetMetadata { + row_type: Some(mock_v1::StructType { fields: vec![] }), + ..Default::default() + }), + ..Default::default() + })) + }); + + // 8 concurrent queries that hold their stream open for 100ms to create sustained concurrency + mock.expect_execute_streaming_sql().returning(|_| { + let (sender, receiver) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + sleep(Duration::from_millis(100)).await; + let _ = sender.send(Ok(setup_select1())).await; + }); + Ok(Response::from(receiver)) + }); + + // Configure dynamic pool with initial=2, min=2, max=6, min_rpc=0.5, max_rpc=1.5, penalty_step=1, cooldown=30ms + let config = DynamicChannelPoolConfig::new() + .with_initial_channels(2) + .with_min_channels(2) + .with_max_channels(6) + .with_min_rpc_per_channel(0.5) + .with_max_rpc_per_channel(1.5) + .with_error_penalty_step(1) + .with_scale_up_cooldown(Duration::from_millis(30)); + let (database_client, spanner, _server) = setup_client_with_dynamic_pool(mock, config).await; + + assert_eq!( + spanner.channel_pool.active_channel_count(), + 2, + "Initial dynamic channel pool size must be 2" + ); + + // Launch 8 concurrent queries + let mut tasks = Vec::new(); + for _ in 0..8 { + let client_clone = database_client.clone(); + tasks.push(tokio::spawn(async move { + let mut result_set = client_clone + .single_use() + .build() + .execute_query(Statement::builder("SELECT 1").build()) + .await + .expect("query should succeed"); + while let Some(row) = result_set.next().await { + let _ = row.expect("row should succeed"); + } + })); + } + + for task in tasks { + task.await.expect("task join should succeed"); + } + + wait_for_channels_greater_than(&spanner, 2, Duration::from_millis(500)).await; + wait_for_in_flight(&spanner, 0, Duration::from_millis(500)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn static_channel_pool_remains_fixed_under_concurrent_queries() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + + mock.expect_execute_streaming_sql().returning(|_| { + let (sender, receiver) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + sleep(Duration::from_millis(60)).await; + let _ = sender.send(Ok(setup_select1())).await; + }); + Ok(Response::from(receiver)) + }); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 2).await; + + assert_eq!( + spanner.channel_pool.active_channel_count(), + 2, + "Static pool must have exactly 2 channels" + ); + + let mut tasks = Vec::new(); + for _ in 0..6 { + let client_clone = database_client.clone(); + tasks.push(tokio::spawn(async move { + let mut result_set = client_clone + .single_use() + .build() + .execute_query(Statement::builder("SELECT 1").build()) + .await + .expect("query succeeds"); + while let Some(row) = result_set.next().await { + let _ = row.expect("row succeeds"); + } + })); + } + + for task in tasks { + task.await.expect("task join succeeds"); + } + + assert_eq!( + spanner.channel_pool.active_channel_count(), + 2, + "Static channel pool must remain strictly at 2 channels under load" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn p2c_avoids_channel_with_error_penalty() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + let call_count = Arc::new(AtomicUsize::new(0)); + let captured_addresses = Arc::new(Mutex::new(Vec::new())); + + let call_count_clone = Arc::clone(&call_count); + let captured_addresses_clone = Arc::clone(&captured_addresses); + + mock.expect_execute_streaming_sql() + .returning(move |request| { + let count = call_count_clone.fetch_add(1, Ordering::SeqCst); + let address = request.remote_addr().expect("remote_addr"); + captured_addresses_clone.lock().expect("lock").push(address); + + if count == 0 { + // First RPC fails with UNAVAILABLE to apply error penalty + Err(Status::unavailable("simulated failure for error penalty")) + } else { + Ok(Response::from(adapt([Ok(setup_select1())]))) + } + }); + + let config = DynamicChannelPoolConfig::new() + .with_initial_channels(2) + .with_min_channels(2) + .with_max_channels(4) + .with_min_rpc_per_channel(2.0) + .with_max_rpc_per_channel(10.0) + .with_error_penalty_step(5) + .with_scale_up_cooldown(Duration::from_secs(5)); + let (database_client, spanner, _server) = setup_client_with_dynamic_pool(mock, config).await; + + // First query fails and penalizes its channel + let _ = database_client + .single_use() + .build() + .execute_query(Statement::builder("SELECT 1").build()) + .await; + + let penalized_address = { + let addrs = captured_addresses.lock().expect("lock"); + assert_eq!(addrs.len(), 1, "First query executed"); + addrs[0] + }; + + // Execute 6 subsequent queries + for _ in 0..6 { + let mut result_set = database_client + .single_use() + .build() + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (result_set.next().await).is_some() {} + } + + let all_addresses = captured_addresses.lock().expect("lock").clone(); + assert_eq!(all_addresses.len(), 7, "All 7 queries executed"); + + // The subsequent 6 queries must avoid the penalized channel whenever P2C compares it + let subsequent_on_penalized = all_addresses[1..] + .iter() + .filter(|addr| **addr == penalized_address) + .count(); + + assert!( + subsequent_on_penalized < 6, + "P2C must deprioritize the penalized channel (penalized channel received {subsequent_on_penalized} out of 6 queries)" + ); + + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio_test_no_panics] +async fn zero_permit_leak_mixed_workload_battery() -> anyhow::Result<()> { + let mut mock = create_session_mock(); + + mock.expect_execute_streaming_sql().returning(|_| { + Ok(Response::from(adapt([ + Ok(setup_chunk(false)), + Ok(setup_chunk(true)), + ]))) + }); + + mock.expect_streaming_read() + .returning(|_| Ok(Response::from(adapt([Ok(setup_select1())])))); + + mock.expect_execute_sql().returning(|_| { + Ok(Response::new(mock_v1::ResultSet { + metadata: Some(mock_v1::ResultSetMetadata { + row_type: Some(mock_v1::StructType { fields: vec![] }), + transaction: Some(mock_v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + }), + ..Default::default() + }), + stats: Some(mock_v1::ResultSetStats { + row_count: Some(mock_v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + })) + }); + + mock.expect_commit().returning(|_| { + Ok(Response::new(mock_v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 1234, + nanos: 0, + }), + ..Default::default() + })) + }); + + mock.expect_rollback().returning(|_| Ok(Response::new(()))); + + let (database_client, spanner, _server) = setup_client_with_static_pool(mock, 4).await; + + // 1. Single-use query full read + let mut rs = database_client + .single_use() + .build() + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + while (rs.next().await).is_some() {} + + // 2. Single-use read full read + let mut rs = database_client + .single_use() + .build() + .execute_read( + ReadRequest::builder("Users", vec!["id".to_string()]) + .with_keys(KeySet::all()) + .build(), + ) + .await?; + while (rs.next().await).is_some() {} + + // 3. Early dropped query stream + let mut rs = database_client + .single_use() + .build() + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + let _ = rs.next().await; + drop(rs); + + // 4. ReadWriteTransaction commit + let rw_tx = ReadWriteTransactionBuilder::new(database_client.clone()) + .build(None) + .await?; + let _ = rw_tx.execute_update("UPDATE Users SET x = 1").await?; + rw_tx.commit().await?; + + // 5. ReadWriteTransaction rollback + let rw_tx = ReadWriteTransactionBuilder::new(database_client.clone()) + .build(None) + .await?; + let _ = rw_tx.execute_update("UPDATE Users SET x = 2").await?; + rw_tx.rollback().await?; + + // 6. ReadWriteTransaction dropped without commit/rollback + let rw_tx = ReadWriteTransactionBuilder::new(database_client.clone()) + .build(None) + .await?; + let _ = rw_tx.execute_update("UPDATE Users SET x = 3").await?; + drop(rw_tx); + + // 7. Write-only transaction write_at_least_once + let _ = database_client + .write_only_transaction() + .build() + .write_at_least_once(vec![Mutation::delete("Users", KeySet::all())]) + .await?; + + // Final Assertion: ZERO leaked permits or guards across entire pool + wait_for_in_flight(&spanner, 0, Duration::from_millis(200)).await; + + assert_eq!( + total_active_rw(&spanner), + 0, + "Total active Read/Write guards across all channels must be exactly 0" + ); + + Ok(()) +} diff --git a/src/spanner/src/channel_pool/mod.rs b/src/spanner/src/channel_pool/mod.rs index f904d2047e..3c756f3848 100644 --- a/src/spanner/src/channel_pool/mod.rs +++ b/src/spanner/src/channel_pool/mod.rs @@ -12,29 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Dynamic Channel Pooling for Spanner. +//! Channel pooling for Spanner. //! //! Provides capacity management, load-balanced channel selection (Power of Two Least Busy), //! health-aware error penalization, caller-owned transaction affinity pinning, and background //! scaling and priming for gRPC channels. -// TODO(dynamic-channel-pooling): Remove allow(dead_code, unused_imports) once integrated into Spanner client. -#![allow(dead_code)] -#![allow(unused_imports)] - pub(crate) mod affinity; -pub(crate) mod config; +mod config; pub(crate) mod entry; pub(crate) mod pool; pub(crate) mod scaler; -#[allow(unused_imports)] -pub(crate) use affinity::TransactionAffinity; -#[allow(unused_imports)] -pub(crate) use config::DynamicChannelPoolConfig; -#[allow(unused_imports)] -pub(crate) use config::{ChannelPoolConfig, StaticChannelPoolConfig}; -#[allow(unused_imports)] +#[cfg(test)] +mod integration_tests; + +pub use config::{ + ChannelPoolConfig, ChannelSelectionStrategy, DynamicChannelPoolConfig, StaticChannelPoolConfig, +}; + +pub(crate) use affinity::{ChannelTarget, TransactionAffinity}; pub(crate) use entry::ChannelLease; -#[allow(unused_imports)] pub(crate) use pool::ChannelPool; diff --git a/src/spanner/src/channel_pool/pool.rs b/src/spanner/src/channel_pool/pool.rs index 305d9108f6..c4780ced1b 100644 --- a/src/spanner/src/channel_pool/pool.rs +++ b/src/spanner/src/channel_pool/pool.rs @@ -17,7 +17,7 @@ //! Provides `ChannelPool`, which unifies both static (fixed-size) and dynamically scaling channel //! pool configurations under a single API for the Spanner client. -use crate::channel_pool::affinity::TransactionAffinity; +use crate::channel_pool::affinity::{ChannelTarget, TransactionAffinity}; use crate::channel_pool::config::{ ChannelPoolConfig, DynamicChannelPoolConfig, MAX_SUPPORTED_CHANNELS, StaticChannelPoolConfig, }; @@ -27,7 +27,7 @@ use crate::client::Channel; use crate::routing::power_of_two_selector::PowerOfTwoSelector; use gaxi::options::ClientConfig; use std::fmt::{Debug, Formatter, Result as FmtResult}; -use std::sync::atomic::{AtomicU64, AtomicUsize}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, RwLock}; use std::time::{Duration, Instant}; use tokio::spawn; @@ -82,6 +82,7 @@ impl ChannelPool { draining_entries: RwLock::new(Vec::new()), next_entry_id, scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -109,6 +110,7 @@ impl ChannelPool { draining_entries: RwLock::new(Vec::new()), next_entry_id, scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -147,6 +149,7 @@ impl ChannelPool { guard.entry.effective_pick_load() as f64 > dynamic_config.max_rpc_per_channel }) { + self.inner.scale_up_requested.store(true, Ordering::Release); self.inner.scale_up_notify.notify_one(); } @@ -160,6 +163,26 @@ impl ChannelPool { self.pick_from_slice(&active_guard) } + /// Returns the 1-based logical channel ID for the target. + pub(crate) fn logical_channel_id_for_target(&self, target: &ChannelTarget<'_>) -> usize { + match target { + ChannelTarget::Affinity(affinity) => { + if let Some(entry_id) = affinity.pinned_entry_id() { + let active_guard = self.inner.active_entries.read().expect("lock poisoned"); + if let Some(entry) = active_guard.iter().find(|entry| entry.id == entry_id) { + return entry.logical_channel_id(); + } + let draining_guard = self.inner.draining_entries.read().expect("lock poisoned"); + if let Some(entry) = draining_guard.iter().find(|entry| entry.id == entry_id) { + return entry.logical_channel_id(); + } + } + 1 + } + ChannelTarget::Any => 1, + } + } + /// Resolves an affinity handle to a leased channel. /// /// # Transaction Affinity Routing Invariants @@ -344,15 +367,9 @@ impl ChannelPool { /// Returns a clone of the first active channel in the pool, if present. /// - /// # Warning - /// - /// This method is intended strictly for internal metadata setup (e.g. configuring - /// fallback gateway connection endpoints during client initialization). - /// - /// **Never** use this method for routing or executing queries or RPCs. Doing so would - /// bypass load balancing and cause traffic to herd onto the first channel. - /// Use [`ChannelPool::pick_channel`] for P2C load-balanced channel selection, or - /// [`ChannelPool::resolve_affinity`] for operations requiring transaction affinity. + /// This helper is intended exclusively for unit tests to verify channel pool initialization + /// and channel availability. It is never used in production request routing. + #[cfg(test)] pub(crate) fn default_channel(&self) -> Option { let active_guard = self.inner.active_entries.read().expect("lock poisoned"); active_guard.first().map(|entry| entry.channel.clone()) @@ -367,6 +384,7 @@ pub(crate) struct ChannelPoolInner { pub(crate) draining_entries: RwLock>>, pub(crate) next_entry_id: AtomicU64, pub(crate) scale_up_notify: Arc, + pub(crate) scale_up_requested: AtomicBool, #[allow(dead_code)] // Retained for RAII drop signaling; read in scaler unit tests via subscribe() pub(crate) shutdown_sender: WatchSender<()>, @@ -1080,6 +1098,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1108,6 +1127,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1267,4 +1287,57 @@ mod tests { "must adopt winning channel 2 on CAS conflict" ); } + + #[tokio::test] + async fn mock_stub_create_session() { + let channel = create_mock_channel(); + let result = channel.inner.create_session().send().await; + assert!(result.is_ok(), "mock session create must succeed"); + } + + #[test] + fn empty_pool_hint_handling() { + let client_config = ClientConfig::default(); + let pool = ChannelPool::new_static( + vec![], + StaticChannelPoolConfig { num_channels: 0 }, + client_config, + ); + + assert!( + pool.pick_channel().is_none(), + "pick_channel on empty pool must return None" + ); + assert_eq!( + pool.logical_channel_id_for_target(&ChannelTarget::Any), + 1, + "logical_channel_id_for_target on ChannelTarget::Any must return 1" + ); + } + + #[test] + fn logical_channel_id_for_target_in_draining() { + let client_config = ClientConfig::default(); + let channel = Arc::new(ChannelEntry::new(10, 3, create_mock_channel())); + let pool = ChannelPool::new_static( + vec![], + StaticChannelPoolConfig { num_channels: 0 }, + client_config, + ); + pool.inner + .draining_entries + .write() + .expect("lock poisoned") + .push(Arc::clone(&channel)); + + let affinity = TransactionAffinity::new_read_write(); + affinity.set_pinned_entry_id_for_test(10); + + let target = ChannelTarget::Affinity(&affinity); + assert_eq!( + pool.logical_channel_id_for_target(&target), + 3, + "logical_channel_id_for_target must find entry in draining pool" + ); + } } diff --git a/src/spanner/src/channel_pool/scaler.rs b/src/spanner/src/channel_pool/scaler.rs index b8b1090422..b7d1cdc7d2 100644 --- a/src/spanner/src/channel_pool/scaler.rs +++ b/src/spanner/src/channel_pool/scaler.rs @@ -158,9 +158,10 @@ pub(crate) async fn scale_up_worker_loop( /// Evaluates scale-up eligibility, enforces capacity bounds, and calculates how many channels to add. /// -/// Returns `0` if capacity is already at `max_channels` or current capacity is sufficient for the -/// observed load. Commits cooldown timestamp to throttle subsequent picker wakeups when at capacity -/// ceiling or when scaling begins. Cooldown sleep is managed by the worker loop before invocation. +/// Returns `0` if capacity is already at `max_channels` or neither individual channel saturation +/// nor aggregate load warrants scale-up. Commits cooldown timestamp to throttle subsequent picker +/// wakeups when at capacity ceiling or when scaling begins. Cooldown sleep is managed by the worker +/// loop before invocation. fn calculate_scale_up_count(inner: &ChannelPoolInner, config: &DynamicChannelPoolConfig) -> usize { let active_guard = inner.active_entries.read().expect("lock poisoned"); let current_len = active_guard.len(); @@ -170,7 +171,7 @@ fn calculate_scale_up_count(inner: &ChannelPoolInner, config: &DynamicChannelPoo return 0; } - // 1. Sizing calculation: + // 1. Sizing and saturation evaluation: // desired_channels = ceil(total_load / target_rpc). // Note: Scale-up uses effective_pick_load() (in-flight + error penalty) to prompt // replacement capacity for failing channels. @@ -178,22 +179,35 @@ fn calculate_scale_up_count(inner: &ChannelPoolInner, config: &DynamicChannelPoo .iter() .map(|entry| entry.effective_pick_load()) .sum(); + let has_saturated_channel = active_guard + .iter() + .any(|entry| (entry.effective_pick_load() as f64) > config.max_rpc_per_channel); + let scale_up_requested = inner.scale_up_requested.swap(false, Ordering::AcqRel); + let is_saturated = (has_saturated_channel || scale_up_requested) && total_load > 0; let desired_channels = config.desired_channel_count(total_load); - if desired_channels <= current_len { + if !is_saturated && desired_channels <= current_len { return 0; } - // 2. Rate limiting: - // Add at most max_scale_up_percent (default 30%, minimum 2 channels) per scale event, + // 2. Rate limiting / doubling step: + // Add at most max_scale_up_percent (default 100%, minimum 2 channels) per scale event, // bounded by max_channels ceiling. let max_to_add_by_percent = ((current_len as f64) * (config.max_scale_up_percent as f64) / 100.0).ceil() as usize; let max_to_add_by_percent = max_to_add_by_percent.max(2); - let channels_to_add = (desired_channels - current_len) - .min(max_to_add_by_percent) - .min(config.max_channels - current_len); + let base_needed = desired_channels.saturating_sub(current_len); + let channels_to_add = if is_saturated { + // Individual channel saturation indicates immediate queuing/contention. + // Guarantee scaling by the scale-up step (doubling), avoiding the + // Little's Law aggregate veto on ultra-low-latency queries. + max_to_add_by_percent + } else { + // Aggregate load requires more capacity, but no single channel is saturated. + base_needed.min(max_to_add_by_percent) + } + .min(config.max_channels - current_len); // Only commit cooldown timestamp if channels are actually being added. // Note: We commit the scale-up cooldown timestamp here before awaiting dialing/priming @@ -544,8 +558,8 @@ mod tests { use spanner_grpc_mock::google::spanner::v1 as mock_v1; use std::fmt::Debug; use std::future::{Future, ready}; - use std::sync::atomic::{AtomicU64, AtomicUsize}; - use std::sync::{Mutex, RwLock}; + use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize}; + use std::sync::{Mutex, RwLock, Weak}; use tokio::sync::Notify; use tokio::sync::watch::channel as watch_channel; use tokio::task::yield_now; @@ -610,6 +624,7 @@ mod tests { ]), next_entry_id: AtomicU64::new(4), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -667,6 +682,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(5), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -746,6 +762,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(4), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -803,6 +820,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(4), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -873,6 +891,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(3), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -891,8 +910,25 @@ mod tests { "Cooldown timestamp must not be set on low load so subsequent bursts are not delayed" ); - // 2. High load: total load 90 -> target rpc 15 -> desired = 6 channels + // 2. Single-channel saturation with low aggregate load (Little's Law case): + // channel_1 receives 25 in-flight (> 20.0 max_rpc), channel_2 has 0 in-flight. + // total load = 25, target rpc = 15 -> desired = ceil(25 / 15) = 2 <= current (2). + // Saturated channel must NOT be vetoed; rate limit (min 2) scales up pool by 2. + channel_1.in_flight_rpcs.store(25, Ordering::Relaxed); + channel_2.in_flight_rpcs.store(0, Ordering::Relaxed); + let count_saturated = calculate_scale_up_count(&inner, &config); + assert_eq!( + count_saturated, 2, + "Saturated channel with low aggregate load must scale up without veto" + ); + assert!( + inner.last_scale_up_time.lock().expect("lock").is_some(), + "Cooldown timestamp must be set after scale-up count > 0" + ); + + // 3. High load with rate limiting: total load 90 -> target rpc 15 -> desired = 6 channels // current = 2, max_to_add = max(ceil(2 * 0.3) = 1, 2) = 2 -> channels_to_add = 2 + *inner.last_scale_up_time.lock().expect("lock") = None; channel_1.in_flight_rpcs.store(45, Ordering::Relaxed); channel_2.in_flight_rpcs.store(45, Ordering::Relaxed); let count = calculate_scale_up_count(&inner, &config); @@ -901,13 +937,50 @@ mod tests { "High load must calculate 2 channels to add based on rate limiting" ); - // Verify cooldown was committed - assert!( - inner.last_scale_up_time.lock().expect("lock").is_some(), - "Cooldown timestamp must be set after scale-up count > 0" + // 4. Default 100% max_scale_up_percent doubles pool size under saturation: + let doubling_config = DynamicChannelPoolConfig { + min_channels: 2, + max_channels: 16, + min_rpc_per_channel: 2.0, + max_rpc_per_channel: 8.0, + max_scale_up_percent: 100, + ..Default::default() + }; + let channel_3 = Arc::new(ChannelEntry::new(3, 3, create_mock_channel())); + let channel_4 = Arc::new(ChannelEntry::new(4, 4, create_mock_channel())); + *inner.active_entries.write().expect("lock") = vec![ + Arc::clone(&channel_1), + Arc::clone(&channel_2), + Arc::clone(&channel_3), + Arc::clone(&channel_4), + ]; + // Saturated channel on 4 channels: doubles pool size by adding 4 channels (4 -> 8) + channel_1.in_flight_rpcs.store(10, Ordering::Relaxed); + channel_2.in_flight_rpcs.store(0, Ordering::Relaxed); + channel_3.in_flight_rpcs.store(0, Ordering::Relaxed); + channel_4.in_flight_rpcs.store(0, Ordering::Relaxed); + let count_doubling = calculate_scale_up_count(&inner, &doubling_config); + assert_eq!( + count_doubling, 4, + "100% max_scale_up_percent must double 4 channels by adding 4" + ); + + // 5. Aggregate load exceeds capacity without any single channel saturated: + // 4 channels, each with 5 in-flight (below 8.0 max_rpc). Total load = 20. + // target rpc = (2 + 8) / 2 = 5 -> desired = ceil(20 / 5) = 4 <= 4 (no scale-up). + // With each at 6 in-flight: Total load = 24 -> desired = ceil(24 / 5) = 5 > 4. + // base_needed = 5 - 4 = 1. Scaler adds 1 channel. + channel_1.in_flight_rpcs.store(6, Ordering::Relaxed); + channel_2.in_flight_rpcs.store(6, Ordering::Relaxed); + channel_3.in_flight_rpcs.store(6, Ordering::Relaxed); + channel_4.in_flight_rpcs.store(6, Ordering::Relaxed); + let count_aggregate = calculate_scale_up_count(&inner, &doubling_config); + assert_eq!( + count_aggregate, 1, + "Aggregate load exceeding capacity without saturated channels adds base_needed" ); - // 3. Pool already at max_channels (8) -> returns 0 and commits cooldown timestamp + // 6. Pool already at max_channels (8) -> returns 0 and commits cooldown timestamp *inner.last_scale_up_time.lock().expect("lock") = None; let mut full_channels = Vec::new(); for index in 1..=8 { @@ -945,6 +1018,7 @@ mod tests { draining_entries: RwLock::new(vec![Arc::clone(&channel_2)]), next_entry_id: AtomicU64::new(3), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1034,6 +1108,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(5), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1127,6 +1202,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(13), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1162,6 +1238,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(23), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1207,6 +1284,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(7), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1275,6 +1353,7 @@ mod tests { ]), next_entry_id: AtomicU64::new(3), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1335,6 +1414,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1371,6 +1451,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1448,6 +1529,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1637,6 +1719,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(1), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1711,6 +1794,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(1), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1768,6 +1852,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1812,6 +1897,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(Some(Instant::now())), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1893,6 +1979,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(Some(Instant::now())), consecutive_low_load_checks: AtomicUsize::new(0), @@ -1979,6 +2066,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(Some(Instant::now())), consecutive_low_load_checks: AtomicUsize::new(0), @@ -2067,6 +2155,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -2117,6 +2206,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -2161,6 +2251,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(3), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -2245,6 +2336,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(Some(Instant::now() - Duration::from_millis(49))), consecutive_low_load_checks: AtomicUsize::new(0), @@ -2297,6 +2389,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -2332,6 +2425,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(2), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -2401,6 +2495,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(3), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -2469,9 +2564,9 @@ mod tests { let channel_3 = Arc::new(ChannelEntry::new(3, 3, create_mock_channel())); let channel_4 = Arc::new(ChannelEntry::new(4, 4, create_mock_channel())); - // Channel 1 receives a heavy burst of 35 in-flight RPCs (> 25 max_rpc), - // but aggregate load (35) across 4 channels is well within total capacity (desired = ceil(35/20) = 2 <= 4). - channel_1.in_flight_rpcs.store(35, Ordering::Relaxed); + // Channel 1 receives an unsaturating burst of 20 in-flight RPCs (below max_rpc of 25.0), + // and aggregate load (20) across 4 channels is well within total capacity (desired = ceil(20/20) = 1 <= 4). + channel_1.in_flight_rpcs.store(20, Ordering::Relaxed); let inner = Arc::new(ChannelPoolInner { config: ChannelPoolConfig::Dynamic(DynamicChannelPoolConfig { @@ -2492,6 +2587,7 @@ mod tests { draining_entries: RwLock::new(Vec::new()), next_entry_id: AtomicU64::new(5), scale_up_notify: Arc::new(Notify::new()), + scale_up_requested: AtomicBool::new(false), shutdown_sender: watch_channel(()).0, last_scale_up_time: Mutex::new(None), consecutive_low_load_checks: AtomicUsize::new(0), @@ -2533,4 +2629,49 @@ mod tests { "scale_up_worker_loop must terminate promptly on pool drop" ); } + + #[tokio::test] + async fn mock_stub_create_session() { + let channel = create_mock_channel(); + let result = channel.inner.create_session().send().await; + assert!(result.is_ok(), "mock session create must succeed"); + } + + #[tokio::test] + async fn scale_up_worker_loop_exits_when_weak_inner_is_none() { + let (_sender, receiver) = watch_channel(()); + let weak_inner = Weak::new(); + scale_up_worker_loop(weak_inner, receiver).await; + } + + #[tokio::test] + async fn scale_down_monitor_loop_exits_when_weak_inner_is_none() { + let (_sender, receiver) = watch_channel(()); + let weak_inner = Weak::new(); + scale_down_monitor_loop(weak_inner, receiver, Duration::from_millis(10)).await; + } + + #[tokio::test] + async fn scale_up_worker_loop_exits_for_static_config() { + let (_sender, receiver) = watch_channel(()); + let pool = ChannelPool::new_static( + vec![create_mock_channel()], + StaticChannelPoolConfig { num_channels: 1 }, + ClientConfig::default(), + ); + let weak_inner = Arc::downgrade(&pool.inner); + scale_up_worker_loop(weak_inner, receiver).await; + } + + #[tokio::test] + async fn scale_down_monitor_loop_exits_for_static_config() { + let (_sender, receiver) = watch_channel(()); + let pool = ChannelPool::new_static( + vec![create_mock_channel()], + StaticChannelPoolConfig { num_channels: 1 }, + ClientConfig::default(), + ); + let weak_inner = Arc::downgrade(&pool.inner); + scale_down_monitor_loop(weak_inner, receiver, Duration::from_millis(1)).await; + } } diff --git a/src/spanner/src/client.rs b/src/spanner/src/client.rs index 76cb3db66d..9a17ab18a3 100644 --- a/src/spanner/src/client.rs +++ b/src/spanner/src/client.rs @@ -16,7 +16,7 @@ use crate::ClientBuilderResult; use crate::RequestOptions; use crate::Result; use crate::channel_pool::{ - ChannelLease, ChannelPool, ChannelPoolConfig, DynamicChannelPoolConfig, + ChannelLease, ChannelPool, ChannelPoolConfig, ChannelTarget, DynamicChannelPoolConfig, StaticChannelPoolConfig, TransactionAffinity, }; use crate::generated::gapic_dataplane::client::Spanner as GapicSpanner; @@ -49,10 +49,7 @@ use http::{ header::{HeaderName, HeaderValue}, }; use std::env; -use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, -}; +use std::sync::Arc; use tokio::task::JoinSet; pub use crate::database_client::DatabaseClient; @@ -73,8 +70,6 @@ use opentelemetry::metrics::MeterProvider; #[derive(Clone, Debug)] pub struct Spanner { pub(crate) channel_pool: ChannelPool, - pub(crate) channels: Vec, - pub(crate) counter: Arc, pub(crate) config: ClientConfig, pub(crate) is_emulator: bool, pub(crate) instance_type: InstanceType, @@ -123,7 +118,7 @@ impl google_cloud_gax::client_builder::internal::ClientFactory for Factory { } let pool_config = resolve_pool_config(&mut config, is_emulator)?; - let (channel_pool, channels) = create_channel_pool(&config, pool_config).await?; + let channel_pool = create_channel_pool(&config, pool_config).await?; #[cfg(feature = "builtin-metrics")] let export_builtin_metrics_to_cloud_monitoring = config @@ -137,8 +132,6 @@ impl google_cloud_gax::client_builder::internal::ClientFactory for Factory { Ok(Spanner { channel_pool, - channels, - counter: Arc::new(AtomicUsize::new(0)), config, is_emulator, instance_type, @@ -158,6 +151,24 @@ pub type ClientBuilder = google_cloud_gax::client_builder::ClientBuilder anyhow::Result<()> { + /// let client = Spanner::builder() + /// .with_channel_pool(DynamicChannelPoolConfig::new()) + /// .build() + /// .await?; + /// # Ok(()) } + /// ``` + fn with_channel_pool>(self, pool_config: C) -> Self; + /// Sets the target [`InstanceType`] (`Cloud` vs `Omni`) for the Spanner client. /// /// # Example @@ -323,6 +334,10 @@ struct ExportBuiltinMetricsToCloudMonitoring(bool); struct ExportBuiltinMetricsToCustomProvider(bool); impl SpannerBuilderExt for ClientBuilder { + fn with_channel_pool>(self, pool_config: C) -> Self { + self.with_extension(pool_config.into()) + } + fn with_instance_type(self, instance_type: InstanceType) -> Self { self.with_extension(instance_type) } @@ -348,19 +363,6 @@ impl SpannerBuilderExt for ClientBuilder { } } -/// Builder extension trait for channel pool configuration. -#[allow(dead_code)] -pub(crate) trait SpannerPoolBuilderExt { - /// Configures the gRPC channel pool for the Spanner client. - fn with_channel_pool>(self, pool_config: C) -> Self; -} - -impl SpannerPoolBuilderExt for ClientBuilder { - fn with_channel_pool>(self, pool_config: C) -> Self { - self.with_extension(pool_config.into()) - } -} - fn parse_emulator_endpoint(endpoint: &str) -> String { match url::Url::parse(endpoint) { Ok(url) if url.has_host() => endpoint.to_string(), @@ -412,15 +414,26 @@ fn resolve_pool_config( config: &mut ClientConfig, is_emulator: bool, ) -> ClientBuilderResult { - resolve_pool_config_with(config, is_emulator, || { - env::var("SPANNER_NUM_CHANNELS").ok() - }) + resolve_pool_config_with( + config, + is_emulator, + || env::var("SPANNER_NUM_CHANNELS").ok(), + || env::var("SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL").ok(), + ) +} + +fn is_truthy(val: &str) -> bool { + matches!( + val.trim().to_ascii_lowercase().as_str(), + "true" | "1" | "yes" | "on" | "t" + ) } fn resolve_pool_config_with( config: &mut ClientConfig, is_emulator: bool, - env_lookup: impl FnOnce() -> Option, + num_channels_lookup: impl FnOnce() -> Option, + dynamic_pool_lookup: impl FnOnce() -> Option, ) -> ClientBuilderResult { if let Some(pool_config) = config.extensions.remove::() { let pool_config = Arc::unwrap_or_clone(pool_config); @@ -443,7 +456,30 @@ fn resolve_pool_config_with( dynamic_config.validate().map_err(BuilderError::transport)?; return Ok(ChannelPoolConfig::Dynamic(dynamic_config)); } - if let Some(num_channels_str) = env_lookup() { + let enable_dynamic = dynamic_pool_lookup() + .as_deref() + .map(is_truthy) + .unwrap_or(false); + + if enable_dynamic { + let mut dynamic_config = DynamicChannelPoolConfig::default(); + if let Some(num_channels_str) = num_channels_lookup() { + let trimmed = num_channels_str.trim(); + if !trimmed.is_empty() { + let num_channels = trimmed.parse::().map_err(BuilderError::transport)?; + if num_channels > dynamic_config.max_channels() { + dynamic_config = dynamic_config.with_max_channels(num_channels); + } + dynamic_config = dynamic_config + .with_initial_channels(num_channels) + .with_min_channels(num_channels); + } + } + dynamic_config.validate().map_err(BuilderError::transport)?; + return Ok(ChannelPoolConfig::Dynamic(dynamic_config)); + } + + if let Some(num_channels_str) = num_channels_lookup() { let trimmed = num_channels_str.trim(); if !trimmed.is_empty() { let num_channels = trimmed.parse::().map_err(BuilderError::transport)?; @@ -461,7 +497,7 @@ fn resolve_pool_config_with( async fn create_channel_pool( config: &ClientConfig, pool_config: ChannelPoolConfig, -) -> ClientBuilderResult<(ChannelPool, Vec)> { +) -> ClientBuilderResult { let num_initial = match &pool_config { ChannelPoolConfig::Static(static_config) => static_config.num_channels, ChannelPoolConfig::Dynamic(dynamic_config) => dynamic_config.initial_channels, @@ -482,13 +518,13 @@ async fn create_channel_pool( let pool = match pool_config { ChannelPoolConfig::Static(static_config) => { - ChannelPool::new_static(channels.clone(), static_config, config.clone()) + ChannelPool::new_static(channels, static_config, config.clone()) } ChannelPoolConfig::Dynamic(dynamic_config) => { - ChannelPool::new_dynamic(channels.clone(), dynamic_config, config.clone()) + ChannelPool::new_dynamic(channels, dynamic_config, config.clone()) } }; - Ok((pool, channels)) + Ok(pool) } #[cfg(feature = "metrics")] @@ -666,6 +702,26 @@ impl Spanner { crate::builder::DatabaseClientBuilder::new(self.clone(), database.into()) } + /// Returns the total number of active channels currently available in the gRPC channel pool. + /// + /// When configured with [`DynamicChannelPoolConfig`], this number can scale between the + /// configured minimum and maximum limits as concurrent in-flight RPC demand fluctuates. + /// + /// When configured with [`StaticChannelPoolConfig`], this returns the fixed number of channels + /// configured in the pool. + /// + /// # Example + /// ``` + /// # use google_cloud_spanner::client::Spanner; + /// # async fn sample() -> anyhow::Result<()> { + /// let client = Spanner::builder().build().await?; + /// let active_channels = client.active_channel_count(); + /// # Ok(()) } + /// ``` + pub fn active_channel_count(&self) -> usize { + self.channel_pool.active_channel_count() + } + /// Creates a new client from the provided stub. /// /// The most common case for calling this function is in tests mocking the @@ -682,14 +738,12 @@ impl Spanner { channel_id: 1, }; let channel_pool = ChannelPool::new_static( - vec![channel.clone()], + vec![channel], StaticChannelPoolConfig::new(1), ClientConfig::default(), ); Self { channel_pool, - channels: vec![channel], - counter: Arc::new(AtomicUsize::new(0)), config: ClientConfig::default(), is_emulator: false, instance_type: InstanceType::Cloud, @@ -731,7 +785,6 @@ impl Spanner { self.instance_type } - #[allow(dead_code)] pub(crate) fn channel_pool(&self) -> &ChannelPool { &self.channel_pool } @@ -742,22 +795,19 @@ impl Spanner { .expect("channel pool must have active channels") } - #[allow(dead_code)] + pub(crate) fn pick_channel_for_target(&self, target: &ChannelTarget<'_>) -> ChannelLease { + match target { + ChannelTarget::Affinity(affinity) => self.resolve_affinity(affinity), + ChannelTarget::Any => self.pick_channel(), + } + } + pub(crate) fn resolve_affinity(&self, affinity: &TransactionAffinity) -> ChannelLease { self.channel_pool .resolve_affinity(affinity) .expect("channel pool must have active channels") } - pub(crate) fn get_channel(&self, hint: usize) -> &Channel { - let idx = hint % self.channels.len(); - &self.channels[idx] - } - - pub(crate) fn next_channel_hint(&self) -> usize { - self.counter.fetch_add(1, Ordering::Relaxed) - } - pub(crate) fn attach_request_id( &self, mut options: RequestOptions, @@ -972,6 +1022,7 @@ mod tests { use google_cloud_gax::error::rpc::Code; use google_cloud_gax::retry_state::RetryState; use google_cloud_test_macros::tokio_test_no_panics; + use scoped_env::ScopedEnv; use serial_test::serial; use spanner_grpc_mock::google::rpc as mock_rpc; use spanner_grpc_mock::google::spanner::v1 as mock_v1; @@ -1002,12 +1053,6 @@ mod tests { assert_not_impl_any!(Spanner: RefUnwindSafe, UnwindSafe); } - impl Spanner { - fn channel_count(&self) -> usize { - self.channel_pool.active_channel_count() - } - } - #[tokio_test_no_panics] #[serial] async fn channel_pool_default_size() { @@ -1024,7 +1069,11 @@ mod tests { .expect("Failed to build client"); let expected_channels = if client.is_emulator() { 1 } else { 4 }; - assert_eq!(client.channel_count(), expected_channels); + assert_eq!( + client.active_channel_count(), + expected_channels, + "default channel pool size must match expected channels" + ); } #[test] @@ -1068,17 +1117,119 @@ mod tests { .await .expect("Failed to build client"); - let hint0 = client.next_channel_hint(); - let hint1 = client.next_channel_hint(); - let hint2 = client.next_channel_hint(); - let hint3 = client.next_channel_hint(); - let hint4 = client.next_channel_hint(); + let channel = client.pick_channel(); + assert!( + channel.channel_id >= 1 && channel.channel_id <= 4, + "picked channel_id should be in 1..=4, got {}", + channel.channel_id + ); + } + + #[tokio_test_no_panics] + #[serial] + async fn database_client_uses_channel_pool() { + use crate::channel_pool::DynamicChannelPoolConfig; + use crate::statement::Statement; + use tokio::sync::Notify; + + let query_started = Arc::new(Notify::new()); + let release_query = Arc::new(Notify::new()); + + let mut mock = MockSpanner::new(); + mock.expect_create_session().returning(|_| { + Ok(Response::new(mock_v1::Session { + name: + "projects/test-project/instances/test-instance/databases/test-db/sessions/123" + .to_string(), + ..Default::default() + })) + }); + + let query_started_clone = Arc::clone(&query_started); + let release_query_clone = Arc::clone(&release_query); + mock.expect_execute_streaming_sql().returning(move |_| { + let query_started = Arc::clone(&query_started_clone); + let release_query = Arc::clone(&release_query_clone); + let (sender, receiver) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + query_started.notify_one(); + release_query.notified().await; + let result_set = mock_v1::PartialResultSet { + metadata: Some(mock_v1::ResultSetMetadata { + row_type: Some(mock_v1::StructType { fields: vec![] }), + transaction: None, + undeclared_parameters: None, + }), + values: vec![], + chunked_value: false, + resume_token: vec![1, 2, 3], + stats: None, + precommit_token: None, + cache_update: None, + last: true, + }; + let _ = sender.send(Ok(result_set)).await; + }); + Ok(Response::new(receiver)) + }); + + let (address, _server) = start("0.0.0.0:0", mock) + .await + .expect("Failed to start mock server"); + + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(1) + .with_min_channels(1) + .with_max_channels(4) + .with_min_rpc_per_channel(0.5) + .with_max_rpc_per_channel(1.5) + .with_error_penalty_step(1); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("Failed to build spanner client"); - assert_eq!(hint0 % 4, 0); - assert_eq!(hint1 % 4, 1); - assert_eq!(hint2 % 4, 2); - assert_eq!(hint3 % 4, 3); - assert_eq!(hint4 % 4, 0); + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-db") + .build() + .await + .expect("Failed to build database client"); + + let query_handle = tokio::spawn(async move { + let transaction = database_client.single_use().build(); + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await + .expect("execute_query must succeed"); + + let row = result_set + .next() + .await + .transpose() + .expect("row must succeed"); + assert!(row.is_none(), "result set should be empty"); + }); + + query_started.notified().await; + + let in_flight_rpcs: u32 = spanner + .channel_pool() + .active_entries() + .iter() + .map(|entry| entry.in_flight()) + .sum(); + + release_query.notify_one(); + query_handle.await.expect("query task must finish cleanly"); + + assert_eq!( + in_flight_rpcs, 1, + "Query executed via DatabaseClient must register as an in-flight RPC in the channel pool" + ); } #[tokio_test_no_panics] @@ -1171,8 +1322,8 @@ mod tests { req.database = "projects/test-project/instances/test-instance/databases/test-db".to_string(); - let session = client - .get_channel(client.next_channel_hint()) + let lease = client.pick_channel(); + let session = lease .inner .create_session() .with_request(req) @@ -1754,13 +1905,6 @@ mod tests { })) }); - mock.expect_begin_transaction().returning(|_| { - Ok(Response::new(mock_v1::Transaction { - id: vec![42], - ..Default::default() - })) - }); - mock.expect_execute_streaming_sql().once().returning(|req| { let metadata = req.metadata(); let timeout = metadata.get("grpc-timeout"); @@ -1929,13 +2073,6 @@ mod tests { })) }); - mock.expect_begin_transaction().returning(|_| { - Ok(Response::new(mock_v1::Transaction { - id: vec![42], - ..Default::default() - })) - }); - // Mock ExecuteSql to first return RESOURCE_EXHAUSTED and then succeed. let mut seq = mockall::Sequence::new(); @@ -2059,13 +2196,6 @@ mod tests { })) }); - mock.expect_begin_transaction().returning(|_| { - Ok(Response::new(Transaction { - id: vec![1, 2, 3], - ..Default::default() - })) - }); - mock.expect_commit().once().returning(|_| { Ok(Response::new(CommitResponse { commit_timestamp: Some(prost_types::Timestamp { @@ -2357,14 +2487,12 @@ mod tests { .expect("Failed to build client"); assert_eq!( - client.channel_count(), + client.active_channel_count(), 4, "default pool size should be 4 channels" ); - // Test with a channel_hint that is larger than the pool size (e.g., hint = 7). - // get_channel(7) maps to channel at index (7 % 4 = 3), which has 1-based channel_id 4. - let channel = client.get_channel(7); + let channel = client.pick_channel(); let options = crate::RequestOptions::default(); let options = client.attach_request_id(options, channel.channel_id); let headers = options @@ -2376,11 +2504,10 @@ mod tests { .to_str() .expect("should be valid ASCII"); - // With 4 channels and hint = 7: (7 % 4) + 1 = 3 + 1 = 4. - // So the prefix should contain ".4." for channel ID 4. assert!( - val.contains(".4."), - "Request ID should contain channel ID 4 for hint 7 with pool size 4, got {val}" + val.contains(&format!(".{}.", channel.channel_id)), + "Request ID should contain channel ID {} with pool size 4, got {val}", + channel.channel_id ); } @@ -2398,7 +2525,7 @@ mod tests { .await .expect("Failed to build client"); - let channel = client.get_channel(0); + let channel = client.pick_channel(); let mut options = crate::RequestOptions::default(); options = client.attach_request_id(options, channel.channel_id); let first_headers = options @@ -2795,7 +2922,7 @@ mod tests { .expect("Failed to build client"); assert_eq!( - client.channel_count(), + client.active_channel_count(), 2, "Client should have exactly 2 channels configured" ); @@ -2837,7 +2964,7 @@ mod tests { .expect("Failed to build client"); assert_eq!( - client.channel_count(), + client.active_channel_count(), 3, "Client should have 3 initial channels configured" ); @@ -2910,7 +3037,7 @@ mod tests { let client = Spanner::from_stub(DummyStub); assert_eq!( - client.channel_count(), + client.active_channel_count(), 1, "Client from stub should have exactly 1 channel" ); @@ -2965,7 +3092,7 @@ mod tests { fn resolve_pool_config() { // Case 1: Default when no env var and no override let mut config = ClientConfig::default(); - let pool_config = resolve_pool_config_with(&mut config, false, || None) + let pool_config = resolve_pool_config_with(&mut config, false, || None, || None) .expect("default pool config should resolve"); assert_eq!( pool_config, @@ -2975,7 +3102,7 @@ mod tests { // Case 2: Emulator defaults to 1 channel when SPANNER_NUM_CHANNELS is not set let mut config = ClientConfig::default(); - let pool_config = resolve_pool_config_with(&mut config, true, || None) + let pool_config = resolve_pool_config_with(&mut config, true, || None, || None) .expect("emulator pool config should resolve to default 1 channel"); assert_eq!( pool_config, @@ -2985,8 +3112,9 @@ mod tests { // Case 2b: SPANNER_NUM_CHANNELS overrides emulator default let mut config = ClientConfig::default(); - let pool_config = resolve_pool_config_with(&mut config, true, || Some("8".to_string())) - .expect("emulator pool config with env var override should resolve"); + let pool_config = + resolve_pool_config_with(&mut config, true, || Some("8".to_string()), || None) + .expect("emulator pool config with env var override should resolve"); assert_eq!( pool_config, ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 8 }), @@ -2995,8 +3123,9 @@ mod tests { // Case 3: SPANNER_NUM_CHANNELS valid integer let mut config = ClientConfig::default(); - let pool_config = resolve_pool_config_with(&mut config, false, || Some("2".to_string())) - .expect("pool config with SPANNER_NUM_CHANNELS=2 should resolve"); + let pool_config = + resolve_pool_config_with(&mut config, false, || Some("2".to_string()), || None) + .expect("pool config with SPANNER_NUM_CHANNELS=2 should resolve"); assert_eq!( pool_config, ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 2 }), @@ -3005,9 +3134,13 @@ mod tests { // Case 4: SPANNER_NUM_CHANNELS unparsable integer string let mut config = ClientConfig::default(); - let error = - resolve_pool_config_with(&mut config, false, || Some("not_a_number".to_string())) - .expect_err("should fail when SPANNER_NUM_CHANNELS is not a valid integer"); + let error = resolve_pool_config_with( + &mut config, + false, + || Some("not_a_number".to_string()), + || None, + ) + .expect_err("should fail when SPANNER_NUM_CHANNELS is not a valid integer"); let debug_error = format!("{error:?}"); assert!( debug_error.contains("InvalidDigit"), @@ -3016,7 +3149,7 @@ mod tests { // Case 5: SPANNER_NUM_CHANNELS zero (validation failure) let mut config = ClientConfig::default(); - let error = resolve_pool_config_with(&mut config, false, || Some("0".to_string())) + let error = resolve_pool_config_with(&mut config, false, || Some("0".to_string()), || None) .expect_err("should fail when SPANNER_NUM_CHANNELS is 0"); let debug_error = format!("{error:?}"); assert!( @@ -3031,8 +3164,9 @@ mod tests { .insert(ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 10, })); - let pool_config = resolve_pool_config_with(&mut config, false, || Some("2".to_string())) - .expect("extension override should resolve"); + let pool_config = + resolve_pool_config_with(&mut config, false, || Some("2".to_string()), || None) + .expect("extension override should resolve"); assert_eq!( pool_config, ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 10 }), @@ -3046,8 +3180,9 @@ mod tests { .insert(ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 8, })); - let pool_config = resolve_pool_config_with(&mut config, true, || Some("2".to_string())) - .expect("extension override should resolve even on emulator"); + let pool_config = + resolve_pool_config_with(&mut config, true, || Some("2".to_string()), || None) + .expect("extension override should resolve even on emulator"); assert_eq!( pool_config, ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 8 }), @@ -3056,8 +3191,9 @@ mod tests { // Case 8: SPANNER_NUM_CHANNELS empty or whitespace string falls back to default let mut config = ClientConfig::default(); - let pool_config = resolve_pool_config_with(&mut config, false, || Some(" ".to_string())) - .expect("whitespace SPANNER_NUM_CHANNELS should resolve to default"); + let pool_config = + resolve_pool_config_with(&mut config, false, || Some(" ".to_string()), || None) + .expect("whitespace SPANNER_NUM_CHANNELS should resolve to default"); assert_eq!( pool_config, ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 4 }), @@ -3070,8 +3206,9 @@ mod tests { config .extensions .insert(ChannelPoolConfig::Dynamic(dynamic_config.clone())); - let pool_config = resolve_pool_config_with(&mut config, false, || Some("2".to_string())) - .expect("dynamic extension override should resolve"); + let pool_config = + resolve_pool_config_with(&mut config, false, || Some("2".to_string()), || None) + .expect("dynamic extension override should resolve"); assert_eq!( pool_config, ChannelPoolConfig::Dynamic(dynamic_config), @@ -3083,8 +3220,9 @@ mod tests { config .extensions .insert(StaticChannelPoolConfig { num_channels: 6 }); - let pool_config = resolve_pool_config_with(&mut config, false, || Some("2".to_string())) - .expect("static struct extension override should resolve"); + let pool_config = + resolve_pool_config_with(&mut config, false, || Some("2".to_string()), || None) + .expect("static struct extension override should resolve"); assert_eq!( pool_config, ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 6 }), @@ -3095,8 +3233,9 @@ mod tests { let mut config = ClientConfig::default(); let dynamic_config = DynamicChannelPoolConfig::default(); config.extensions.insert(dynamic_config.clone()); - let pool_config = resolve_pool_config_with(&mut config, false, || Some("2".to_string())) - .expect("dynamic struct extension override should resolve"); + let pool_config = + resolve_pool_config_with(&mut config, false, || Some("2".to_string()), || None) + .expect("dynamic struct extension override should resolve"); assert_eq!( pool_config, ChannelPoolConfig::Dynamic(dynamic_config), @@ -3108,8 +3247,9 @@ mod tests { config.extensions.insert(Arc::new(ChannelPoolConfig::Static( StaticChannelPoolConfig { num_channels: 7 }, ))); - let pool_config = resolve_pool_config_with(&mut config, false, || Some("2".to_string())) - .expect("Arc extension override should resolve"); + let pool_config = + resolve_pool_config_with(&mut config, false, || Some("2".to_string()), || None) + .expect("Arc extension override should resolve"); assert_eq!( pool_config, ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 7 }), @@ -3121,7 +3261,7 @@ mod tests { config .extensions .insert(StaticChannelPoolConfig { num_channels: 0 }); - let error = resolve_pool_config_with(&mut config, false, || None) + let error = resolve_pool_config_with(&mut config, false, || None, || None) .expect_err("should fail when StaticChannelPoolConfig in extensions is invalid"); let debug_error = format!("{error:?}"); assert!( @@ -3136,7 +3276,7 @@ mod tests { ..Default::default() }; config.extensions.insert(invalid_dynamic); - let error = resolve_pool_config_with(&mut config, false, || None) + let error = resolve_pool_config_with(&mut config, false, || None, || None) .expect_err("should fail when DynamicChannelPoolConfig in extensions is invalid"); let debug_error = format!("{error:?}"); assert!( @@ -3149,13 +3289,152 @@ mod tests { config.extensions.insert(Arc::new(ChannelPoolConfig::Static( StaticChannelPoolConfig { num_channels: 0 }, ))); - let error = resolve_pool_config_with(&mut config, false, || None) + let error = resolve_pool_config_with(&mut config, false, || None, || None) .expect_err("should fail when Arc in extensions is invalid"); let debug_error = format!("{error:?}"); assert!( debug_error.contains("num_channels must be at least 1"), "error should indicate invalid channels: {debug_error}" ); + + // Case 16: SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=true resolves to default dynamic pool + let mut config = ClientConfig::default(); + let pool_config = + resolve_pool_config_with(&mut config, false, || None, || Some("true".to_string())) + .expect("SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=true should resolve"); + assert_eq!( + pool_config, + ChannelPoolConfig::Dynamic(DynamicChannelPoolConfig::default()), + "SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=true should produce default dynamic pool" + ); + + // Case 17: SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL accepts truthy values ("1", "yes", "on", "t") + for truthy in ["1", "yes", "on", "t", "TRUE", "Yes "] { + let mut config = ClientConfig::default(); + let pool_config = + resolve_pool_config_with(&mut config, false, || None, || Some(truthy.to_string())) + .expect("truthy value should resolve"); + assert_eq!( + pool_config, + ChannelPoolConfig::Dynamic(DynamicChannelPoolConfig::default()), + "'{truthy}' must enable dynamic channel pool" + ); + } + + // Case 18: SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=false falls back to default static pool + let mut config = ClientConfig::default(); + let pool_config = + resolve_pool_config_with(&mut config, false, || None, || Some("false".to_string())) + .expect("SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=false should resolve"); + assert_eq!( + pool_config, + ChannelPoolConfig::Static(StaticChannelPoolConfig::default()), + "SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=false should fall back to static pool" + ); + + // Case 19: Programmatic static configuration takes precedence over SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=true + let mut config = ClientConfig::default(); + config.extensions.insert(StaticChannelPoolConfig::new(2)); + let pool_config = + resolve_pool_config_with(&mut config, false, || None, || Some("true".to_string())) + .expect("programmatic override should resolve"); + assert_eq!( + pool_config, + ChannelPoolConfig::Static(StaticChannelPoolConfig { num_channels: 2 }), + "programmatic static config takes precedence over SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL" + ); + + // Case 20: Programmatic dynamic configuration takes precedence over SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=false + let mut config = ClientConfig::default(); + let custom_dynamic = DynamicChannelPoolConfig::new().with_min_channels(2); + config.extensions.insert(custom_dynamic.clone()); + let pool_config = + resolve_pool_config_with(&mut config, false, || None, || Some("false".to_string())) + .expect("programmatic dynamic override should resolve"); + assert_eq!( + pool_config, + ChannelPoolConfig::Dynamic(custom_dynamic), + "programmatic dynamic config takes precedence over SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL" + ); + + // Case 21: SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=true with SPANNER_NUM_CHANNELS=8 + let mut config = ClientConfig::default(); + let pool_config = resolve_pool_config_with( + &mut config, + false, + || Some("8".to_string()), + || Some("true".to_string()), + ) + .expect("dynamic pool with custom channels should resolve"); + assert_eq!( + pool_config, + ChannelPoolConfig::Dynamic( + DynamicChannelPoolConfig::default() + .with_initial_channels(8) + .with_min_channels(8) + ), + "SPANNER_NUM_CHANNELS sets initial and min channels when dynamic pool is enabled" + ); + + // Case 22: SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=true with SPANNER_NUM_CHANNELS > max_channels (e.g. 300) fails validation + let mut config = ClientConfig::default(); + let error = resolve_pool_config_with( + &mut config, + false, + || Some("300".to_string()), + || Some("true".to_string()), + ) + .expect_err("should fail when SPANNER_NUM_CHANNELS exceeds maximum supported limit"); + let debug_error = format!("{error:?}"); + assert!( + debug_error.contains("max_channels cannot exceed maximum supported limit"), + "error should indicate max_channels limit exceeded: {debug_error}" + ); + + // Case 23: SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=true with invalid SPANNER_NUM_CHANNELS format + let mut config = ClientConfig::default(); + let error = resolve_pool_config_with( + &mut config, + false, + || Some("invalid_number".to_string()), + || Some("true".to_string()), + ) + .expect_err("should fail when SPANNER_NUM_CHANNELS is not a valid integer"); + let debug_error = format!("{error:?}"); + assert!( + debug_error.contains("InvalidDigit"), + "error should indicate invalid integer format: {debug_error}" + ); + + // Case 24: SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=true with whitespace SPANNER_NUM_CHANNELS + let mut config = ClientConfig::default(); + let pool_config = resolve_pool_config_with( + &mut config, + false, + || Some(" ".to_string()), + || Some("true".to_string()), + ) + .expect("whitespace SPANNER_NUM_CHANNELS should be ignored"); + assert_eq!( + pool_config, + ChannelPoolConfig::Dynamic(DynamicChannelPoolConfig::default()), + "whitespace SPANNER_NUM_CHANNELS must fall back to default dynamic pool" + ); + + // Case 25: SPANNER_ENABLE_DYNAMIC_CHANNEL_POOL=true with SPANNER_NUM_CHANNELS=0 fails validation + let mut config = ClientConfig::default(); + let error = resolve_pool_config_with( + &mut config, + false, + || Some("0".to_string()), + || Some("true".to_string()), + ) + .expect_err("should fail when dynamic channels configured as 0"); + let debug_error = format!("{error:?}"); + assert!( + debug_error.contains("min_channels must be at least 1"), + "error should indicate invalid channels for dynamic pool: {debug_error}" + ); } #[tokio_test_no_panics] @@ -3248,4 +3527,80 @@ mod tests { "Channel pool must have 3 channels from Arc extension" ); } + + #[test] + fn parse_timeout_units() { + let mut metadata = MetadataMap::new(); + metadata.insert( + "grpc-timeout", + "500m".parse().expect("valid grpc-timeout header value"), + ); + assert_eq!( + parse_timeout(&metadata), + 500_000, + "parse_timeout for 500m must equal 500000" + ); + + metadata.insert( + "grpc-timeout", + "500000n".parse().expect("valid grpc-timeout header value"), + ); + assert_eq!( + parse_timeout(&metadata), + 500, + "parse_timeout for 500000n must equal 500" + ); + } + + #[test] + #[serial] + fn emulator_detection_from_env() { + let _environment_guard = ScopedEnv::set("SPANNER_EMULATOR_HOST", "localhost:9010"); + let mut client_configuration = ClientConfig::default(); + let is_emulator = detect_and_configure_emulator(&mut client_configuration); + assert!( + is_emulator, + "must detect emulator from SPANNER_EMULATOR_HOST environment variable" + ); + } + + #[tokio::test] + async fn admin_builders_configuration() { + let mock = MockSpanner::new(); + let (address, _server) = start("0.0.0.0:0", mock) + .await + .expect("Failed to start mock server"); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .build() + .await + .expect("build client must succeed"); + + let _database_admin_builder = spanner.database_admin_builder(); + let _instance_admin_builder = spanner.instance_admin_builder(); + } + + #[tokio::test] + async fn channel_create_with_tracing() { + let mock = MockSpanner::new(); + let (address, _server) = start("0.0.0.0:0", mock) + .await + .expect("Failed to start mock server"); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_tracing() + .build() + .await + .expect("build client with tracing must succeed"); + + assert_eq!( + spanner.channel_pool.active_channel_count(), + 4, + "channel pool should have 4 channels by default" + ); + } } diff --git a/src/spanner/src/database_client.rs b/src/spanner/src/database_client.rs index ada1ab41b4..4b632bf009 100644 --- a/src/spanner/src/database_client.rs +++ b/src/spanner/src/database_client.rs @@ -14,6 +14,7 @@ use crate::batch_read_only_transaction::BatchReadOnlyTransactionBuilder; use crate::batch_write_transaction::BatchWriteTransactionBuilder; +use crate::channel_pool::ChannelTarget; use crate::client::Spanner; use crate::model::transaction_options::Mode; use crate::model::transaction_options::read_only::TimestampBound; @@ -47,7 +48,7 @@ use crate::routing::latency_registry::LatencyRegistry; use crate::routing::location_router::{LocationRouter, RoutingContext}; use crate::routing::server_connection::ServerConnection; use crate::server_streaming::builder::{BatchWrite, ExecuteStreamingSql, StreamingRead}; -use crate::server_streaming::stream::TransactionIdCallback; +use crate::server_streaming::stream::{StreamLifetimeGuard, TransactionIdCallback}; use crate::session_maintainer::ManagedSessionMaintainer; use crate::transaction_runner::TransactionRunnerBuilder; use crate::write_only_transaction::WriteOnlyTransactionBuilder; @@ -99,16 +100,22 @@ macro_rules! define_db_rpc { $pre_route:path, $post_hook:path ) => { - pub(crate) async fn $method( + pub(crate) async fn $method<'a>( &self, mut request: $request_type, options: RequestOptions, - channel_hint: usize, + channel_target: impl Into>, ) -> Result<$response_type> { + let channel_target = channel_target.into(); let (connection, routing_context) = $pre_route(self, &mut request); + let channel_lease = connection + .is_none() + .then(|| self.spanner.pick_channel_for_target(&channel_target)); let channel = match &connection { Some(connection) => connection.channel(), - None => self.spanner.get_channel(channel_hint), + None => channel_lease + .as_deref() + .expect("channel lease must be present when connection is absent"), }; let _request_guard = connection .as_ref() @@ -119,6 +126,9 @@ macro_rules! define_db_rpc { .spanner .$method(request, options, channel, &self.o11y) .await; + if let Some(lease) = &channel_lease { + lease.record_call_result(&result); + } let latency = start.elapsed(); self.record_routing_feedback(connection.as_ref(), group_uid, latency, &result); $post_hook(self, routing_context, connection.as_ref(), &result); @@ -131,23 +141,27 @@ macro_rules! define_db_rpc { macro_rules! define_db_streaming_rpc { ($method:ident, $expect_method:ident, $request_type:ty, $builder_type:ty) => { - pub(crate) fn $method( + pub(crate) fn $method<'a>( &self, request: $request_type, options: RequestOptions, - channel_hint: usize, + channel_target: impl Into>, ) -> $builder_type { - let channel = self.spanner.get_channel(channel_hint); - self.spanner.$method(request, options, channel) + let channel_target = channel_target.into(); + let lease = self.spanner.pick_channel_for_target(&channel_target); + let builder = self.spanner.$method(request, options, &lease); + let lifetime_guard: StreamLifetimeGuard = Arc::new(lease.guard); + builder.with_lifetime_guard(lifetime_guard) } }; ($method:ident, $expect_method:ident, $request_type:ty, $builder_type:ty, $extract_key:expr) => { - pub(crate) fn $method( + pub(crate) fn $method<'a>( &self, mut request: $request_type, options: RequestOptions, - channel_hint: usize, + channel_target: impl Into>, ) -> $builder_type { + let channel_target = channel_target.into(); let is_read_write_begin = is_read_write_begin(request.transaction.as_ref()); // Step 1: When location-aware routing is disabled (standard Cloud Spanner), // `self.location_routing` is `None` so `$extract_key` is skipped immediately. @@ -168,18 +182,27 @@ macro_rules! define_db_streaming_rpc { // Step 3: Select the gRPC channel: // - If location-aware routing resolved a direct node connection (`Some(connection)`), use `connection.channel()`. - // - Otherwise (location routing disabled, unkeyed query/read, or cold cache), fall back to round-robin - // load-balancing across the client's channel pool via `self.spanner.get_channel(channel_hint)`. - // This fallback is a fast O(1) slice index without any heap allocation, cloning, or lock acquisition. - let channel = match &connection { - Some(connection) => connection.channel(), - None => self.spanner.get_channel(channel_hint), + // - Otherwise, lease a channel from the client's channel pool via `self.spanner.pick_channel_for_target(&channel_target)` + // and attach its active RPC lifetime guard to track in-flight concurrency. + let (builder, lifetime_guard) = match connection.as_ref() { + Some(connection) => ( + self.spanner.$method(request, options, connection.channel()), + None, + ), + None => { + let lease = self.spanner.pick_channel_for_target(&channel_target); + let builder = self.spanner.$method(request, options, &lease); + let guard: StreamLifetimeGuard = Arc::new(lease.guard); + (builder, Some(guard)) + } }; let callback = self.streaming_transaction_id_callback(is_read_write_begin, connection.as_ref()); - self.spanner - .$method(request, options, channel) - .with_transaction_id_callback(callback) + let mut builder = builder.with_transaction_id_callback(callback); + if let Some(guard) = lifetime_guard { + builder = builder.with_lifetime_guard(guard); + } + builder } }; } @@ -279,17 +302,17 @@ impl DatabaseClient { self.spanner.is_emulator() } - pub(crate) fn next_channel_hint(&self) -> usize { - self.spanner.next_channel_hint() - } - - pub(crate) fn attach_request_id( + pub(crate) fn attach_request_id<'a>( &self, options: RequestOptions, - channel_hint: usize, + target: impl Into>, ) -> RequestOptions { - let channel = self.spanner.get_channel(channel_hint); - self.spanner.attach_request_id(options, channel.channel_id) + let target = target.into(); + let channel_id = self + .spanner + .channel_pool() + .logical_channel_id_for_target(&target); + self.spanner.attach_request_id(options, channel_id) } for_all_unary_db_rpcs!(define_db_rpc); @@ -516,6 +539,28 @@ impl DatabaseClient { BatchWriteTransactionBuilder::new(self.clone()) } + /// Returns the total number of active channels currently available in the underlying gRPC channel pool. + /// + /// When configured with [`DynamicChannelPoolConfig`](crate::channel_pool::DynamicChannelPoolConfig), + /// this number can scale between the configured minimum and maximum limits as concurrent + /// RPC demand fluctuates. + /// + /// When configured with [`StaticChannelPoolConfig`](crate::channel_pool::StaticChannelPoolConfig), + /// this returns the fixed number of channels configured in the pool. + /// + /// # Example + /// ``` + /// # use google_cloud_spanner::client::Spanner; + /// # async fn sample() -> anyhow::Result<()> { + /// let spanner = Spanner::builder().build().await?; + /// let db = spanner.database_client("projects/p/instances/i/databases/d").build().await?; + /// let active_channels = db.active_channel_count(); + /// # Ok(()) } + /// ``` + pub fn active_channel_count(&self) -> usize { + self.spanner.active_channel_count() + } + pub(crate) fn session_name(&self) -> String { self.session_maintainer.session_name() } @@ -1323,11 +1368,7 @@ impl LocationRoutingState { .unwrap_or(DEFAULT_ENDPOINT) .to_string(); - let default_channel = spanner - .channels - .first() - .cloned() - .expect("Spanner client must have at least one channel"); + let default_channel = (*spanner.pick_channel()).clone(); let default_connection = ServerConnection::new_default(default_endpoint, default_channel); let connection_cache = Arc::new(ConnectionCache::new(default_connection)); @@ -2481,7 +2522,7 @@ mod tests { assert_eq!(hit_connection.address(), node_address); // 4. Mark node on cooldown: falls back to default gateway (resolves to None so the - // request dispatches over the client's channel pool via channel_hint) + // request dispatches over the client's channel pool) router.cooldown_tracker().record_failure(node_address); let fallback_connection = db_client.resolve_routing_connection(&hit_context); assert!( @@ -3061,27 +3102,29 @@ mod tests { } #[tokio_test_no_panics] - async fn transaction_affinity_gateway_fallback_preserves_channel_hint_across_statements() { + async fn transaction_affinity_gateway_fallback_preserves_channel_across_statements() { + use crate::channel_pool::TransactionAffinity; use std::sync::Mutex; let captured_requests = Arc::new(Mutex::new(Vec::new())); let mut mock = create_test_mock(); - fn extract_request_id(metadata: &MetadataMap) -> String { + fn extract_channel_id(metadata: &MetadataMap) -> String { metadata .get("x-goog-spanner-request-id") .and_then(|id| id.to_str().ok()) + .and_then(|header| header.split('.').nth(3)) .unwrap_or_default() .to_string() } let captured_clone = Arc::clone(&captured_requests); mock.expect_begin_transaction().returning(move |request| { - let request_id = extract_request_id(request.metadata()); + let channel_id = extract_channel_id(request.metadata()); captured_clone .lock() .expect("lock should succeed") - .push(("begin_transaction", request_id)); + .push(("begin_transaction", channel_id)); Ok(Response::new(mock_v1::Transaction { id: b"tx-gw-session-1".to_vec(), @@ -3091,11 +3134,11 @@ mod tests { let captured_clone = Arc::clone(&captured_requests); mock.expect_execute_sql().returning(move |request| { - let request_id = extract_request_id(request.metadata()); + let channel_id = extract_channel_id(request.metadata()); captured_clone .lock() .expect("lock should succeed") - .push(("execute_sql", request_id)); + .push(("execute_sql", channel_id)); Ok(Response::new(mock_v1::ResultSet::default())) }); @@ -3103,22 +3146,22 @@ mod tests { let captured_clone = Arc::clone(&captured_requests); mock.expect_execute_streaming_sql() .returning(move |request| { - let request_id = extract_request_id(request.metadata()); + let channel_id = extract_channel_id(request.metadata()); captured_clone .lock() .expect("lock should succeed") - .push(("execute_streaming_sql", request_id)); + .push(("execute_streaming_sql", channel_id)); Ok(Response::from(adapt([]))) }); let captured_clone = Arc::clone(&captured_requests); mock.expect_commit().returning(move |request| { - let request_id = extract_request_id(request.metadata()); + let channel_id = extract_channel_id(request.metadata()); captured_clone .lock() .expect("lock should succeed") - .push(("commit", request_id)); + .push(("commit", channel_id)); Ok(Response::new(mock_v1::CommitResponse::default())) }); @@ -3143,18 +3186,16 @@ mod tests { assert!(database_client.is_location_aware_routing_enabled()); - // Verify across multiple channel affinities (channel_hint 2 -> slot .3., channel_hint 1 -> slot .2.): - // 1. Statements within a transaction remain pinned to the same channel slot. - // 2. Different transactions distribute across distinct channel slots rather than - // collapsing onto Channel 0 (slot .1.). - for channel_hint in [2usize, 1usize] { - let expected_channel_id = format!(".{}.", channel_hint + 1); + // Verify across multiple transaction affinities: + // Statements within a transaction remain pinned to the same channel slot. + for _ in 0..2 { + let affinity = Arc::new(TransactionAffinity::new_read_write()); // 1. BeginTransaction (unkeyed read-write options) let begin_request = BeginTransactionRequest::default() .set_options(TransactionOptions::default().set_read_write(ReadWrite::default())); let transaction = database_client - .begin_transaction(begin_request, RequestOptions::default(), channel_hint) + .begin_transaction(begin_request, RequestOptions::default(), &affinity) .await .expect("begin_transaction should succeed"); @@ -3168,21 +3209,21 @@ mod tests { let selector = TransactionSelector::new().set_id(transaction_id.clone()); let sql_request = ExecuteSqlRequest::default().set_transaction(selector.clone()); database_client - .execute_sql(sql_request, RequestOptions::default(), channel_hint) + .execute_sql(sql_request, RequestOptions::default(), &affinity) .await .expect("execute_sql should succeed"); // 3. ExecuteStreamingSql with the returned transaction ID let streaming_request = ExecuteSqlRequest::default().set_transaction(selector); let _ = database_client - .execute_streaming_sql(streaming_request, RequestOptions::default(), channel_hint) + .execute_streaming_sql(streaming_request, RequestOptions::default(), &affinity) .send() .await; // 4. Commit with the transaction ID let commit_request = CommitRequest::default().set_transaction_id(transaction_id); database_client - .commit(commit_request, RequestOptions::default(), channel_hint) + .commit(commit_request, RequestOptions::default(), &affinity) .await .expect("commit should succeed"); @@ -3200,17 +3241,22 @@ mod tests { 4, "expected 4 calls: begin_transaction, execute_sql, execute_streaming_sql, commit" ); - for (rpc_name, request_id) in calls { - assert!( - request_id.contains(&expected_channel_id), - "RPC {rpc_name} for transaction with channel_hint {channel_hint} must route via channel {expected_channel_id}, got {request_id}" + let pinned_channel_id = calls[0].1.clone(); + assert!( + !pinned_channel_id.is_empty(), + "pinned channel ID must not be empty" + ); + for (rpc_name, channel_id) in calls { + assert_eq!( + channel_id, pinned_channel_id, + "RPC {rpc_name} for transaction must route via the pinned channel {pinned_channel_id}, got {channel_id}" ); } } } #[tokio_test_no_panics] - async fn channel_pool_round_robin_for_all_rpcs_when_location_routing_disabled() { + async fn channel_pool_routes_all_rpcs_when_location_routing_disabled() { use std::sync::Mutex; let captured_requests = Arc::new(Mutex::new(Vec::new())); @@ -3296,55 +3342,46 @@ mod tests { assert!(!db_client.is_location_aware_routing_enabled()); - // Verify round-robin channel distribution 1..=4 for all mapped RPCs across 4 hints - for channel_hint in 0..4 { - let expected_channel_id = format!(".{}.", channel_hint + 1); - - macro_rules! call_unary_rpc { - ($method:ident, $expect_method:ident, $request_type:ident, $response_type:ty $(, $extra:expr)*) => { - let _ = db_client - .$method( - $request_type::default(), - RequestOptions::default(), - channel_hint, - ) - .await; - }; - } - for_all_unary_db_rpcs!(call_unary_rpc); - - macro_rules! call_streaming_rpc { - ($method:ident, $expect_method:ident, $request_type:ident, $builder_type:ident $(, $extract_key:expr)?) => { - let _ = db_client - .$method( - $request_type::default(), - RequestOptions::default(), - channel_hint, - ) - .send() - .await; - }; - } - for_all_streaming_db_rpcs!(call_streaming_rpc); + macro_rules! call_unary_rpc { + ($method:ident, $expect_method:ident, $request_type:ident, $response_type:ty $(, $extra:expr)*) => { + let _ = db_client + .$method( + $request_type::default(), + RequestOptions::default(), + ChannelTarget::Any, + ) + .await; + }; + } + for_all_unary_db_rpcs!(call_unary_rpc); - let calls = captured_requests.lock().expect("lock").clone(); - captured_requests.lock().expect("lock").clear(); - assert_eq!( - calls.len(), - 10, - "each RPC method must be called once per hint" + macro_rules! call_streaming_rpc { + ($method:ident, $expect_method:ident, $request_type:ident, $builder_type:ident $(, $extract_key:expr)?) => { + let _ = db_client + .$method( + $request_type::default(), + RequestOptions::default(), + ChannelTarget::Any, + ) + .send() + .await; + }; + } + for_all_streaming_db_rpcs!(call_streaming_rpc); + + let calls = captured_requests.lock().expect("lock").clone(); + assert_eq!(calls.len(), 10, "each RPC method must be called once"); + for (rpc_name, request_id) in calls { + assert!( + !request_id.is_empty(), + "RPC {rpc_name} must use a valid channel ID from the pool, got {request_id}" ); - for (rpc_name, request_id) in calls { - assert!( - request_id.contains(&expected_channel_id), - "RPC {rpc_name} with channel_hint {channel_hint} must use channel ID {expected_channel_id}, got {request_id}" - ); - } } } #[tokio_test_no_panics] - async fn streaming_rpcs_round_robin_when_location_routing_enabled_without_routing_key() { + async fn streaming_rpcs_route_via_channel_pool_when_location_routing_enabled_without_routing_key() + { use std::sync::Mutex; let captured_requests = Arc::new(Mutex::new(Vec::new())); @@ -3392,37 +3429,32 @@ mod tests { assert!(db_client.is_location_aware_routing_enabled()); - for channel_hint in 0..4 { - let expected_channel_id = format!(".{}.", channel_hint + 1); - - // execute_streaming_sql (no routing key) - let _ = db_client - .execute_streaming_sql( - ExecuteSqlRequest::default(), - RequestOptions::default(), - channel_hint, - ) - .send() - .await; + // execute_streaming_sql (no routing key) + let _ = db_client + .execute_streaming_sql( + ExecuteSqlRequest::default(), + RequestOptions::default(), + ChannelTarget::Any, + ) + .send() + .await; - // streaming_read with KeySet::all() (no routing key) - let mut key_set = KeySet::new(); - key_set.all = true; - let read_request = ReadRequest::new().set_table("Users").set_key_set(key_set); - let _ = db_client - .streaming_read(read_request, RequestOptions::default(), channel_hint) - .send() - .await; + // streaming_read with KeySet::all() (no routing key) + let mut key_set = KeySet::new(); + key_set.all = true; + let read_request = ReadRequest::new().set_table("Users").set_key_set(key_set); + let _ = db_client + .streaming_read(read_request, RequestOptions::default(), ChannelTarget::Any) + .send() + .await; - let calls = captured_requests.lock().expect("lock").clone(); - captured_requests.lock().expect("lock").clear(); - assert_eq!(calls.len(), 2); - for (rpc_name, request_id) in calls { - assert!( - request_id.contains(&expected_channel_id), - "Even when location routing is enabled, {rpc_name} without routing key must round-robin onto channel {expected_channel_id}, got {request_id}" - ); - } + let calls = captured_requests.lock().expect("lock").clone(); + assert_eq!(calls.len(), 2, "must capture exactly two RPC calls"); + for (rpc_name, request_id) in calls { + assert!( + !request_id.is_empty(), + "Even when location routing is enabled, {rpc_name} without routing key must route via channel pool with request id, got {request_id}" + ); } } @@ -3488,7 +3520,11 @@ mod tests { // 1. Cold start: routes to gateway, attaches discovery routing hint with operation_uid let _ = database_client - .execute_streaming_sql(request.clone(), RequestOptions::default(), 0) + .execute_streaming_sql( + request.clone(), + RequestOptions::default(), + ChannelTarget::Any, + ) .send() .await; @@ -3563,7 +3599,7 @@ mod tests { // 4. Cache hit: routes directly to tablet mock with full routing hint let _ = database_client - .execute_streaming_sql(request, RequestOptions::default(), 0) + .execute_streaming_sql(request, RequestOptions::default(), ChannelTarget::Any) .send() .await; @@ -3664,7 +3700,11 @@ mod tests { // 1. Cold start: without cached recipe or range, routes to gateway and attaches bootstrap routing hint let _ = database_client - .streaming_read(read_request.clone(), RequestOptions::default(), 0) + .streaming_read( + read_request.clone(), + RequestOptions::default(), + ChannelTarget::Any, + ) .send() .await; @@ -3737,7 +3777,7 @@ mod tests { // 4. Cache hit: routes directly to tablet mock with full routing hint let _ = database_client - .streaming_read(read_request, RequestOptions::default(), 0) + .streaming_read(read_request, RequestOptions::default(), ChannelTarget::Any) .send() .await; @@ -4042,7 +4082,11 @@ mod tests { .set_key_set(key_set.clone()); let _ = database_client - .streaming_read(read_request.clone(), RequestOptions::default(), 0) + .streaming_read( + read_request.clone(), + RequestOptions::default(), + ChannelTarget::Any, + ) .send() .await; @@ -4099,7 +4143,7 @@ mod tests { // 3. Subsequent streaming read for table Users -> RoutingHint must be populated and attached let _ = database_client - .streaming_read(read_request, RequestOptions::default(), 0) + .streaming_read(read_request, RequestOptions::default(), ChannelTarget::Any) .send() .await; @@ -5261,7 +5305,7 @@ mod tests { .set_options(TransactionOptions::new().set_read_write(ReadWrite::new())) .set_mutation_key(user_mutation.clone().build_proto()); let _ = database_client - .begin_transaction(begin_request, RequestOptions::default(), 0) + .begin_transaction(begin_request, RequestOptions::default(), ChannelTarget::Any) .await .expect("begin_transaction must succeed"); @@ -5297,7 +5341,11 @@ mod tests { .set_transaction_id(Bytes::from_static(b"tx-e2e-1")) .set_mutations(vec![user_mutation.clone().build_proto()]); let _ = database_client - .commit(commit_request, RequestOptions::default(), 0) + .commit( + commit_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await .expect("commit must succeed"); @@ -5332,7 +5380,11 @@ mod tests { .set_single_use_transaction(TransactionOptions::new().set_read_write(ReadWrite::new())) .set_mutations(vec![user_mutation.clone().build_proto()]); let _ = database_client - .commit(single_use_commit_request, RequestOptions::default(), 0) + .commit( + single_use_commit_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await .expect("single-use commit must succeed"); @@ -5362,7 +5414,11 @@ mod tests { .set_session("projects/p/instances/i/databases/d/sessions/s1") .set_options(TransactionOptions::new().set_read_write(ReadWrite::new())); let unkeyed_response = database_client - .begin_transaction(unkeyed_begin_request, RequestOptions::default(), 0) + .begin_transaction( + unkeyed_begin_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await .expect("unkeyed begin_transaction must succeed"); diff --git a/src/spanner/src/lib.rs b/src/spanner/src/lib.rs index a669334b9f..1f835f2d3b 100644 --- a/src/spanner/src/lib.rs +++ b/src/spanner/src/lib.rs @@ -22,6 +22,8 @@ // Public domain modules. +/// Configuration types for the gRPC channel pool. +pub mod channel_pool; /// Key and key range definition types. pub mod key; /// Write mutations and transaction commit binders. @@ -81,7 +83,6 @@ pub mod stub { pub(crate) mod batch_dml; pub(crate) mod batch_read_only_transaction; pub(crate) mod batch_write_transaction; -pub(crate) mod channel_pool; pub(crate) mod database_client; pub(crate) mod from_value; pub(crate) mod observability; diff --git a/src/spanner/src/partitioned_dml_transaction.rs b/src/spanner/src/partitioned_dml_transaction.rs index 05ebafc971..78082f6048 100644 --- a/src/spanner/src/partitioned_dml_transaction.rs +++ b/src/spanner/src/partitioned_dml_transaction.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::channel_pool::TransactionAffinity; +use crate::channel_pool::{ChannelTarget, TransactionAffinity}; use crate::client::amend_request_options_for_lar; use crate::database_client::DatabaseClient; use crate::google::spanner::v1::result_set_stats::RowCount::RowCountLowerBound; @@ -181,7 +181,6 @@ impl PartitionedDmlTransaction { ..Default::default() }; let base_request = statement.into_request(); - let channel_hint = self.client.next_channel_hint(); let client = self.client; let is_emulator = client.is_emulator(); @@ -193,9 +192,10 @@ impl PartitionedDmlTransaction { let client = client.clone(); async move { - let _affinity = Arc::new(TransactionAffinity::new_read_write()); + let affinity = Arc::new(TransactionAffinity::new_read_write()); + let target = ChannelTarget::Affinity(&affinity); let transaction = client - .begin_transaction(begin_request, gax_options.clone(), channel_hint) + .begin_transaction(begin_request, gax_options.clone(), target) .await?; let execute_request = @@ -209,7 +209,7 @@ impl PartitionedDmlTransaction { }); let stream_builder = - client.execute_streaming_sql(execute_request, gax_options, channel_hint); + client.execute_streaming_sql(execute_request, gax_options, target); let stream = stream_builder.send().await?; extract_lower_bound_update_count_from_stream(stream, &client).await diff --git a/src/spanner/src/read_only_transaction.rs b/src/spanner/src/read_only_transaction.rs index e6a8ae7716..4151c1c53b 100644 --- a/src/spanner/src/read_only_transaction.rs +++ b/src/spanner/src/read_only_transaction.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::channel_pool::TransactionAffinity; +use crate::channel_pool::{ChannelTarget, TransactionAffinity}; use crate::database_client::DatabaseClient; use crate::error::internal_error; use crate::model::TransactionOptions; @@ -101,7 +101,6 @@ impl SingleUseReadOnlyTransactionBuilder { .set_single_use(TransactionOptions::default().set_read_only(read_only)); let session_name = self.client.session_name(); - let channel_hint = self.client.next_channel_hint(); SingleUseReadOnlyTransaction { context: ReadContext { session_name, @@ -112,7 +111,6 @@ impl SingleUseReadOnlyTransactionBuilder { ), precommit_token_tracker: PrecommitTokenTracker::new_noop(), transaction_tag: None, - channel_hint, begin_transaction_request_options: None, affinity: None, }, @@ -401,15 +399,19 @@ impl MultiUseReadOnlyTransactionBuilder { let options = TransactionOptions::default().set_read_only(read_only); let session_name = self.client.session_name(); - let channel_hint = self.client.next_channel_hint(); + let affinity = self + .affinity + .unwrap_or_else(|| Arc::new(TransactionAffinity::new_read_only())); + let selector = match self.begin_transaction_option { BeginTransactionOption::ExplicitBegin => { + let target = ChannelTarget::from(&affinity); let response = execute_begin_transaction( &self.client, session_name.clone(), options, None, - channel_hint, + target, self.begin_gax_options.clone().unwrap_or_default(), None, ) @@ -425,11 +427,6 @@ impl MultiUseReadOnlyTransactionBuilder { )), }; - let affinity = Some( - self.affinity - .unwrap_or_else(|| Arc::new(TransactionAffinity::new_read_only())), - ); - Ok(MultiUseReadOnlyTransaction { context: ReadContext { session_name, @@ -437,9 +434,8 @@ impl MultiUseReadOnlyTransactionBuilder { transaction_selector: selector, precommit_token_tracker: PrecommitTokenTracker::new_noop(), transaction_tag: None, - channel_hint, begin_transaction_request_options: self.begin_gax_options, - affinity, + affinity: Some(affinity), }, }) } @@ -556,12 +552,12 @@ impl MultiUseReadOnlyTransaction { } /// Executes an explicit `BeginTransaction` RPC on Spanner. -pub(crate) async fn execute_begin_transaction( - client: &crate::database_client::DatabaseClient, +pub(crate) async fn execute_begin_transaction<'a>( + client: &DatabaseClient, session_name: String, - options: crate::model::TransactionOptions, + options: TransactionOptions, transaction_tag: Option, - channel_hint: usize, + channel_target: impl Into>, request_options: crate::RequestOptions, mutation_key: Option, ) -> crate::Result { @@ -575,7 +571,7 @@ pub(crate) async fn execute_begin_transaction( } client - .begin_transaction(request, request_options, channel_hint) + .begin_transaction(request, request_options, channel_target) .await } @@ -677,12 +673,10 @@ pub(crate) struct ExplicitBeginParams { pub(crate) client: crate::database_client::DatabaseClient, pub(crate) session_name: String, pub(crate) transaction_tag: Option, - pub(crate) channel_hint: usize, pub(crate) request_options: crate::RequestOptions, pub(crate) is_stream_fallback: bool, pub(crate) precommit_token_tracker: crate::precommit::PrecommitTokenTracker, pub(crate) mutation_key: Option, - #[allow(dead_code)] pub(crate) affinity: Option>, } @@ -752,12 +746,13 @@ impl ReadContextTransactionSelector { // Only the leader thread will reach this point to perform the explicit begin. // Waiters are blocked in `poll_selector_status` waiting for the result, // and already completed states return early above. + let target = ChannelTarget::from(params.affinity.as_deref()); let response = match execute_begin_transaction( ¶ms.client, params.session_name, options, params.transaction_tag, - params.channel_hint, + target, params.request_options, params.mutation_key, ) @@ -983,7 +978,6 @@ pub(crate) struct ReadContext { pub(crate) transaction_selector: ReadContextTransactionSelector, pub(crate) precommit_token_tracker: PrecommitTokenTracker, pub(crate) transaction_tag: Option, - pub(crate) channel_hint: usize, pub(crate) begin_transaction_request_options: Option, pub(crate) affinity: Option>, } @@ -1034,7 +1028,6 @@ impl ReadContext { client: self.client.clone(), session_name: self.session_name.clone(), transaction_tag: self.transaction_tag.clone(), - channel_hint: self.channel_hint, request_options: options, is_stream_fallback, precommit_token_tracker: self.precommit_token_tracker.clone(), @@ -1046,7 +1039,6 @@ impl ReadContext { } /// Returns a reference to the transaction channel affinity handle, if set. - #[allow(dead_code)] pub(crate) fn affinity(&self) -> Option<&TransactionAffinity> { self.affinity.as_deref() } @@ -1085,9 +1077,10 @@ macro_rules! execute_stream_with_retry { ($self:expr, $request:ident, $gax_options:ident, $rpc_method:ident, $operation_variant:path, $method_name:expr) => {{ let operation_start_time = Instant::now(); let mut attempt_start_time = operation_start_time; + let target = ChannelTarget::from($self.affinity()); let stream = match $self .client - .$rpc_method($request.clone(), $gax_options.clone(), $self.channel_hint) + .$rpc_method($request.clone(), $gax_options.clone(), target) .send() .await { @@ -1145,7 +1138,7 @@ macro_rules! execute_stream_with_retry { attempt_start_time = Instant::now(); match $self .client - .$rpc_method($request.clone(), $gax_options.clone(), $self.channel_hint) + .$rpc_method($request.clone(), $gax_options.clone(), target) .send() .await { @@ -1173,7 +1166,6 @@ macro_rules! execute_stream_with_retry { session_name: $self.session_name.clone(), transaction_tag: $self.transaction_tag.clone(), operation: $operation_variant($request), - channel_hint: $self.channel_hint, gax_options: $gax_options, method_name: $method_name, attempt_start_time: Some(attempt_start_time), @@ -1191,9 +1183,10 @@ impl ReadContext { seqno: Option, ) -> crate::Result { let statement = statement.into(); + let target = ChannelTarget::from(self.affinity()); let gax_options = self .client - .attach_request_id(statement.gax_options().clone(), self.channel_hint); + .attach_request_id(statement.gax_options().clone(), target); let mut request = statement .into_request() .set_session(self.session_name.clone()) @@ -1216,9 +1209,10 @@ impl ReadContext { read: T, ) -> crate::Result { let read = read.into(); + let target = ChannelTarget::from(self.affinity()); let gax_options = self .client - .attach_request_id(read.gax_options.clone(), self.channel_hint); + .attach_request_id(read.gax_options.clone(), target); let mut request = read .into_request() .set_session(self.session_name.clone()) @@ -3602,7 +3596,6 @@ pub(crate) mod tests { transaction_selector: selector, precommit_token_tracker: crate::read_only_transaction::PrecommitTokenTracker::new(), transaction_tag: None, - channel_hint: 0, begin_transaction_request_options: None, affinity: None, }; @@ -3786,10 +3779,10 @@ pub(crate) mod tests { transaction .affinity() .expect("affinity present") - .set_entry_id(202); + .set_entry_id(1); assert_eq!( affinity.pinned_entry_id(), - Some(202), + Some(1), "Affinity handle passed to builder must observe the pinned channel ID" ); @@ -3798,7 +3791,7 @@ pub(crate) mod tests { .await?; assert_eq!( result_set.affinity().pinned_entry_id(), - Some(202), + Some(1), "ResultSet generated from MultiUse transaction must share the same pinned affinity" ); diff --git a/src/spanner/src/read_write_transaction.rs b/src/spanner/src/read_write_transaction.rs index bbe1a66ede..4e6838c0af 100644 --- a/src/spanner/src/read_write_transaction.rs +++ b/src/spanner/src/read_write_transaction.rs @@ -15,7 +15,7 @@ use crate::Error; use crate::RequestOptions; use crate::batch::BatchDml; -use crate::channel_pool::TransactionAffinity; +use crate::channel_pool::{ChannelTarget, TransactionAffinity}; use crate::client::amend_request_options_for_lar; use crate::database_client::DatabaseClient; use crate::error::internal_error; @@ -171,15 +171,16 @@ impl ReadWriteTransactionBuilder { async fn begin( &self, session_name: String, - channel_hint: usize, + affinity: &TransactionAffinity, request_options: crate::RequestOptions, ) -> crate::Result { + let target = ChannelTarget::from(affinity); let response = crate::read_only_transaction::execute_begin_transaction( &self.client, session_name, self.options.clone(), self.transaction_tag.clone(), - channel_hint, + target, request_options, None, ) @@ -192,10 +193,13 @@ impl ReadWriteTransactionBuilder { } pub(crate) async fn build( - self, + mut self, deadline: Option, ) -> crate::Result { - let channel_hint = self.client.next_channel_hint(); + let affinity = self + .affinity + .take() + .unwrap_or_else(|| Arc::new(TransactionAffinity::new_read_write())); let transaction_selector = match self.begin_transaction_option { BeginTransactionOption::ExplicitBegin => { let mut options = self.begin_gax_options.clone().unwrap_or_default(); @@ -205,7 +209,7 @@ impl ReadWriteTransactionBuilder { &mut options, ); - self.begin(self.session_name.clone(), channel_hint, options) + self.begin(self.session_name.clone(), &affinity, options) .await? } BeginTransactionOption::InlineBegin => ReadContextTransactionSelector::Lazy(Arc::new( @@ -213,11 +217,6 @@ impl ReadWriteTransactionBuilder { )), }; - let affinity = Some( - self.affinity - .unwrap_or_else(|| Arc::new(TransactionAffinity::new_read_write())), - ); - Ok(ReadWriteTransaction { context: ReadContext { session_name: self.session_name, @@ -225,9 +224,8 @@ impl ReadWriteTransactionBuilder { transaction_selector, precommit_token_tracker: PrecommitTokenTracker::new(), transaction_tag: self.transaction_tag, - channel_hint, begin_transaction_request_options: None, - affinity, + affinity: Some(affinity), }, seqno: Arc::new(AtomicI64::new(1)), max_commit_delay: self.max_commit_delay, @@ -240,7 +238,6 @@ impl ReadWriteTransactionBuilder { }) } - #[allow(dead_code)] pub(crate) fn with_affinity(mut self, affinity: Arc) -> Self { self.affinity = Some(affinity); self @@ -328,14 +325,11 @@ macro_rules! execute_with_retry { let mut guard = LazyTransactionStartGuard::new($self.context.transaction_selector.clone(), is_starting); + let target = ChannelTarget::from($self.context.affinity()); let response_result = $self .context .client - .$rpc_method( - $request.clone(), - $gax_options.clone(), - $self.context.channel_hint, - ) + .$rpc_method($request.clone(), $gax_options.clone(), target) .await; let service_error = response_result @@ -662,10 +656,11 @@ impl ReadWriteTransaction { let mut gax_options = self.commit_gax_options.clone().unwrap_or_default(); self.amend_gax_options(&mut gax_options); + let target = ChannelTarget::from(self.affinity()); let response = self .context .client - .commit(request, gax_options, self.context.channel_hint) + .commit(request, gax_options, target) .await?; let response = @@ -681,12 +676,16 @@ impl ReadWriteTransaction { self.context .client - .commit(retry_commit_req, gax_options, self.context.channel_hint) + .commit(retry_commit_req, gax_options, target) .await? } else { response }; + if let Some(affinity) = self.affinity() { + affinity.release_rw_guard(); + } + Ok(response) } @@ -703,11 +702,16 @@ impl ReadWriteTransaction { let mut gax_options = RequestOptions::default(); self.amend_gax_options(&mut gax_options); + let target = ChannelTarget::from(self.affinity()); self.context .client - .rollback(request, gax_options, self.context.channel_hint) + .rollback(request, gax_options, target) .await?; + if let Some(affinity) = self.affinity() { + affinity.release_rw_guard(); + } + Ok(()) } @@ -720,7 +724,6 @@ impl ReadWriteTransaction { } /// Returns a reference to the transaction channel affinity handle, if set. - #[allow(dead_code)] pub(crate) fn affinity(&self) -> Option<&TransactionAffinity> { self.context.affinity() } @@ -786,6 +789,7 @@ impl RetryPolicy for TransactionBoundedRetryPolicy { #[cfg(test)] mod tests { use super::*; + use crate::client::{Spanner, SpannerBuilderExt}; use crate::error::BatchUpdateError; use crate::read_only_transaction::tests::{create_session_mock, setup_db_client}; use crate::result_set::tests::{adapt, string_val}; @@ -4304,10 +4308,10 @@ mod tests { transaction .affinity() .expect("affinity present") - .set_entry_id(101); + .set_entry_id(1); assert_eq!( affinity.pinned_entry_id(), - Some(101), + Some(1), "Affinity handle passed to builder must observe the pinned channel ID" ); @@ -4316,7 +4320,7 @@ mod tests { .await?; assert_eq!( result_set.affinity().pinned_entry_id(), - Some(101), + Some(1), "ResultSet generated from ReadWrite transaction must share the same pinned affinity" ); @@ -4348,4 +4352,640 @@ mod tests { Ok(()) } + + async fn setup_db_client_with_dynamic_pool( + mock: spanner_grpc_mock::MockSpanner, + initial_channels: usize, + max_channels: usize, + ) -> (DatabaseClient, Spanner, tokio::task::JoinHandle<()>) { + use crate::channel_pool::DynamicChannelPoolConfig; + use crate::client::Spanner; + use google_cloud_auth::credentials::anonymous::Builder as Anonymous; + + let (address, server) = spanner_grpc_mock::start("127.0.0.1:0", mock) + .await + .expect("Failed to start mock server"); + + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(initial_channels) + .with_min_channels(initial_channels) + .with_max_channels(max_channels); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("Failed to build client"); + + let database_client = spanner + .database_client("projects/p/instances/i/databases/d") + .build() + .await + .expect("Failed to create DatabaseClient"); + + (database_client, spanner, server) + } + + #[tokio_test_no_panics] + async fn read_write_transaction_affinity_under_dynamic_pool_inline_begin() -> anyhow::Result<()> + { + run_read_write_transaction_affinity_under_dynamic_pool(BeginTransactionOption::InlineBegin) + .await + } + + #[tokio_test_no_panics] + async fn read_write_transaction_affinity_under_dynamic_pool_explicit_begin() + -> anyhow::Result<()> { + run_read_write_transaction_affinity_under_dynamic_pool( + BeginTransactionOption::ExplicitBegin, + ) + .await + } + + async fn run_read_write_transaction_affinity_under_dynamic_pool( + begin_transaction_option: BeginTransactionOption, + ) -> anyhow::Result<()> { + use crate::batch_dml::BatchDml; + use crate::statement::Statement; + + let mut mock = create_session_mock(); + let remote_addresses = Arc::new(Mutex::new(Vec::new())); + + if begin_transaction_option == BeginTransactionOption::ExplicitBegin { + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_begin_transaction() + .once() + .returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + })) + }); + } + + // 1. Query: execute_streaming_sql + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + let mut metadata = v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + ..Default::default() + }; + if begin_transaction_option == BeginTransactionOption::InlineBegin { + metadata.transaction = Some(v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + }); + } + let partial_result_set = v1::PartialResultSet { + metadata: Some(metadata), + ..Default::default() + }; + Ok(tonic::Response::from(adapt([Ok(partial_result_set)]))) + }); + + // 2. DML: execute_sql + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_sql().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(v1::ResultSet { + metadata: Some(v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + ..Default::default() + }), + stats: Some(v1::ResultSetStats { + row_count: Some(v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + })) + }); + + // 3. Batch DML: execute_batch_dml + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_batch_dml() + .once() + .returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(v1::ExecuteBatchDmlResponse { + result_sets: vec![v1::ResultSet { + stats: Some(v1::ResultSetStats { + row_count: Some(v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + }], + ..Default::default() + })) + }); + + // 4. Commit: commit + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_commit().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 1000, + nanos: 0, + }), + ..Default::default() + })) + }); + + let (database_client, _spanner, _server) = + setup_db_client_with_dynamic_pool(mock, 4, 8).await; + + let transaction = ReadWriteTransactionBuilder::new(database_client) + .with_begin_transaction_option(begin_transaction_option) + .build(None) + .await + .expect("Failed to build transaction"); + + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + let _ = result_set.next().await; + + let count = transaction + .execute_update("UPDATE Users SET Name = 'Alice' WHERE Id = 1") + .await?; + assert_eq!(count, 1, "Expected 1 row updated"); + + let batch_result = transaction + .execute_batch_update( + BatchDml::builder() + .add_statement("UPDATE Users SET Name = 'Bob' WHERE Id = 2") + .build(), + ) + .await?; + assert_eq!(batch_result.len(), 1, "Expected 1 batch statement executed"); + + let commit_result = transaction.commit().await?; + assert_eq!( + commit_result + .commit_timestamp + .expect("timestamp should be present") + .seconds(), + 1000, + "Commit timestamp mismatch" + ); + + let addresses = remote_addresses.lock().expect("mutex lock"); + let expected_rpc_count = + if begin_transaction_option == BeginTransactionOption::ExplicitBegin { + 5 + } else { + 4 + }; + assert_eq!( + addresses.len(), + expected_rpc_count, + "Expected {} RPCs executed", + expected_rpc_count + ); + let first_address = addresses[0]; + for (index, address) in addresses.iter().enumerate() { + assert_eq!( + *address, first_address, + "RPC at index {} must use the same channel as first RPC ({})", + index, first_address + ); + } + + Ok(()) + } + + #[tokio_test_no_panics] + async fn read_write_transaction_affinity_preserved_across_dynamic_pool_scale_up() + -> anyhow::Result<()> { + use crate::channel_pool::entry::ChannelEntry; + use crate::statement::Statement; + + let mut mock = create_session_mock(); + let remote_addresses = Arc::new(Mutex::new(Vec::new())); + + // Statement 1: Query (Inline begin) + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + let metadata = v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + transaction: Some(v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + }), + ..Default::default() + }; + let partial_result_set = v1::PartialResultSet { + metadata: Some(metadata), + ..Default::default() + }; + Ok(tonic::Response::from(adapt([Ok(partial_result_set)]))) + }); + + // Statement 2: Update + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_sql().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(v1::ResultSet { + metadata: Some(v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + ..Default::default() + }), + stats: Some(v1::ResultSetStats { + row_count: Some(v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + })) + }); + + // Commit + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_commit().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 1000, + nanos: 0, + }), + ..Default::default() + })) + }); + + // Start with 2 channels, pool can scale up to 8 + let (database_client, spanner, _server) = + setup_db_client_with_dynamic_pool(mock, 2, 8).await; + + let transaction = ReadWriteTransactionBuilder::new(database_client) + .build(None) + .await + .expect("Failed to build transaction"); + + // Execute Statement 1 + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + let _ = result_set.next().await; + + let pinned_entry_id = transaction + .affinity() + .expect("affinity must be present") + .pinned_entry_id(); + assert!( + pinned_entry_id.is_some(), + "Transaction affinity must be pinned after first statement" + ); + + // Simulate dynamic scale-up mid-transaction: add new channels to active_entries + { + let default_channel = spanner + .channel_pool() + .default_channel() + .expect("default channel must exist"); + let mut active_write = spanner + .channel_pool() + .inner + .active_entries + .write() + .expect("lock active_entries"); + let current_len = active_write.len(); + let next_entry_id = current_len as u64 + 1; + // Add 2 new channels, doubling active channels from 2 to 4 + active_write.push(Arc::new(ChannelEntry::new( + next_entry_id, + current_len + 1, + default_channel.clone(), + ))); + active_write.push(Arc::new(ChannelEntry::new( + next_entry_id + 1, + current_len + 2, + default_channel, + ))); + assert_eq!( + active_write.len(), + 4, + "Pool must now have 4 active channels" + ); + } + + // Execute Statement 2: Update + let count = transaction + .execute_update("UPDATE Users SET Name = 'Alice' WHERE Id = 1") + .await?; + assert_eq!(count, 1, "Expected 1 row updated"); + + // Commit + let _ = transaction.commit().await?; + + // Verify that all RPCs used the exact same channel remote address + let addresses = remote_addresses.lock().expect("mutex lock"); + assert_eq!(addresses.len(), 3, "Expected 3 RPCs executed"); + let initial_address = addresses[0]; + assert_eq!( + addresses[1], initial_address, + "Statement 2 must use the pinned channel despite pool scale-up" + ); + assert_eq!( + addresses[2], initial_address, + "Commit must use the pinned channel despite pool scale-up" + ); + + Ok(()) + } + + #[tokio_test_no_panics] + async fn read_write_transaction_affinity_preserved_during_channel_draining() + -> anyhow::Result<()> { + use crate::channel_pool::entry::ChannelState; + use crate::channel_pool::scaler::sweep_draining_channels; + use crate::statement::Statement; + + let mut mock = create_session_mock(); + let remote_addresses = Arc::new(Mutex::new(Vec::new())); + + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + let metadata = v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + transaction: Some(v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + }), + ..Default::default() + }; + let partial_result_set = v1::PartialResultSet { + metadata: Some(metadata), + ..Default::default() + }; + Ok(tonic::Response::from(adapt([Ok(partial_result_set)]))) + }); + + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_sql().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(v1::ResultSet { + metadata: Some(v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + ..Default::default() + }), + stats: Some(v1::ResultSetStats { + row_count: Some(v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + })) + }); + + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_commit().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 1000, + nanos: 0, + }), + ..Default::default() + })) + }); + + let (database_client, spanner, _server) = + setup_db_client_with_dynamic_pool(mock, 4, 8).await; + + let transaction = ReadWriteTransactionBuilder::new(database_client) + .build(None) + .await + .expect("Failed to build transaction"); + + // Statement 1: Query pins the channel + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + let _ = result_set.next().await; + drop(result_set); + + let pinned_entry_id = transaction + .affinity() + .expect("affinity must be present") + .pinned_entry_id() + .expect("channel must be pinned after query"); + + // Locate pinned channel entry and verify active_rw_count is 1 + let pinned_entry = { + let active_guard = spanner + .channel_pool() + .inner + .active_entries + .read() + .expect("lock active_entries"); + active_guard + .iter() + .find(|entry| entry.id == pinned_entry_id) + .map(Arc::clone) + .expect("pinned entry must exist in active_entries") + }; + assert_eq!( + pinned_entry.active_rw_count(), + 1, + "Pinned entry must have active_rw_count == 1 while transaction is open" + ); + + // Move the pinned channel to Draining (simulating scale-down) + { + let mut active_write = spanner + .channel_pool() + .inner + .active_entries + .write() + .expect("lock active_entries"); + active_write.retain(|entry| entry.id != pinned_entry_id); + + pinned_entry.set_state(ChannelState::Draining); + let mut draining_write = spanner + .channel_pool() + .inner + .draining_entries + .write() + .expect("lock draining_entries"); + draining_write.push(Arc::clone(&pinned_entry)); + } + + // Run sweep_draining_channels. Because active_rw_count is 1, Channel X must NOT be closed! + sweep_draining_channels(&spanner.channel_pool().inner, StdDuration::from_millis(0)); + assert!( + pinned_entry.is_draining(), + "Pinned entry must remain in Draining state while R/W transaction is active" + ); + assert_eq!( + spanner.channel_pool().draining_channel_count(), + 1, + "Draining channel count must still be 1" + ); + + // Execute Statement 2: Update while the channel is draining! + let count = transaction + .execute_update("UPDATE Users SET Name = 'Alice' WHERE Id = 1") + .await?; + assert_eq!(count, 1, "Expected 1 row updated"); + + // Commit while the channel is draining! + let _ = transaction.commit().await?; + + // Verify that all 3 RPCs executed on the same channel + let addresses = remote_addresses.lock().expect("mutex lock"); + assert_eq!(addresses.len(), 3, "Expected 3 RPCs executed"); + assert_eq!( + addresses[1], addresses[0], + "Statement 2 must execute on the draining pinned channel" + ); + assert_eq!( + addresses[2], addresses[0], + "Commit must execute on the draining pinned channel" + ); + + // Now that the transaction is done and dropped, active_rw_count must be 0 + assert_eq!( + pinned_entry.active_rw_count(), + 0, + "active_rw_count must drop to 0 after transaction completes" + ); + + // Sweeping now must close the channel and remove it from draining entries! + sweep_draining_channels(&spanner.channel_pool().inner, StdDuration::from_millis(0)); + assert!( + pinned_entry.is_closed(), + "Draining entry must transition to Closed once transaction completes" + ); + assert_eq!( + spanner.channel_pool().draining_channel_count(), + 0, + "Draining channels must be empty after sweep" + ); + + Ok(()) + } + + #[tokio_test_no_panics] + async fn read_write_transaction_affinity_preserved_on_rollback() -> anyhow::Result<()> { + use crate::statement::Statement; + + let mut mock = create_session_mock(); + let remote_addresses = Arc::new(Mutex::new(Vec::new())); + + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + let metadata = v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + transaction: Some(v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + }), + ..Default::default() + }; + let partial_result_set = v1::PartialResultSet { + metadata: Some(metadata), + ..Default::default() + }; + Ok(tonic::Response::from(adapt([Ok(partial_result_set)]))) + }); + + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_rollback().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + let request = request.into_inner(); + assert_eq!(request.transaction_id, vec![1, 2, 3]); + Ok(tonic::Response::new(())) + }); + + let (database_client, _spanner, _server) = + setup_db_client_with_dynamic_pool(mock, 4, 8).await; + + let transaction = ReadWriteTransactionBuilder::new(database_client) + .build(None) + .await + .expect("Failed to build transaction"); + + // Statement 1: Query pins the channel + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + let _ = result_set.next().await; + + // Explicit rollback + transaction.rollback().await?; + + let addresses = remote_addresses.lock().expect("mutex lock"); + assert_eq!(addresses.len(), 2, "Expected 2 RPCs executed"); + assert_eq!( + addresses[1], addresses[0], + "Rollback must use the same channel as preceding statements" + ); + + Ok(()) + } } diff --git a/src/spanner/src/request_id.rs b/src/spanner/src/request_id.rs index c680b58a7f..79f97892a4 100644 --- a/src/spanner/src/request_id.rs +++ b/src/spanner/src/request_id.rs @@ -238,22 +238,24 @@ mod tests { }; // Execute first RPC (attempt 1 -> UNAVAILABLE, attempt 2 -> success) + let lease1 = client.pick_channel(); let _session1 = client .create_session( request.clone(), crate::RequestOptions::default(), - client.get_channel(0), + &lease1, &Observability::disabled_arc(), ) .await .expect("first create_session should succeed after retry"); // Execute second RPC (attempt 1 -> success) + let lease2 = client.pick_channel(); let _session2 = client .create_session( request, crate::RequestOptions::default(), - client.get_channel(0), + &lease2, &Observability::disabled_arc(), ) .await diff --git a/src/spanner/src/result_set.rs b/src/spanner/src/result_set.rs index 8531b70b31..e95e9c22ec 100644 --- a/src/spanner/src/result_set.rs +++ b/src/spanner/src/result_set.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::channel_pool::TransactionAffinity; +use crate::channel_pool::{ChannelTarget, TransactionAffinity}; use crate::database_client::DatabaseClient; use crate::error::internal_error; use crate::google::spanner::v1::{self, PartialResultSet}; @@ -85,7 +85,6 @@ pub struct ResultSet { max_buffered_partial_result_sets: usize, retry_count: usize, transaction_selector: Option, - channel_hint: usize, gax_options: GaxRequestOptions, method_name: &'static str, headers: HeaderMap, @@ -110,7 +109,6 @@ pub(crate) struct ResultSetParams { pub session_name: String, pub transaction_tag: Option, pub operation: StreamOperation, - pub channel_hint: usize, pub gax_options: GaxRequestOptions, pub method_name: &'static str, pub attempt_start_time: Option, @@ -148,7 +146,6 @@ impl ResultSet { session_name, transaction_tag, operation, - channel_hint, gax_options, method_name, attempt_start_time, @@ -181,7 +178,6 @@ impl ResultSet { max_buffered_partial_result_sets: MAX_BUFFERED_PARTIAL_RESULT_SETS, retry_count: 0, transaction_selector, - channel_hint, gax_options, tokio_handle: Handle::try_current().ok(), method_name, @@ -597,7 +593,6 @@ impl ResultSet { client: self.client.clone(), session_name: self.session_name.clone(), transaction_tag: self.transaction_tag.clone(), - channel_hint: self.channel_hint, request_options: self.gax_options.clone(), is_stream_fallback: true, precommit_token_tracker: self.precommit_token_tracker.clone(), @@ -752,6 +747,7 @@ impl ResultSet { self.attempt_recorded = false; self.headers.clear(); + let target = ChannelTarget::from(&self.affinity); let stream_result = match &mut self.operation { StreamOperation::Query(req) => { req.resume_token = self.last_resume_token.clone(); @@ -759,7 +755,7 @@ impl ResultSet { .clone() .or_else(|| req.transaction.take()); self.client - .execute_streaming_sql(req.clone(), self.gax_options.clone(), self.channel_hint) + .execute_streaming_sql(req.clone(), self.gax_options.clone(), target) .send() .await } @@ -769,7 +765,7 @@ impl ResultSet { .clone() .or_else(|| req.transaction.take()); self.client - .streaming_read(req.clone(), self.gax_options.clone(), self.channel_hint) + .streaming_read(req.clone(), self.gax_options.clone(), target) .send() .await } @@ -1979,7 +1975,11 @@ pub(crate) mod tests { .set_sql("SELECT 1".to_string()); let stream = db_client - .execute_streaming_sql(req.clone(), GaxRequestOptions::default(), 0) + .execute_streaming_sql( + req.clone(), + GaxRequestOptions::default(), + ChannelTarget::Any, + ) .send() .await?; @@ -1991,7 +1991,6 @@ pub(crate) mod tests { session_name: "session".to_string(), transaction_tag: None, operation: StreamOperation::Query(req), - channel_hint: 0, gax_options: GaxRequestOptions::default(), method_name: "ExecuteStreamingSql", attempt_start_time: None, diff --git a/src/spanner/src/routing/mock_tests.rs b/src/spanner/src/routing/mock_tests.rs index d976209bf2..1ae52bc799 100644 --- a/src/spanner/src/routing/mock_tests.rs +++ b/src/spanner/src/routing/mock_tests.rs @@ -26,6 +26,7 @@ //! - Proactive background cache synchronization via `CacheSubscriber`. use crate::RequestOptions; +use crate::channel_pool::ChannelTarget; use crate::client::{Spanner, SpannerBuilderExt}; use crate::database_client::DatabaseClient; use crate::key; @@ -2829,7 +2830,11 @@ async fn unary_commit_routes_to_affinity_address_and_clears_affinity() -> anyhow .set_transaction_id(Bytes::copy_from_slice(transaction_id)); let response = database_client - .commit(commit_request, RequestOptions::default(), 0) + .commit( + commit_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await?; assert!( @@ -2908,7 +2913,11 @@ async fn unary_rollback_routes_to_affinity_address_and_clears_affinity() -> anyh .set_transaction_id(Bytes::copy_from_slice(transaction_id)); database_client - .rollback(rollback_request, RequestOptions::default(), 0) + .rollback( + rollback_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await?; assert!( @@ -3004,7 +3013,11 @@ async fn unary_single_use_commit_routes_to_leader_tablet_replica() -> anyhow::Re .set_mutations(vec![mutation.build_proto()]); let response = database_client - .commit(commit_request, RequestOptions::default(), 0) + .commit( + commit_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await?; assert!( @@ -3102,7 +3115,7 @@ async fn unary_begin_transaction_with_mutation_key_routes_to_leader_and_records_ .set_mutation_key(mutation.build_proto()); let response = database_client - .begin_transaction(begin_request, RequestOptions::default(), 0) + .begin_transaction(begin_request, RequestOptions::default(), ChannelTarget::Any) .await?; assert!( @@ -3202,7 +3215,7 @@ async fn unary_begin_transaction_with_read_only_options_does_not_record_affinity .set_mutation_key(mutation.build_proto()); let response = database_client - .begin_transaction(begin_request, RequestOptions::default(), 0) + .begin_transaction(begin_request, RequestOptions::default(), ChannelTarget::Any) .await?; assert!( @@ -3275,7 +3288,11 @@ async fn unary_execute_sql_routes_to_affinity_address() -> anyhow::Result<()> { ); let _ = database_client - .execute_sql(execute_sql_request, RequestOptions::default(), 0) + .execute_sql( + execute_sql_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await?; assert!( @@ -3334,7 +3351,11 @@ async fn unary_execute_sql_with_inline_begin_rw_records_affinity() -> anyhow::Re ); let _ = database_client - .execute_sql(execute_sql_request, RequestOptions::default(), 0) + .execute_sql( + execute_sql_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await?; assert!( @@ -3405,7 +3426,11 @@ async fn unary_execute_batch_dml_routes_to_affinity_address() -> anyhow::Result< .set_seqno(1); let _ = database_client - .execute_batch_dml(batch_dml_request, RequestOptions::default(), 0) + .execute_batch_dml( + batch_dml_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await?; assert!( @@ -3470,7 +3495,11 @@ async fn unary_execute_batch_dml_with_inline_begin_rw_records_affinity() -> anyh .set_seqno(1); let _ = database_client - .execute_batch_dml(batch_dml_request, RequestOptions::default(), 0) + .execute_batch_dml( + batch_dml_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await?; assert!( @@ -3527,7 +3556,11 @@ async fn streaming_execute_sql_with_inline_begin_rw_records_affinity() -> anyhow ); let mut stream = database_client - .execute_streaming_sql(execute_sql_request, RequestOptions::default(), 0) + .execute_streaming_sql( + execute_sql_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .send() .await?; @@ -3634,7 +3667,11 @@ async fn streaming_execute_sql_with_multiple_chunks_records_affinity_and_yields_ ); let mut stream = database_client - .execute_streaming_sql(execute_sql_request, RequestOptions::default(), 0) + .execute_streaming_sql( + execute_sql_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .send() .await?; @@ -3722,7 +3759,7 @@ async fn streaming_read_with_inline_begin_rw_records_affinity() -> anyhow::Resul ); let mut stream = database_client - .streaming_read(read_request, RequestOptions::default(), 0) + .streaming_read(read_request, RequestOptions::default(), ChannelTarget::Any) .send() .await?; @@ -3792,7 +3829,11 @@ async fn streaming_execute_sql_with_inline_begin_ro_does_not_record_affinity() - ); let mut stream = database_client - .execute_streaming_sql(execute_sql_request, RequestOptions::default(), 0) + .execute_streaming_sql( + execute_sql_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .send() .await?; @@ -3861,7 +3902,11 @@ async fn streaming_execute_sql_with_error_does_not_record_affinity() -> anyhow:: ); let mut stream = database_client - .execute_streaming_sql(execute_sql_request, RequestOptions::default(), 0) + .execute_streaming_sql( + execute_sql_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .send() .await?; @@ -4309,7 +4354,7 @@ async fn streaming_read_with_inline_begin_rw_routed_to_tablet_records_tablet_aff ); let mut stream = database_client - .streaming_read(read_request, RequestOptions::default(), 0) + .streaming_read(read_request, RequestOptions::default(), ChannelTarget::Any) .send() .await?; @@ -4396,7 +4441,11 @@ async fn unary_partition_read_routes_to_tablet_node() -> anyhow::Result<()> { .set_key_set(key_set.into_proto()); let _ = database_client - .partition_read(partition_read_request, RequestOptions::default(), 0) + .partition_read( + partition_read_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await?; assert!( @@ -4460,7 +4509,11 @@ async fn unary_partition_query_with_transaction_id_routes_to_affinity_address() ); let _ = database_client - .partition_query(partition_query_request, RequestOptions::default(), 0) + .partition_query( + partition_query_request, + RequestOptions::default(), + ChannelTarget::Any, + ) .await?; assert!( @@ -5026,7 +5079,7 @@ async fn end_to_end_unary_execute_sql_with_key_recipe_routes_to_tablet_replica() let request1 = statement.clone().into_request(); let _ = database_client - .execute_sql(request1, RequestOptions::default(), 0) + .execute_sql(request1, RequestOptions::default(), ChannelTarget::Any) .await?; assert!( @@ -5062,7 +5115,7 @@ async fn end_to_end_unary_execute_sql_with_key_recipe_routes_to_tablet_replica() // 4. Second execution (cache hit): routes directly to tablet mock with attached routing hint let request2 = statement.into_request(); let _ = database_client - .execute_sql(request2, RequestOptions::default(), 0) + .execute_sql(request2, RequestOptions::default(), ChannelTarget::Any) .await?; assert!( @@ -5418,7 +5471,7 @@ async fn unary_execute_sql_with_directed_read_options_and_key_recipe_routes_to_d let request1 = statement.clone().into_request(); let _ = database_client - .execute_sql(request1, RequestOptions::default(), 0) + .execute_sql(request1, RequestOptions::default(), ChannelTarget::Any) .await?; assert!( @@ -5466,7 +5519,7 @@ async fn unary_execute_sql_with_directed_read_options_and_key_recipe_routes_to_d // Cache hit: routes directly to mock_east replica matching directed read options let request2 = statement.into_request(); let _ = database_client - .execute_sql(request2, RequestOptions::default(), 0) + .execute_sql(request2, RequestOptions::default(), ChannelTarget::Any) .await?; assert!( @@ -5987,7 +6040,7 @@ async fn unary_rpc_feedback_success_records_latency_and_repairs_cooldown() -> an .execute_sql( statement.clone().into_request(), RequestOptions::default(), - 0, + ChannelTarget::Any, ) .await?; @@ -6024,7 +6077,7 @@ async fn unary_rpc_feedback_success_records_latency_and_repairs_cooldown() -> an .execute_sql( statement.clone().into_request(), RequestOptions::default(), - 0, + ChannelTarget::Any, ) .await?; } @@ -6143,7 +6196,7 @@ async fn unary_rpc_feedback_resource_exhausted_places_endpoint_on_cooldown_and_p .execute_sql( statement.clone().into_request(), RequestOptions::default(), - 0, + ChannelTarget::Any, ) .await?; @@ -6178,7 +6231,7 @@ async fn unary_rpc_feedback_resource_exhausted_places_endpoint_on_cooldown_and_p .execute_sql( statement.clone().into_request(), RequestOptions::default(), - 0, + ChannelTarget::Any, ) .await; assert!( @@ -6204,7 +6257,11 @@ async fn unary_rpc_feedback_resource_exhausted_places_endpoint_on_cooldown_and_p // Third execution: since tablet is on cooldown, router falls back to gateway let fallback_result = database_client - .execute_sql(statement.into_request(), RequestOptions::default(), 0) + .execute_sql( + statement.into_request(), + RequestOptions::default(), + ChannelTarget::Any, + ) .await; assert!( fallback_result.is_ok(), @@ -6275,7 +6332,7 @@ async fn unary_rpc_feedback_unavailable_places_endpoint_on_cooldown() -> anyhow: .execute_sql( statement.clone().into_request(), RequestOptions::default(), - 0, + ChannelTarget::Any, ) .await?; @@ -6298,7 +6355,7 @@ async fn unary_rpc_feedback_unavailable_places_endpoint_on_cooldown() -> anyhow: let mut options = RequestOptions::default(); options.set_retry_policy(NeverRetry); let result = database_client - .execute_sql(statement.into_request(), options, 0) + .execute_sql(statement.into_request(), options, ChannelTarget::Any) .await; assert!( result.is_err(), @@ -6351,7 +6408,11 @@ async fn unary_rpc_feedback_gateway_fallback_never_placed_on_cooldown() -> anyho .add_param("account_id", 42i64) .build(); let result = database_client - .execute_sql(statement.into_request(), RequestOptions::default(), 0) + .execute_sql( + statement.into_request(), + RequestOptions::default(), + ChannelTarget::Any, + ) .await; assert!( result.is_err(), @@ -6435,7 +6496,7 @@ async fn unary_rpc_feedback_non_retryable_error_does_not_place_endpoint_on_coold .execute_sql( statement.clone().into_request(), RequestOptions::default(), - 0, + ChannelTarget::Any, ) .await?; @@ -6456,7 +6517,11 @@ async fn unary_rpc_feedback_non_retryable_error_does_not_place_endpoint_on_coold // Direct execution to tablet fails with non-retryable InvalidArgument let result = database_client - .execute_sql(statement.into_request(), RequestOptions::default(), 0) + .execute_sql( + statement.into_request(), + RequestOptions::default(), + ChannelTarget::Any, + ) .await; assert!( result.is_err(), @@ -6576,7 +6641,7 @@ async fn unary_rpc_feedback_begin_transaction_records_latency_and_repairs_cooldo .set_mutation_key(mutation.build_proto()); let response = database_client - .begin_transaction(begin_request, RequestOptions::default(), 0) + .begin_transaction(begin_request, RequestOptions::default(), ChannelTarget::Any) .await?; assert!( @@ -6680,7 +6745,11 @@ async fn unary_rpc_feedback_direct_affinity_zero_group_uid_handles_cooldown_and_ ); let result_1 = database_client - .execute_sql(execute_request_1, RequestOptions::default(), 0) + .execute_sql( + execute_request_1, + RequestOptions::default(), + ChannelTarget::Any, + ) .await; assert!( result_1.is_ok(), @@ -6707,7 +6776,11 @@ async fn unary_rpc_feedback_direct_affinity_zero_group_uid_handles_cooldown_and_ ); let result_2 = database_client - .execute_sql(execute_request_2, RequestOptions::default(), 0) + .execute_sql( + execute_request_2, + RequestOptions::default(), + ChannelTarget::Any, + ) .await; assert!( result_2.is_err(), @@ -6737,7 +6810,11 @@ async fn unary_rpc_feedback_direct_affinity_zero_group_uid_handles_cooldown_and_ ); let result_3 = database_client - .execute_sql(execute_request_3, RequestOptions::default(), 0) + .execute_sql( + execute_request_3, + RequestOptions::default(), + ChannelTarget::Any, + ) .await; assert!( result_3.is_ok(), diff --git a/src/spanner/src/server_streaming/builder.rs b/src/spanner/src/server_streaming/builder.rs index 006a2b42b4..99c889f00b 100644 --- a/src/spanner/src/server_streaming/builder.rs +++ b/src/spanner/src/server_streaming/builder.rs @@ -415,7 +415,8 @@ mod tests { .await .expect("spanner client should build"); - let grpc_client = spanner.channels[0] + let grpc_client = spanner + .pick_channel() .grpc_client .clone() .expect("grpc client should exist"); diff --git a/src/spanner/src/server_streaming/stream.rs b/src/spanner/src/server_streaming/stream.rs index 08e8d3448c..0d622aedfa 100644 --- a/src/spanner/src/server_streaming/stream.rs +++ b/src/spanner/src/server_streaming/stream.rs @@ -168,6 +168,7 @@ pub(crate) type CacheUpdateStream = SpannerServerStream; #[cfg(test)] mod tests { use super::*; + use crate::channel_pool::ChannelTarget; use crate::model::ExecuteSqlRequest; use crate::read_only_transaction::tests::{create_session_mock, setup_db_client}; use gaxi::grpc::tonic::{Response, Status}; @@ -254,7 +255,7 @@ mod tests { .set_session(db_client.session_name()) .set_sql("SELECT 1"); let stream = db_client - .execute_streaming_sql(request, RequestOptions::default(), 0) + .execute_streaming_sql(request, RequestOptions::default(), ChannelTarget::Any) .with_lifetime_guard(guard) .send() .await?; @@ -298,7 +299,7 @@ mod tests { .set_session(db_client.session_name()) .set_sql("SELECT 1"); let mut stream = db_client - .execute_streaming_sql(request, RequestOptions::default(), 0) + .execute_streaming_sql(request, RequestOptions::default(), ChannelTarget::Any) .with_lifetime_guard(guard) .send() .await?; @@ -345,7 +346,7 @@ mod tests { .set_session(db_client.session_name()) .set_sql("SELECT 1"); let mut stream = db_client - .execute_streaming_sql(request, RequestOptions::default(), 0) + .execute_streaming_sql(request, RequestOptions::default(), ChannelTarget::Any) .with_lifetime_guard(guard) .send() .await?; @@ -391,7 +392,7 @@ mod tests { .set_session(db_client.session_name()) .set_sql("SELECT 1"); let mut stream = db_client - .execute_streaming_sql(request, RequestOptions::default(), 0) + .execute_streaming_sql(request, RequestOptions::default(), ChannelTarget::Any) .send() .await?; @@ -430,7 +431,7 @@ mod tests { .set_session(db_client.session_name()) .set_sql("SELECT 1"); let mut stream = db_client - .execute_streaming_sql(request, RequestOptions::default(), 0) + .execute_streaming_sql(request, RequestOptions::default(), ChannelTarget::Any) .send() .await? .with_transaction_id_callback(callback); @@ -512,7 +513,7 @@ mod tests { .set_session(db_client.session_name()) .set_sql("SELECT 1"); let mut stream = db_client - .execute_streaming_sql(request, RequestOptions::default(), 0) + .execute_streaming_sql(request, RequestOptions::default(), ChannelTarget::Any) .send() .await? .with_transaction_id_callback(callback); @@ -544,7 +545,7 @@ mod tests { .set_session(db_client.session_name()) .set_sql("SELECT 1"); let mut stream = db_client - .execute_streaming_sql(request, RequestOptions::default(), 0) + .execute_streaming_sql(request, RequestOptions::default(), ChannelTarget::Any) .send() .await? .with_transaction_id_callback(callback); diff --git a/src/spanner/src/session_maintainer.rs b/src/spanner/src/session_maintainer.rs index 5168c853f1..e389080b20 100644 --- a/src/spanner/src/session_maintainer.rs +++ b/src/spanner/src/session_maintainer.rs @@ -66,6 +66,10 @@ impl ManagedSessionMaintainer { let session = Self::create_session(&spanner, &database_name, &database_role, &options, &o11y).await?; + spanner + .channel_pool() + .set_prime_session(session.name.clone()); + let maintainer = Arc::new(ManagedSessionMaintainer { spanner, session: RwLock::new(ManagedSession { @@ -113,6 +117,10 @@ impl ManagedSessionMaintainer { ) .await?; + self.spanner + .channel_pool() + .set_prime_session(new_session.name.clone()); + let mut guard = self.session.write().expect("failed to write session"); *guard = ManagedSession { session: Arc::new(new_session), diff --git a/src/spanner/src/transaction_runner.rs b/src/spanner/src/transaction_runner.rs index abd37257af..5b2cdcb31d 100644 --- a/src/spanner/src/transaction_runner.rs +++ b/src/spanner/src/transaction_runner.rs @@ -2367,4 +2367,187 @@ mod tests { Ok(()) } + + async fn setup_db_client_with_dynamic_pool( + mock: spanner_grpc_mock::MockSpanner, + initial_channels: usize, + max_channels: usize, + ) -> (DatabaseClient, tokio::task::JoinHandle<()>) { + use crate::channel_pool::DynamicChannelPoolConfig; + use crate::client::{Spanner, SpannerBuilderExt}; + use google_cloud_auth::credentials::anonymous::Builder as Anonymous; + + let (address, server) = spanner_grpc_mock::start("127.0.0.1:0", mock) + .await + .expect("Failed to start mock server"); + + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(initial_channels) + .with_min_channels(initial_channels) + .with_max_channels(max_channels); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("Failed to build client"); + + let database_client = spanner + .database_client("projects/p/instances/i/databases/d") + .build() + .await + .expect("Failed to create DatabaseClient"); + + (database_client, server) + } + + #[tokio_test_no_panics] + async fn transaction_runner_aborted_retry_routes_all_attempts_to_same_channel() + -> anyhow::Result<()> { + use crate::result_set::tests::adapt; + use crate::statement::Statement; + + let mut mock = create_session_mock(); + let remote_addresses = Arc::new(Mutex::new(Vec::new())); + + // Attempt 1: Statement 1 (ExecuteStreamingSql) succeeds + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + let metadata = v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + transaction: Some(v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + }), + ..Default::default() + }; + let partial_result_set = v1::PartialResultSet { + metadata: Some(metadata), + ..Default::default() + }; + Ok(tonic::Response::from(adapt([Ok(partial_result_set)]))) + }); + + // Attempt 1: Statement 2 (ExecuteSql) fails with Aborted + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_sql().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Err(tonic::Status::new( + tonic::Code::Aborted, + "Transaction was aborted", + )) + }); + + // Attempt 2: Statement 1 (ExecuteStreamingSql) succeeds + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_streaming_sql() + .once() + .returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + let metadata = v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + transaction: Some(v1::Transaction { + id: vec![4, 5, 6], + ..Default::default() + }), + ..Default::default() + }; + let partial_result_set = v1::PartialResultSet { + metadata: Some(metadata), + ..Default::default() + }; + Ok(tonic::Response::from(adapt([Ok(partial_result_set)]))) + }); + + // Attempt 2: Statement 2 (ExecuteSql) succeeds + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_execute_sql().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(v1::ResultSet { + metadata: Some(v1::ResultSetMetadata { + row_type: Some(v1::StructType { fields: vec![] }), + ..Default::default() + }), + stats: Some(v1::ResultSetStats { + row_count: Some(v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + })) + }); + + // Attempt 2: Commit succeeds + let remote_addresses_clone = remote_addresses.clone(); + mock.expect_commit().once().returning(move |request| { + remote_addresses_clone.lock().expect("mutex lock").push( + request + .remote_addr() + .expect("remote_addr should be available"), + ); + Ok(tonic::Response::new(CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 2000, + nanos: 0, + }), + ..Default::default() + })) + }); + + let (database_client, _server) = setup_db_client_with_dynamic_pool(mock, 4, 8).await; + + let runner = database_client.read_write_transaction().build().await?; + let result = runner + .run(|transaction: ReadWriteTransaction| async move { + let mut result_set = transaction + .execute_query(Statement::builder("SELECT 1").build()) + .await?; + let _ = result_set.next().await; + let count = transaction + .execute_update("UPDATE Users SET Name = 'Alice' WHERE Id = 1") + .await?; + Ok(count) + }) + .await?; + + assert_eq!(result.result, 1, "Expected update count of 1 on retry"); + + let addresses = remote_addresses.lock().expect("mutex lock"); + assert_eq!( + addresses.len(), + 5, + "Expected 5 total RPCs across both attempts (2 on attempt 1, 3 on attempt 2)" + ); + + let initial_address = addresses[0]; + for (index, address) in addresses.iter().enumerate() { + assert_eq!( + *address, initial_address, + "RPC at index {} must use the same channel as attempt 1 ({})", + index, initial_address + ); + } + + Ok(()) + } } diff --git a/src/spanner/src/write_only_transaction.rs b/src/spanner/src/write_only_transaction.rs index da061eb113..deeed3627a 100644 --- a/src/spanner/src/write_only_transaction.rs +++ b/src/spanner/src/write_only_transaction.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::channel_pool::TransactionAffinity; +use crate::channel_pool::{ChannelTarget, TransactionAffinity}; use crate::client::{DatabaseClient, amend_request_options_for_lar}; use crate::model::request_options::Priority; use crate::model::transaction_options::ReadWrite; @@ -413,7 +413,6 @@ impl WriteOnlyTransaction { let client = self.client; let session_name = self.session_name.clone(); let previous_transaction_id = Arc::new(Mutex::new(Bytes::new())); - let channel_hint = client.next_channel_hint(); let affinity = Arc::new(TransactionAffinity::new_read_write()); let max_commit_delay = self.max_commit_delay; @@ -429,7 +428,7 @@ impl WriteOnlyTransaction { let previous_transaction_id = previous_transaction_id.clone(); let begin_gax_options = begin_gax_options.clone(); let commit_gax_options = commit_gax_options.clone(); - let _affinity = Arc::clone(&affinity); + let affinity = Arc::clone(&affinity); async move { let previous_id: Bytes = previous_transaction_id.lock().unwrap().clone(); @@ -449,8 +448,9 @@ impl WriteOnlyTransaction { .set_request_options(req_options.clone()) .set_or_clear_mutation_key(mutation_key.clone()); + let target = ChannelTarget::Affinity(&affinity); let tx = client - .begin_transaction(begin_req, begin_gax_options, channel_hint) + .begin_transaction(begin_req, begin_gax_options, target) .await?; *previous_transaction_id.lock().unwrap() = tx.id.clone(); @@ -465,7 +465,7 @@ impl WriteOnlyTransaction { ); let response = client - .commit(commit_req, commit_gax_options.clone(), channel_hint) + .commit(commit_req, commit_gax_options.clone(), target) .await?; // If a commit_response with a precommit_token is returned, then we need to @@ -482,7 +482,7 @@ impl WriteOnlyTransaction { ); client - .commit(retry_commit_req, commit_gax_options, channel_hint) + .commit(retry_commit_req, commit_gax_options, target) .await } else { Ok(response) @@ -545,7 +545,6 @@ impl WriteOnlyTransaction { .set_or_clear_max_commit_delay(self.max_commit_delay) .set_return_commit_stats(self.return_commit_stats); let client = self.client; - let channel_hint = client.next_channel_hint(); let is_emulator = client.is_emulator(); let action = || { @@ -555,7 +554,7 @@ impl WriteOnlyTransaction { async move { client - .commit(request, commit_gax_options, channel_hint) + .commit(request, commit_gax_options, ChannelTarget::Any) .await } }; diff --git a/src/spanner/tests/channel_pool_contention.rs b/src/spanner/tests/channel_pool_contention.rs new file mode 100644 index 0000000000..badf750d28 --- /dev/null +++ b/src/spanner/tests/channel_pool_contention.rs @@ -0,0 +1,1732 @@ +// 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. + +use gaxi::grpc::tonic::Response; +use google_cloud_auth::credentials::anonymous::Builder as Anonymous; +use google_cloud_spanner::channel_pool::{DynamicChannelPoolConfig, StaticChannelPoolConfig}; +use google_cloud_spanner::client::{DatabaseClient, Spanner, SpannerBuilderExt}; +use google_cloud_spanner::statement::Statement; +use google_cloud_spanner::transaction::{ReadWriteTransaction, TimestampBound}; +use google_cloud_test_macros::tokio_test_no_panics; +use prost_types::Value as ProtoValue; +use prost_types::value::Kind as ProtoValueKind; +use rand::random_range; +use serial_test::serial; +use spanner_grpc_mock::google::spanner::v1 as mock_v1; +use spanner_grpc_mock::google::spanner::v1::struct_type::Field; +use spanner_grpc_mock::google::spanner::v1::{StructType, Type, TypeCode}; +use spanner_grpc_mock::{MockSpanner, start}; +use std::collections::HashMap; +use std::net::SocketAddr; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; +use tokio::sync::{Mutex, Semaphore, mpsc}; +use tokio::task::JoinSet; +use tokio::time::sleep; + +/// Simulates GFE and SpanFE connection-level queuing and processing contention. +/// +/// For each physical TCP connection (identified by client socket address), allows at most +/// `max_parallel_per_channel` concurrent requests to execute in parallel. Any excess requests +/// are held in an asynchronous FIFO queue until an active request completes. +#[derive(Clone, Debug)] +pub(crate) struct ChannelContentionManager { + channels: Arc>>>, + max_parallel_per_channel: usize, + base_latency_min_micros: u64, + base_latency_noise_max_micros: u64, + total_queued_requests: Arc, + total_completed_requests: Arc, +} + +impl ChannelContentionManager { + pub(crate) fn new( + max_parallel_per_channel: usize, + base_latency_min_micros: u64, + base_latency_noise_max_micros: u64, + ) -> Self { + Self { + channels: Arc::new(Mutex::new(HashMap::new())), + max_parallel_per_channel, + base_latency_min_micros, + base_latency_noise_max_micros, + total_queued_requests: Arc::new(AtomicUsize::new(0)), + total_completed_requests: Arc::new(AtomicUsize::new(0)), + } + } + + /// Acquires a slot for execution on the connection, waiting in the FIFO queue if all + /// parallel slots are currently occupied. Simulates execution latency with random noise. + pub(crate) async fn execute_with_queue(&self, client_address: SocketAddr) { + let semaphore = { + let mut registry = self.channels.lock().await; + registry + .entry(client_address) + .or_insert_with(|| Arc::new(Semaphore::new(self.max_parallel_per_channel))) + .clone() + }; + + // If all permits are currently in use, the request waits in the semaphore's FIFO queue. + if semaphore.available_permits() == 0 { + self.total_queued_requests.fetch_add(1, Ordering::Relaxed); + } + + let _permit = semaphore + .acquire() + .await + .expect("semaphore must never be closed"); + + let noise_micros = if self.base_latency_noise_max_micros > 0 { + random_range(0..=self.base_latency_noise_max_micros) + } else { + 0 + }; + let execution_delay = Duration::from_micros(self.base_latency_min_micros + noise_micros); + sleep(execution_delay).await; + + self.total_completed_requests + .fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn total_queued_count(&self) -> usize { + self.total_queued_requests.load(Ordering::Relaxed) + } + + pub(crate) fn total_completed_count(&self) -> usize { + self.total_completed_requests.load(Ordering::Relaxed) + } +} + +async fn start_contention_server( + contention_manager: ChannelContentionManager, +) -> (String, tokio::task::JoinHandle<()>) { + let mut mock = MockSpanner::new(); + + mock.expect_create_session().returning(|_| { + Ok(Response::new(mock_v1::Session { + name: + "projects/test-project/instances/test-instance/databases/test-database/sessions/s1" + .to_string(), + multiplexed: true, + ..Default::default() + })) + }); + + let contention_manager_clone = contention_manager.clone(); + mock.expect_execute_streaming_sql() + .returning(move |request| { + let client_address = request + .remote_addr() + .expect("remote client address must be present on TCP transport"); + let is_begin = request + .get_ref() + .transaction + .as_ref() + .and_then(|selector| selector.selector.as_ref()) + .is_some_and(|selector| { + matches!(selector, mock_v1::transaction_selector::Selector::Begin(_)) + }); + let transaction_metadata = if is_begin { + Some(mock_v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + }) + } else { + None + }; + + let result_set = mock_v1::PartialResultSet { + metadata: Some(mock_v1::ResultSetMetadata { + row_type: Some(StructType { + fields: vec![ + Field { + name: "id".to_string(), + r#type: Some(Type { + code: TypeCode::Int64 as i32, + ..Default::default() + }), + }, + Field { + name: "value".to_string(), + r#type: Some(Type { + code: TypeCode::String as i32, + ..Default::default() + }), + }, + ], + }), + transaction: transaction_metadata, + undeclared_parameters: None, + }), + values: vec![ + ProtoValue { + kind: Some(ProtoValueKind::StringValue("1".to_string())), + }, + ProtoValue { + kind: Some(ProtoValueKind::StringValue("test-value".to_string())), + }, + ], + chunked_value: false, + resume_token: vec![1, 2, 3], + stats: None, + precommit_token: None, + cache_update: None, + last: true, + }; + + let (transmitter, receiver) = mpsc::channel(1); + let contention_manager = contention_manager_clone.clone(); + tokio::spawn(async move { + contention_manager.execute_with_queue(client_address).await; + let _ = transmitter.send(Ok(result_set)).await; + }); + Ok(Response::new(receiver)) + }); + + let contention_manager_clone = contention_manager.clone(); + mock.expect_execute_sql().returning(move |request| { + let client_address = request + .remote_addr() + .expect("remote client address must be present on TCP transport"); + let contention_manager = contention_manager_clone.clone(); + tokio::spawn(async move { + contention_manager.execute_with_queue(client_address).await; + }); + Ok(Response::new(mock_v1::ResultSet { + metadata: Some(mock_v1::ResultSetMetadata { + row_type: Some(StructType { fields: vec![] }), + ..Default::default() + }), + stats: Some(mock_v1::ResultSetStats { + row_count: Some(mock_v1::result_set_stats::RowCount::RowCountExact(1)), + ..Default::default() + }), + ..Default::default() + })) + }); + + let contention_manager_clone = contention_manager.clone(); + mock.expect_begin_transaction().returning(move |request| { + let client_address = request + .remote_addr() + .expect("remote client address must be present on TCP transport"); + let contention_manager = contention_manager_clone.clone(); + tokio::spawn(async move { + contention_manager.execute_with_queue(client_address).await; + }); + Ok(Response::new(mock_v1::Transaction { + id: vec![1, 2, 3], + ..Default::default() + })) + }); + + let contention_manager_clone = contention_manager.clone(); + mock.expect_commit().returning(move |request| { + let client_address = request + .remote_addr() + .expect("remote client address must be present on TCP transport"); + let contention_manager = contention_manager_clone.clone(); + tokio::spawn(async move { + contention_manager.execute_with_queue(client_address).await; + }); + Ok(Response::new(mock_v1::CommitResponse { + commit_timestamp: Some(prost_types::Timestamp { + seconds: 12345, + nanos: 0, + }), + ..Default::default() + })) + }); + + mock.expect_rollback().returning(|_| Ok(Response::new(()))); + + start("127.0.0.1:0", mock) + .await + .expect("mock server must start") +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn dynamic_channel_pool_scales_up_and_reduces_contention() { + // 1. Mock server setup: 4 parallel requests max per channel connection. + // Base latency: 1.5ms + random noise in [0.0ms, 0.5ms]. + let contention_manager = ChannelContentionManager::new(4, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + // 2. Client setup with DynamicChannelPoolConfig: + // Initial channels: 4, Max channels: 32, scale_up_cooldown: 10ms, doubling upon saturation. + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(4) + .with_min_channels(4) + .with_max_channels(32) + .with_min_rpc_per_channel(2.0) + .with_max_rpc_per_channel(8.0) + .with_scale_up_cooldown(Duration::from_millis(10)) + .with_max_scale_up_percent(100); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + assert_eq!( + spanner.active_channel_count(), + 4, + "initial channel count must be 4" + ); + + // 3. Workload generation: 48 concurrent worker tasks generating continuous load. + // On 4 initial channels, 48 requests creates 12 concurrent requests per channel, + // which exceeds max_parallel_per_channel (4) and max_rpc_per_channel (8.0), + // driving saturation and triggering automatic scale-up. + let total_workers = 48; + let is_stopped = Arc::new(AtomicBool::new(false)); + let mut join_set = JoinSet::new(); + + for _ in 0..total_workers { + let client_clone = database_client.clone(); + let stopped = is_stopped.clone(); + join_set.spawn(async move { + let mut latencies = Vec::new(); + while !stopped.load(Ordering::Relaxed) { + let start_time = Instant::now(); + let statement = Statement::builder("SELECT 1").build(); + let transaction = client_clone + .single_use() + .set_timestamp_bound(TimestampBound::exact_staleness(Duration::from_secs(15))) + .build(); + let mut result_set = transaction + .execute_query(statement) + .await + .expect("execute_query must succeed"); + while let Some(row) = result_set + .next() + .await + .transpose() + .expect("row must succeed") + { + let _: i64 = row.get(0_usize); + } + latencies.push(start_time.elapsed()); + } + latencies + }); + } + + // Monitor pool scaling with a deadline. + let scale_deadline = Instant::now() + Duration::from_secs(5); + let mut final_channels; + loop { + final_channels = spanner.active_channel_count(); + if final_channels > 4 || Instant::now() >= scale_deadline { + break; + } + sleep(Duration::from_millis(10)).await; + } + + // Stop all initial burst worker tasks and collect their latencies. + is_stopped.store(true, Ordering::Relaxed); + let mut burst_latencies = Vec::new(); + while let Some(join_result) = join_set.join_next().await { + let worker_latencies = join_result.expect("worker task must succeed"); + burst_latencies.extend(worker_latencies); + } + + // 4. Phase 2: With the scaled-up pool (>=8 channels), run 16 concurrent workers + // (2 requests per channel on average). Because concurrency per channel (2) <= 4 (server limit), + // queuing is completely eliminated and latency returns to base execution time (1.5ms - 2.0ms). + let post_scale_workers = 16; + let mut post_scale_latencies = Vec::new(); + for _ in 0..10 { + let mut wave_set = JoinSet::new(); + for _ in 0..post_scale_workers { + let client_clone = database_client.clone(); + wave_set.spawn(async move { + let start_time = Instant::now(); + let statement = Statement::builder("SELECT 1").build(); + let transaction = client_clone + .single_use() + .set_timestamp_bound(TimestampBound::exact_staleness(Duration::from_secs(15))) + .build(); + let mut result_set = transaction + .execute_query(statement) + .await + .expect("execute_query must succeed"); + while let Some(row) = result_set + .next() + .await + .transpose() + .expect("row must succeed") + { + let _: i64 = row.get(0_usize); + } + start_time.elapsed() + }); + } + while let Some(join_result) = wave_set.join_next().await { + post_scale_latencies.push(join_result.expect("worker task must succeed")); + } + } + + // 5. Assertions: + // A. The contention manager must have experienced queuing during the initial burst. + assert!( + contention_manager.total_queued_count() > 0, + "contention manager should have observed queued requests during initial burst" + ); + + // B. Total requests completed must be positive. + assert!( + contention_manager.total_completed_count() > 0, + "requests should have completed" + ); + + // C. The dynamic channel pool must have scaled up from 4 channels. + assert!( + final_channels > 4, + "dynamic channel pool must scale up beyond initial 4 channels under contention, got {final_channels}" + ); + + // D. Post-scale p50 latency must be significantly relieved compared to the contended phase. + post_scale_latencies.sort(); + let post_p50 = post_scale_latencies[post_scale_latencies.len() / 2]; + let post_p95 = post_scale_latencies[(post_scale_latencies.len() * 95) / 100]; + let post_p99 = post_scale_latencies[(post_scale_latencies.len() * 99) / 100]; + + burst_latencies.sort(); + let burst_p50 = burst_latencies[burst_latencies.len() / 2]; + let burst_p95 = burst_latencies[(burst_latencies.len() * 95) / 100]; + let burst_p99 = burst_latencies[(burst_latencies.len() * 99) / 100]; + + eprintln!( + "\n=======================================================\n[Dynamic Channel Pool Contention Benchmark Result]\nChannels scaled: 4 -> {final_channels}\nServer Queue Limit: 4 parallel requests per channel\nBase Latency: 1.5ms + [0.0ms, 0.5ms] noise\n-------------------------------------------------------\nInitial Contended Burst (4 channels, 48 concurrent workers):\n Requests: {}\n Queued on server: {}\n p50 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n-------------------------------------------------------\nPost Scale-Up Distribution ({final_channels} channels, {post_scale_workers} concurrent workers):\n Requests: {}\n p50 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n=======================================================", + burst_latencies.len(), + contention_manager.total_queued_count(), + burst_p50, + burst_p95, + burst_p99, + post_scale_latencies.len(), + post_p50, + post_p95, + post_p99 + ); + + assert!( + post_p50 < burst_p50, + "post-scale p50 latency ({post_p50:?}) should be lower than contended burst p50 ({burst_p50:?})" + ); +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn static_channel_pool_remains_at_fixed_size_under_contention() { + let contention_manager = ChannelContentionManager::new(4, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + let static_config = StaticChannelPoolConfig::new(4); + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(static_config) + .build() + .await + .expect("client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + assert_eq!( + spanner.active_channel_count(), + 4, + "static pool must have 4 channels" + ); + + let total_workers = 16; + let requests_per_worker = 4; + let mut join_set = JoinSet::new(); + + for _ in 0..total_workers { + let client_clone = database_client.clone(); + join_set.spawn(async move { + for _ in 0..requests_per_worker { + let start_time = Instant::now(); + let statement = Statement::builder("SELECT 1").build(); + let transaction = client_clone + .single_use() + .set_timestamp_bound(TimestampBound::exact_staleness(Duration::from_secs(15))) + .build(); + let mut result_set = transaction + .execute_query(statement) + .await + .expect("execute_query must succeed"); + while let Some(row) = result_set + .next() + .await + .transpose() + .expect("row must succeed") + { + let _: i64 = row.get(0_usize); + } + let _ = start_time.elapsed(); + } + }); + } + + while let Some(join_result) = join_set.join_next().await { + join_result.expect("worker task must succeed"); + } + + // Static pool must strictly remain at 4 channels + assert_eq!( + spanner.active_channel_count(), + 4, + "static channel pool must remain fixed at 4 channels" + ); +} + +#[derive(Debug, Clone)] +struct BenchmarkRunMetrics { + total_operations: usize, + queued_operations: usize, + mean_latency: Duration, + p50_latency: Duration, + p90_latency: Duration, + p95_latency: Duration, + p99_latency: Duration, + operations_over_15ms: usize, + operations_over_15ms_percent: f64, +} + +impl BenchmarkRunMetrics { + fn calculate(queued_operations: usize, mut latencies: Vec) -> Self { + latencies.sort(); + let total_operations = latencies.len(); + let total_duration: Duration = latencies.iter().sum(); + let mean_latency = if total_operations > 0 { + total_duration / total_operations as u32 + } else { + Duration::ZERO + }; + let p50_latency = if total_operations > 0 { + latencies[total_operations / 2] + } else { + Duration::ZERO + }; + let p90_latency = if total_operations > 0 { + latencies[(total_operations * 90) / 100] + } else { + Duration::ZERO + }; + let p95_latency = if total_operations > 0 { + latencies[(total_operations * 95) / 100] + } else { + Duration::ZERO + }; + let p99_latency = if total_operations > 0 { + latencies[(total_operations * 99) / 100] + } else { + Duration::ZERO + }; + let operations_over_15ms = latencies + .iter() + .filter(|&&lat| lat >= Duration::from_millis(15)) + .count(); + let operations_over_15ms_percent = if total_operations > 0 { + (operations_over_15ms as f64 / total_operations as f64) * 100.0 + } else { + 0.0 + }; + + Self { + total_operations, + queued_operations, + mean_latency, + p50_latency, + p90_latency, + p95_latency, + p99_latency, + operations_over_15ms, + operations_over_15ms_percent, + } + } +} + +async fn warm_up_database_client(database_client: &DatabaseClient) { + let statement = Statement::builder("SELECT id, value FROM test WHERE id = @id") + .add_param("id", 1i64) + .build(); + let transaction = database_client + .single_use() + .set_timestamp_bound(TimestampBound::exact_staleness(Duration::from_secs(15))) + .build(); + let mut result_set = transaction + .execute_query(statement) + .await + .expect("warm_up execute_query must succeed"); + while let Some(row) = result_set + .next() + .await + .transpose() + .expect("warm_up row must succeed") + { + let _: i64 = row.get(0_usize); + let _: String = row.get(1_usize); + } +} + +async fn run_database_client_point_select_workload( + database_client: &DatabaseClient, + worker_count: usize, + duration: Duration, +) -> Vec { + let is_stopped = Arc::new(AtomicBool::new(false)); + let mut join_set = JoinSet::new(); + + for _ in 0..worker_count { + let client_clone = database_client.clone(); + let stopped = is_stopped.clone(); + join_set.spawn(async move { + let mut worker_latencies = Vec::new(); + while !stopped.load(Ordering::Relaxed) { + let start_time = Instant::now(); + let statement = Statement::builder("SELECT id, value FROM test WHERE id = @id") + .add_param("id", 1i64) + .build(); + let transaction = client_clone + .single_use() + .set_timestamp_bound(TimestampBound::exact_staleness(Duration::from_secs(15))) + .build(); + let mut result_set = transaction + .execute_query(statement) + .await + .expect("execute_query must succeed"); + while let Some(row) = result_set + .next() + .await + .transpose() + .expect("row must succeed") + { + let _: i64 = row.get(0_usize); + let _: String = row.get(1_usize); + } + worker_latencies.push(start_time.elapsed()); + } + worker_latencies + }); + } + + sleep(duration).await; + is_stopped.store(true, Ordering::Relaxed); + + let mut all_latencies = Vec::new(); + while let Some(join_result) = join_set.join_next().await { + let worker_latencies = join_result.expect("worker task must succeed"); + all_latencies.extend(worker_latencies); + } + all_latencies +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_spiky_workload_database_client_under_default_dynamic_pool() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(4) + .with_min_channels(4) + .with_max_channels(256) + .with_min_rpc_per_channel(1.0) + .with_max_rpc_per_channel(2.0) + .with_error_penalty_step(1) + .with_scale_up_cooldown(Duration::from_secs(4)) + .with_max_scale_up_percent(100); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("spanner client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + assert_eq!( + spanner.active_channel_count(), + 4, + "initial dynamic channel pool count must be 4" + ); + + warm_up_database_client(&database_client).await; + + // Phase 1: Calm/Normal State (2 concurrent workers for 300ms) + let calm_latencies = + run_database_client_point_select_workload(&database_client, 2, Duration::from_millis(300)) + .await; + + let calm_queued_count = contention_manager.total_queued_count(); + let calm_metrics = BenchmarkRunMetrics::calculate(calm_queued_count, calm_latencies); + + assert_eq!( + calm_queued_count, 0, + "calm phase must not cause server queuing" + ); + assert!( + calm_metrics.p50_latency < Duration::from_millis(6), + "calm p50 latency must be near baseline (<6ms)" + ); + + // Phase 2: Spiky Burst State (12 concurrent workers for 800ms) + // With 4 channels and 2 permits per channel, 12 workers creates 3 requests per channel (> 2 permits). + let burst_start_queued = contention_manager.total_queued_count(); + let burst_latencies = + run_database_client_point_select_workload(&database_client, 12, Duration::from_millis(800)) + .await; + + let burst_queued_count = contention_manager.total_queued_count() - burst_start_queued; + let final_channels = spanner.active_channel_count(); + let burst_metrics = BenchmarkRunMetrics::calculate(burst_queued_count, burst_latencies); + + eprintln!( + "\n===============================================================================\n[Replication 1: Dynamic Pool Under DatabaseClient Point-Select Spiky Load]\nChannels: 4 -> {final_channels} (Scales to 8 at onset, then blocked by 4s cooldown!)\nServer Permit Limit: 2 parallel requests per connection\nBase Latency: 1.5ms + [0.0ms, 0.5ms] noise\n-------------------------------------------------------------------------------\nNormal State (2 workers, 300ms):\n Operations: {}\n Server Queued: {}\n p50 latency: {:?}\n p95 latency: {:?}\nBurst State (12 workers, 800ms):\n Operations: {}\n Server Queued: {}\n Mean latency: {:?}\n p50 latency: {:?}\n p90 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n Operations > 15ms: {} ({:.2}%)\n===============================================================================", + calm_metrics.total_operations, + calm_metrics.queued_operations, + calm_metrics.p50_latency, + calm_metrics.p95_latency, + burst_metrics.total_operations, + burst_metrics.queued_operations, + burst_metrics.mean_latency, + burst_metrics.p50_latency, + burst_metrics.p90_latency, + burst_metrics.p95_latency, + burst_metrics.p99_latency, + burst_metrics.operations_over_15ms, + burst_metrics.operations_over_15ms_percent + ); + + // ROOT CAUSE 2 VERIFICATION: + // With DatabaseClient now wired to dynamic channel pool, the pool doubles from 4 to 8 at burst onset, + // but the 4-second scale-up cooldown prevents scaling any further (e.g. to 16, 32, 64) during the 800ms burst. + assert_eq!( + final_channels, 8, + "Dynamic pool must scale from 4 to 8 at onset, but remain throttled at 8 due to 4s cooldown" + ); + + assert!( + burst_queued_count > 0, + "Burst must produce server-side queuing on the 4 overloaded connections" + ); + assert!( + burst_metrics.p95_latency > Duration::from_millis(5), + "p95 latency must be degraded due to queuing behind the 4 channels" + ); +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_spiky_workload_database_client_under_64_channel_static_pool() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + // Static 64 channels (matching benchmark Configuration C). + let static_config = StaticChannelPoolConfig::new(64); + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(static_config) + .build() + .await + .expect("spanner client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + assert_eq!( + spanner.active_channel_count(), + 64, + "static 64-channel pool must have 64 channels" + ); + + warm_up_database_client(&database_client).await; + + let calm_latencies = + run_database_client_point_select_workload(&database_client, 2, Duration::from_millis(300)) + .await; + let calm_queued_count = contention_manager.total_queued_count(); + let calm_metrics = BenchmarkRunMetrics::calculate(calm_queued_count, calm_latencies); + + let burst_start_queued = contention_manager.total_queued_count(); + let burst_latencies = + run_database_client_point_select_workload(&database_client, 12, Duration::from_millis(800)) + .await; + + let burst_queued_count = contention_manager.total_queued_count() - burst_start_queued; + let burst_metrics = BenchmarkRunMetrics::calculate(burst_queued_count, burst_latencies); + + eprintln!( + "\n===============================================================================\n[Replication 2: Static 64 Channels Under DatabaseClient Point-Select Spiky Load]\nChannels: 64 (Fixed)\nServer Permit Limit: 2 parallel requests per connection\nBase Latency: 1.5ms + [0.0ms, 0.5ms] noise\n-------------------------------------------------------------------------------\nNormal State (2 workers, 300ms):\n Operations: {}\n Server Queued: {}\n p50 latency: {:?}\nBurst State (12 workers, 800ms):\n Operations: {}\n Server Queued: {}\n Mean latency: {:?}\n p50 latency: {:?}\n p90 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n Operations > 15ms: {} ({:.2}%)\n===============================================================================", + calm_metrics.total_operations, + calm_metrics.queued_operations, + calm_metrics.p50_latency, + burst_metrics.total_operations, + burst_metrics.queued_operations, + burst_metrics.mean_latency, + burst_metrics.p50_latency, + burst_metrics.p90_latency, + burst_metrics.p95_latency, + burst_metrics.p99_latency, + burst_metrics.operations_over_15ms, + burst_metrics.operations_over_15ms_percent + ); + + // With 64 channels, 12 concurrent workers distribute across separate channels (<= 1 req/channel). + // Server-side queuing is completely zero, exactly reproducing Configuration C in the benchmark! + assert_eq!( + burst_queued_count, 0, + "64 channels must experience zero server queuing during the 12-worker burst" + ); + assert!( + burst_metrics.p95_latency < Duration::from_millis(50), + "P95 latency with 64 channels must remain bounded in debug mode (<50ms)" + ); +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_dynamic_pool_4s_cooldown_bottleneck_under_short_burst() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(4) + .with_min_channels(4) + .with_max_channels(64) + .with_min_rpc_per_channel(1.0) + .with_max_rpc_per_channel(2.0) + .with_error_penalty_step(1) + .with_scale_up_cooldown(Duration::from_secs(4)) + .with_max_scale_up_percent(100); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("spanner build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + assert_eq!( + spanner.active_channel_count(), + 4, + "initial channels must be 4" + ); + + let total_workers = 16; + let is_stopped = Arc::new(AtomicBool::new(false)); + let mut join_set = JoinSet::new(); + + for _ in 0..total_workers { + let client_clone = database_client.clone(); + let stopped = is_stopped.clone(); + join_set.spawn(async move { + let mut latencies = Vec::new(); + while !stopped.load(Ordering::Relaxed) { + let start_time = Instant::now(); + let statement = Statement::builder("SELECT 1").build(); + let transaction = client_clone + .single_use() + .set_timestamp_bound(TimestampBound::exact_staleness(Duration::from_secs(15))) + .build(); + let mut result_set = transaction + .execute_query(statement) + .await + .expect("execute_query must succeed"); + while let Some(row) = result_set + .next() + .await + .transpose() + .expect("row must succeed") + { + let _: i64 = row.get(0_usize); + } + latencies.push(start_time.elapsed()); + } + latencies + }); + } + + // Run burst for 800ms (shorter than the 4s cooldown) + sleep(Duration::from_millis(800)).await; + is_stopped.store(true, Ordering::Relaxed); + + let mut burst_latencies = Vec::new(); + while let Some(join_result) = join_set.join_next().await { + burst_latencies.extend(join_result.expect("worker task must succeed")); + } + + let final_channels = spanner.active_channel_count(); + let metrics = + BenchmarkRunMetrics::calculate(contention_manager.total_queued_count(), burst_latencies); + + eprintln!( + "\n===============================================================================\n[Replication 3: Direct Dynamic Pool With 4s Cooldown Under 800ms Burst]\nChannels: 4 -> {final_channels} (Doubled once at t=0, then throttled by 4s cooldown!)\nServer Permit Limit: 2 parallel requests per connection\nBase Latency: 1.5ms + [0.0ms, 0.5ms] noise\n-------------------------------------------------------------------------------\nBurst State (16 workers, 800ms):\n Operations: {}\n Server Queued: {}\n Mean latency: {:?}\n p50 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n Operations > 15ms: {} ({:.2}%)\n===============================================================================", + metrics.total_operations, + metrics.queued_operations, + metrics.mean_latency, + metrics.p50_latency, + metrics.p95_latency, + metrics.p99_latency, + metrics.operations_over_15ms, + metrics.operations_over_15ms_percent + ); + + // At t=0, the pool detected saturation and scaled 4 -> 8. + // But because scale_up_cooldown is 4 seconds, during the entire 800ms burst, + // it was legally prohibited from scaling further! + // So it was stuck at 8 channels, where 16 workers create 2 req/channel. + assert_eq!( + final_channels, 8, + "Pool must scale from 4 to 8 at burst onset, but remain throttled at 8 due to 4s cooldown" + ); +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_dynamic_pool_fast_cooldown_recovers_under_burst() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + // Dynamic pool with responsive 50ms cooldown and 100% geometric scale-up. + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(4) + .with_min_channels(4) + .with_max_channels(64) + .with_min_rpc_per_channel(1.0) + .with_max_rpc_per_channel(2.0) + .with_error_penalty_step(1) + .with_scale_up_cooldown(Duration::from_millis(50)) + .with_max_scale_up_percent(100); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("spanner build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + assert_eq!( + spanner.active_channel_count(), + 4, + "initial channels must be 4" + ); + + let total_workers = 24; + let is_stopped = Arc::new(AtomicBool::new(false)); + let mut join_set = JoinSet::new(); + + for _ in 0..total_workers { + let client_clone = database_client.clone(); + let stopped = is_stopped.clone(); + join_set.spawn(async move { + let mut latencies = Vec::new(); + while !stopped.load(Ordering::Relaxed) { + let start_time = Instant::now(); + let statement = Statement::builder("SELECT 1").build(); + let transaction = client_clone + .single_use() + .set_timestamp_bound(TimestampBound::exact_staleness(Duration::from_secs(15))) + .build(); + let mut result_set = transaction + .execute_query(statement) + .await + .expect("execute_query must succeed"); + while let Some(row) = result_set + .next() + .await + .transpose() + .expect("row must succeed") + { + let _: i64 = row.get(0_usize); + } + latencies.push(start_time.elapsed()); + } + latencies + }); + } + + // Run burst for 800ms (ample time for 50ms cooldowns: 4 -> 8 -> 16 -> 32) + sleep(Duration::from_millis(800)).await; + is_stopped.store(true, Ordering::Relaxed); + + let mut burst_latencies = Vec::new(); + while let Some(join_result) = join_set.join_next().await { + burst_latencies.extend(join_result.expect("worker task must succeed")); + } + + let final_channels = spanner.active_channel_count(); + let metrics = + BenchmarkRunMetrics::calculate(contention_manager.total_queued_count(), burst_latencies); + + eprintln!( + "\n===============================================================================\n[Replication 4: Dynamic Pool With Responsive 50ms Cooldown Under 800ms Burst]\nChannels: 4 -> {final_channels} (Scaled multiple steps rapidly!)\nServer Permit Limit: 2 parallel requests per connection\nBase Latency: 1.5ms + [0.0ms, 0.5ms] noise\n-------------------------------------------------------------------------------\nBurst State (24 workers, 800ms):\n Operations: {}\n Server Queued: {}\n Mean latency: {:?}\n p50 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n Operations > 15ms: {} ({:.2}%)\n===============================================================================", + metrics.total_operations, + metrics.queued_operations, + metrics.mean_latency, + metrics.p50_latency, + metrics.p95_latency, + metrics.p99_latency, + metrics.operations_over_15ms, + metrics.operations_over_15ms_percent + ); + + // With a 50ms cooldown and 24 workers, the pool scales beyond 8 channels to at least 16 channels. + assert!( + final_channels >= 16, + "Dynamic pool with 50ms cooldown must scale up to at least 16 channels, got {final_channels}" + ); +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_database_client_scales_dynamically_with_fast_cooldown() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(4) + .with_min_channels(4) + .with_max_channels(64) + .with_min_rpc_per_channel(1.0) + .with_max_rpc_per_channel(2.0) + .with_error_penalty_step(1) + .with_scale_up_cooldown(Duration::from_millis(50)) + .with_max_scale_up_percent(100); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("spanner client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + warm_up_database_client(&database_client).await; + + assert_eq!( + spanner.active_channel_count(), + 4, + "initial dynamic channel pool count must be 4" + ); + + // Run burst using DatabaseClient with 24 concurrent workers for 800ms + let burst_latencies = + run_database_client_point_select_workload(&database_client, 24, Duration::from_millis(800)) + .await; + + let final_channels = spanner.active_channel_count(); + let queued_count = contention_manager.total_queued_count(); + let burst_metrics = BenchmarkRunMetrics::calculate(queued_count, burst_latencies); + + eprintln!( + "\n===============================================================================\n[Replication 5: DatabaseClient Under Dynamic Pool With 50ms Cooldown]\nChannels: 4 -> {final_channels} (DatabaseClient successfully drives dynamic scaling!)\nServer Permit Limit: 2 parallel requests per connection\nBase Latency: 1.5ms + [0.0ms, 0.5ms] noise\n-------------------------------------------------------------------------------\nBurst State (24 workers, 800ms):\n Operations: {}\n Server Queued: {}\n Mean latency: {:?}\n p50 latency: {:?}\n p90 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n Operations > 15ms: {} ({:.2}%)\n===============================================================================", + burst_metrics.total_operations, + burst_metrics.queued_operations, + burst_metrics.mean_latency, + burst_metrics.p50_latency, + burst_metrics.p90_latency, + burst_metrics.p95_latency, + burst_metrics.p99_latency, + burst_metrics.operations_over_15ms, + burst_metrics.operations_over_15ms_percent + ); + + assert!( + final_channels >= 16, + "DatabaseClient queries must drive dynamic channel pool scaling beyond 8 to at least 16 channels, got {final_channels}" + ); +} + +async fn warm_up_select_update(database_client: &DatabaseClient) { + let runner = database_client + .read_write_transaction() + .build() + .await + .expect("warm_up build read_write_transaction"); + + runner + .run(|transaction: ReadWriteTransaction| async move { + let statement = Statement::builder("SELECT id FROM test WHERE id = @id") + .add_param("id", 1i64) + .build(); + let mut result_set = transaction.execute_query(statement).await?; + let _ = result_set.next().await; + drop(result_set); + + let update_statement = + Statement::builder("UPDATE test SET value = @value WHERE id = @id") + .add_param("value", "warmup-value") + .add_param("id", 1i64) + .build(); + transaction.execute_update(update_statement).await?; + Ok(()) + }) + .await + .expect("warm_up select_update must succeed"); +} + +async fn run_database_client_select_update_workload( + database_client: &DatabaseClient, + worker_count: usize, + duration: Duration, +) -> Vec { + let is_stopped = Arc::new(AtomicBool::new(false)); + let mut join_set = JoinSet::new(); + + for _ in 0..worker_count { + let client_clone = database_client.clone(); + let stopped = is_stopped.clone(); + join_set.spawn(async move { + let mut worker_latencies = Vec::new(); + while !stopped.load(Ordering::Relaxed) { + let start_time = Instant::now(); + let runner = client_clone + .read_write_transaction() + .build() + .await + .expect("build read_write_transaction runner"); + + runner + .run(|transaction: ReadWriteTransaction| async move { + let random_id = random_range(1i64..=1_000_000i64); + let select_statement = + Statement::builder("SELECT id FROM test WHERE id = @id") + .add_param("id", random_id) + .build(); + + let mut result_set = transaction.execute_query(select_statement).await?; + let exists = result_set.next().await.transpose()?.is_some(); + drop(result_set); + + if exists { + let update_statement = + Statement::builder("UPDATE test SET value = @value WHERE id = @id") + .add_param("value", "updated-test-value") + .add_param("id", random_id) + .build(); + transaction.execute_update(update_statement).await?; + } else { + let insert_statement = Statement::builder( + "INSERT INTO test (id, value) VALUES (@id, @value)", + ) + .add_param("id", random_id) + .add_param("value", "inserted-test-value") + .build(); + transaction.execute_update(insert_statement).await?; + } + + Ok(()) + }) + .await + .expect("select_update transaction must succeed"); + + worker_latencies.push(start_time.elapsed()); + } + worker_latencies + }); + } + + sleep(duration).await; + is_stopped.store(true, Ordering::Relaxed); + + let mut all_latencies = Vec::new(); + while let Some(join_result) = join_set.join_next().await { + let worker_latencies = join_result.expect("worker task must succeed"); + all_latencies.extend(worker_latencies); + } + all_latencies +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_select_update_database_client_under_default_dynamic_pool() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(4) + .with_min_channels(4) + .with_max_channels(256) + .with_min_rpc_per_channel(1.0) + .with_max_rpc_per_channel(2.0) + .with_error_penalty_step(1) + .with_scale_up_cooldown(Duration::from_secs(4)) + .with_max_scale_up_percent(100); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("spanner client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + warm_up_select_update(&database_client).await; + + assert_eq!( + spanner.active_channel_count(), + 4, + "initial dynamic channel pool count must be 4" + ); + + // Run burst of select-then-update transactions with 24 concurrent workers for 800ms + let burst_latencies = run_database_client_select_update_workload( + &database_client, + 24, + Duration::from_millis(800), + ) + .await; + + let final_channels = spanner.active_channel_count(); + let queued_count = contention_manager.total_queued_count(); + let burst_metrics = BenchmarkRunMetrics::calculate(queued_count, burst_latencies); + + eprintln!( + "\n===============================================================================\n[Replication 6: Select-Then-Update DatabaseClient Under Default Dynamic Pool (4s Cooldown)]\nChannels: 4 -> {final_channels} (Default 4s cooldown limits scale-up during 800ms burst)\nServer Permit Limit: 2 parallel requests per connection\nBase Latency: 1.5ms + [0.0ms, 0.5ms] noise per RPC (3 RPCs per txn: select, update, commit)\n-------------------------------------------------------------------------------\nBurst State (24 workers, 800ms):\n Transactions: {}\n Server Queued Requests: {}\n Mean latency: {:?}\n p50 latency: {:?}\n p90 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n Transactions > 15ms: {} ({:.2}%)\n===============================================================================", + burst_metrics.total_operations, + burst_metrics.queued_operations, + burst_metrics.mean_latency, + burst_metrics.p50_latency, + burst_metrics.p90_latency, + burst_metrics.p95_latency, + burst_metrics.p99_latency, + burst_metrics.operations_over_15ms, + burst_metrics.operations_over_15ms_percent + ); + + assert!( + burst_metrics.total_operations > 0, + "Expected at least one select-update transaction to execute" + ); +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_select_update_database_client_under_64_channel_static_pool() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + let static_config = StaticChannelPoolConfig::new(64); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(static_config) + .build() + .await + .expect("spanner client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + warm_up_select_update(&database_client).await; + + assert_eq!( + spanner.active_channel_count(), + 64, + "static pool must have 64 channels" + ); + + // Run burst of select-then-update transactions with 24 concurrent workers for 800ms + let burst_latencies = run_database_client_select_update_workload( + &database_client, + 24, + Duration::from_millis(800), + ) + .await; + + let queued_count = contention_manager.total_queued_count(); + let burst_metrics = BenchmarkRunMetrics::calculate(queued_count, burst_latencies); + + eprintln!( + "\n===============================================================================\n[Replication 7: Select-Then-Update DatabaseClient Under Static 64-Channel Pool]\nChannels: 64 (Static, abundant channels distribute 24 transactions with minimal queuing)\nServer Permit Limit: 2 parallel requests per connection\nBase Latency: 1.5ms + [0.0ms, 0.5ms] noise per RPC (3 RPCs per txn: select, update, commit)\n-------------------------------------------------------------------------------\nBurst State (24 workers, 800ms):\n Transactions: {}\n Server Queued Requests: {}\n Mean latency: {:?}\n p50 latency: {:?}\n p90 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n Transactions > 15ms: {} ({:.2}%)\n===============================================================================", + burst_metrics.total_operations, + burst_metrics.queued_operations, + burst_metrics.mean_latency, + burst_metrics.p50_latency, + burst_metrics.p90_latency, + burst_metrics.p95_latency, + burst_metrics.p99_latency, + burst_metrics.operations_over_15ms, + burst_metrics.operations_over_15ms_percent + ); + + assert!( + burst_metrics.total_operations > 0, + "Expected at least one select-update transaction to execute" + ); +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_select_update_database_client_scales_dynamically_with_fast_cooldown() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(4) + .with_min_channels(4) + .with_max_channels(64) + .with_min_rpc_per_channel(1.0) + .with_max_rpc_per_channel(2.0) + .with_error_penalty_step(1) + .with_scale_up_cooldown(Duration::from_millis(50)) + .with_max_scale_up_percent(100); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("spanner client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + warm_up_select_update(&database_client).await; + + assert_eq!( + spanner.active_channel_count(), + 4, + "initial dynamic channel pool count must be 4" + ); + + // Run burst of select-then-update transactions with 24 concurrent workers for 800ms + let burst_latencies = run_database_client_select_update_workload( + &database_client, + 24, + Duration::from_millis(800), + ) + .await; + + let final_channels = spanner.active_channel_count(); + let queued_count = contention_manager.total_queued_count(); + let burst_metrics = BenchmarkRunMetrics::calculate(queued_count, burst_latencies); + + eprintln!( + "\n===============================================================================\n[Replication 8: Select-Then-Update DatabaseClient Under Dynamic Pool With 50ms Cooldown]\nChannels: 4 -> {final_channels} (Select-Update transactions successfully drive dynamic scaling!)\nServer Permit Limit: 2 parallel requests per connection\nBase Latency: 1.5ms + [0.0ms, 0.5ms] noise per RPC (3 RPCs per txn: select, update, commit)\n-------------------------------------------------------------------------------\nBurst State (24 workers, 800ms):\n Transactions: {}\n Server Queued Requests: {}\n Mean latency: {:?}\n p50 latency: {:?}\n p90 latency: {:?}\n p95 latency: {:?}\n p99 latency: {:?}\n Transactions > 15ms: {} ({:.2}%)\n===============================================================================", + burst_metrics.total_operations, + burst_metrics.queued_operations, + burst_metrics.mean_latency, + burst_metrics.p50_latency, + burst_metrics.p90_latency, + burst_metrics.p95_latency, + burst_metrics.p99_latency, + burst_metrics.operations_over_15ms, + burst_metrics.operations_over_15ms_percent + ); + + assert!( + final_channels >= 16, + "Select-Update transactions must drive dynamic channel pool scaling beyond 8 to at least 16 channels, got {final_channels}" + ); +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_select_update_database_client_2min_spiky_default_dynamic_pool() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + // Default dynamic pool configuration with 4s scale_up_cooldown + let dynamic_config = DynamicChannelPoolConfig::new() + .with_initial_channels(4) + .with_min_channels(4) + .with_max_channels(64) + .with_min_rpc_per_channel(1.0) + .with_max_rpc_per_channel(2.0) + .with_error_penalty_step(1) + .with_scale_up_cooldown(Duration::from_secs(4)) + .with_max_scale_up_percent(100); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(dynamic_config) + .build() + .await + .expect("spanner client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + warm_up_select_update(&database_client).await; + + assert_eq!( + spanner.active_channel_count(), + 4, + "initial dynamic channel pool count must be 4" + ); + + // 2-Minute Spiky Workload: 6 Cycles of [10s calm (2 workers), 10s burst (24 workers)] + let mut all_calm_latencies = Vec::new(); + let mut all_burst_latencies = Vec::new(); + let mut steady_state_burst_latencies = Vec::new(); + + eprintln!("\n==============================================================================="); + eprintln!("[Replication 9: 2-Minute Spiky Select-Then-Update Workload (Default 4s Cooldown)]"); + eprintln!("Load Pattern: 6 Cycles of [10s Calm (2 workers) + 10s Burst (24 workers)] = 120s"); + eprintln!("Dynamic Pool: Initial 4, Min 4, Max 64, 4s Scale-Up Cooldown, 100% Geometric Scale"); + eprintln!( + "Server Limits: 2 permits per channel, 1.5ms base + [0.0, 0.5ms] jitter (3 RPCs/txn)" + ); + eprintln!("==============================================================================="); + + for cycle in 1..=6 { + // --- Calm Phase (10s, 2 workers) --- + let calm_queued_start = contention_manager.total_queued_count(); + let calm_latencies = run_database_client_select_update_workload( + &database_client, + 2, + Duration::from_secs(10), + ) + .await; + let calm_queued = contention_manager.total_queued_count() - calm_queued_start; + let calm_channels = spanner.active_channel_count(); + let calm_metrics = BenchmarkRunMetrics::calculate(calm_queued, calm_latencies.clone()); + all_calm_latencies.extend(calm_latencies); + + eprintln!( + "Cycle {} Calm (10s, 2 workers): Txns: {:>5} | Queued: {:>4} | Channels: {:>2} | p50: {:>8.2?} | p95: {:>8.2?}", + cycle, + calm_metrics.total_operations, + calm_metrics.queued_operations, + calm_channels, + calm_metrics.p50_latency, + calm_metrics.p95_latency, + ); + + // --- Burst Phase (10s, 24 workers) --- + let burst_queued_start = contention_manager.total_queued_count(); + let burst_latencies = run_database_client_select_update_workload( + &database_client, + 24, + Duration::from_secs(10), + ) + .await; + let burst_queued = contention_manager.total_queued_count() - burst_queued_start; + let burst_channels = spanner.active_channel_count(); + let burst_metrics = BenchmarkRunMetrics::calculate(burst_queued, burst_latencies.clone()); + all_burst_latencies.extend(burst_latencies.clone()); + + if cycle > 1 { + steady_state_burst_latencies.extend(burst_latencies); + } + + eprintln!( + "Cycle {} Burst (10s, 24 workers): Txns: {:>5} | Queued: {:>4} | Channels: {:>2} | p50: {:>8.2?} | p90: {:>8.2?} | p95: {:>8.2?} | p99: {:>8.2?} | >15ms: {:>5.2}%", + cycle, + burst_metrics.total_operations, + burst_metrics.queued_operations, + burst_channels, + burst_metrics.p50_latency, + burst_metrics.p90_latency, + burst_metrics.p95_latency, + burst_metrics.p99_latency, + burst_metrics.operations_over_15ms_percent, + ); + } + + let final_channels = spanner.active_channel_count(); + let total_queued = contention_manager.total_queued_count(); + let overall_burst_metrics = BenchmarkRunMetrics::calculate(total_queued, all_burst_latencies); + let steady_burst_metrics = BenchmarkRunMetrics::calculate(0, steady_state_burst_latencies); + + eprintln!("\n-------------------------------------------------------------------------------"); + eprintln!("[2-Minute Aggregate Summary: Dynamic Pool (4s Cooldown)]"); + eprintln!("Final Channels: {}", final_channels); + eprintln!( + "Total Server Queued Requests across 2 minutes: {}", + total_queued + ); + eprintln!("\nSteady-State Bursts (Cycles 2-6 Combined, Channels = 64):"); + eprintln!( + " Transactions: {}", + steady_burst_metrics.total_operations + ); + eprintln!( + " Mean latency: {:?}", + steady_burst_metrics.mean_latency + ); + eprintln!( + " p50 latency: {:?}", + steady_burst_metrics.p50_latency + ); + eprintln!( + " p90 latency: {:?}", + steady_burst_metrics.p90_latency + ); + eprintln!( + " p95 latency: {:?}", + steady_burst_metrics.p95_latency + ); + eprintln!( + " p99 latency: {:?}", + steady_burst_metrics.p99_latency + ); + eprintln!( + " Transactions >15ms: {} ({:.2}%)", + steady_burst_metrics.operations_over_15ms, + steady_burst_metrics.operations_over_15ms_percent + ); + eprintln!("\nOverall All Bursts (Cycles 1-6 Combined, including Cold Ramp-up):"); + eprintln!( + " Transactions: {}", + overall_burst_metrics.total_operations + ); + eprintln!( + " Mean latency: {:?}", + overall_burst_metrics.mean_latency + ); + eprintln!( + " p50 latency: {:?}", + overall_burst_metrics.p50_latency + ); + eprintln!( + " p90 latency: {:?}", + overall_burst_metrics.p90_latency + ); + eprintln!( + " p95 latency: {:?}", + overall_burst_metrics.p95_latency + ); + eprintln!( + " p99 latency: {:?}", + overall_burst_metrics.p99_latency + ); + eprintln!( + " Transactions >15ms: {} ({:.2}%)", + overall_burst_metrics.operations_over_15ms, + overall_burst_metrics.operations_over_15ms_percent + ); + eprintln!("==============================================================================="); + + assert!( + final_channels >= 32, + "Pool must scale to at least 32 channels during 2-minute spiky workload with 10s bursts, got {final_channels}" + ); +} + +#[tokio_test_no_panics] +#[serial] +#[ignore] +async fn replication_select_update_database_client_2min_spiky_64_channel_static_pool() { + let contention_manager = ChannelContentionManager::new(2, 1500, 500); + let (address, _server) = start_contention_server(contention_manager.clone()).await; + + let static_config = StaticChannelPoolConfig::new(64); + + let spanner = Spanner::builder() + .with_endpoint(address) + .with_credentials(Anonymous::new().build()) + .with_channel_pool(static_config) + .build() + .await + .expect("spanner client build must succeed"); + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-database") + .build() + .await + .expect("database_client build must succeed"); + + warm_up_select_update(&database_client).await; + + assert_eq!( + spanner.active_channel_count(), + 64, + "static channel pool count must be 64" + ); + + // 2-Minute Spiky Workload: 6 Cycles of [10s calm (2 workers), 10s burst (24 workers)] + let mut all_burst_latencies = Vec::new(); + + eprintln!("\n==============================================================================="); + eprintln!("[Replication 10: 2-Minute Spiky Select-Then-Update Workload (Static 64 Channels)]"); + eprintln!("Load Pattern: 6 Cycles of [10s Calm (2 workers) + 10s Burst (24 workers)] = 120s"); + eprintln!("Static Pool: 64 Channels fixed (abundant channels)"); + eprintln!( + "Server Limits: 2 permits per channel, 1.5ms base + [0.0, 0.5ms] jitter (3 RPCs/txn)" + ); + eprintln!("==============================================================================="); + + for cycle in 1..=6 { + // --- Calm Phase (10s, 2 workers) --- + let calm_queued_start = contention_manager.total_queued_count(); + let calm_latencies = run_database_client_select_update_workload( + &database_client, + 2, + Duration::from_secs(10), + ) + .await; + let calm_queued = contention_manager.total_queued_count() - calm_queued_start; + let calm_metrics = BenchmarkRunMetrics::calculate(calm_queued, calm_latencies); + + eprintln!( + "Cycle {} Calm (10s, 2 workers): Txns: {:>5} | Queued: {:>4} | Channels: 64 | p50: {:>8.2?} | p95: {:>8.2?}", + cycle, + calm_metrics.total_operations, + calm_metrics.queued_operations, + calm_metrics.p50_latency, + calm_metrics.p95_latency, + ); + + // --- Burst Phase (10s, 24 workers) --- + let burst_queued_start = contention_manager.total_queued_count(); + let burst_latencies = run_database_client_select_update_workload( + &database_client, + 24, + Duration::from_secs(10), + ) + .await; + let burst_queued = contention_manager.total_queued_count() - burst_queued_start; + let burst_metrics = BenchmarkRunMetrics::calculate(burst_queued, burst_latencies.clone()); + all_burst_latencies.extend(burst_latencies); + + eprintln!( + "Cycle {} Burst (10s, 24 workers): Txns: {:>5} | Queued: {:>4} | Channels: 64 | p50: {:>8.2?} | p90: {:>8.2?} | p95: {:>8.2?} | p99: {:>8.2?} | >15ms: {:>5.2}%", + cycle, + burst_metrics.total_operations, + burst_metrics.queued_operations, + burst_metrics.p50_latency, + burst_metrics.p90_latency, + burst_metrics.p95_latency, + burst_metrics.p99_latency, + burst_metrics.operations_over_15ms_percent, + ); + } + + let total_queued = contention_manager.total_queued_count(); + let overall_burst_metrics = BenchmarkRunMetrics::calculate(total_queued, all_burst_latencies); + + eprintln!("\n-------------------------------------------------------------------------------"); + eprintln!("[Static 64-Channel Pool 2-Minute Aggregate Summary]"); + eprintln!("Channels: 64"); + eprintln!( + "Total Server Queued Requests across 2 minutes: {}", + total_queued + ); + eprintln!("All Bursts Combined:"); + eprintln!( + " Transactions: {}", + overall_burst_metrics.total_operations + ); + eprintln!( + " Mean latency: {:?}", + overall_burst_metrics.mean_latency + ); + eprintln!( + " p50 latency: {:?}", + overall_burst_metrics.p50_latency + ); + eprintln!( + " p90 latency: {:?}", + overall_burst_metrics.p90_latency + ); + eprintln!( + " p95 latency: {:?}", + overall_burst_metrics.p95_latency + ); + eprintln!( + " p99 latency: {:?}", + overall_burst_metrics.p99_latency + ); + eprintln!( + " Transactions >15ms: {} ({:.2}%)", + overall_burst_metrics.operations_over_15ms, + overall_burst_metrics.operations_over_15ms_percent + ); + assert_eq!( + spanner.active_channel_count(), + 64, + "static channel pool count must remain 64" + ); + eprintln!("==============================================================================="); +}