diff --git a/src/spanner/src/database_client.rs b/src/spanner/src/database_client.rs index 3aacf24d31..c26f6bf53c 100644 --- a/src/spanner/src/database_client.rs +++ b/src/spanner/src/database_client.rs @@ -31,6 +31,7 @@ use crate::partitioned_dml_transaction::PartitionedDmlTransactionBuilder; use crate::read_only_transaction::{ MultiUseReadOnlyTransactionBuilder, SingleUseReadOnlyTransactionBuilder, }; +use crate::retry_delay::{extract_retry_delay_from_error, extract_status_code_from_error}; use crate::routing::cache_subscriber::CacheSubscriber; use crate::routing::cache_updater::CacheUpdater; use crate::routing::connection_cache::ConnectionCache; @@ -51,9 +52,10 @@ use crate::transaction_runner::TransactionRunnerBuilder; use crate::write_only_transaction::WriteOnlyTransactionBuilder; use crate::{RequestOptions, Result}; use bytes::Bytes; +use google_cloud_gax::error::rpc::Code; use std::env; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; /// A client for interacting with a specific Spanner database. /// @@ -83,7 +85,6 @@ pub struct DatabaseClient { spanner: Spanner, pub(crate) session_maintainer: Arc, pub(crate) leader_aware_routing_enabled: bool, - #[allow(dead_code)] // TODO: Used by request routing interceptors in subsequent PRs pub(crate) location_routing: Option>, pub(crate) o11y: Arc, } @@ -108,10 +109,17 @@ macro_rules! define_db_rpc { Some(connection) => connection.channel(), None => self.spanner.get_channel(channel_hint), }; + let _request_guard = connection + .as_ref() + .map(ServerConnection::acquire_request_guard); + let group_uid = request.routing_group_uid(); + let start = Instant::now(); let result = self .spanner .$method(request, options, channel, &self.o11y) .await; + let latency = start.elapsed(); + self.record_routing_feedback(connection.as_ref(), group_uid, latency, &result); $post_hook(self, routing_context, connection.as_ref(), &result); let response = result?; response.observe(self); @@ -293,7 +301,6 @@ impl DatabaseClient { /// pointing directly to the target node. /// - If route resolution falls back to the default gateway connection, returns `None` so the /// request is dispatched over the client's channel pool using the transaction's assigned channel affinity. - #[allow(dead_code)] // TODO(#6236): Used by request routing in subsequent PRs pub(crate) fn resolve_routing_connection( &self, context: &RoutingContext, @@ -545,8 +552,61 @@ impl DatabaseClient { .map(|routing| routing.location_router.latency_registry()) } + /// Records feedback from an RPC execution into the [`LocationRouter`] and [`LatencyRegistry`]. + /// + /// - For successful direct node executions, repairs cooldown failure tiers via `record_success` + /// and records round-trip latency via `record_latency`. + /// - For failures on direct nodes (`Code::ResourceExhausted` or `Code::Unavailable`), places the + /// endpoint on cooldown with any server-recommended retry delay, and inflates the latency score + /// with an error penalty. + /// - Gateway fallback connections (`is_none()` or `is_default()`) are never placed on cooldown. + fn record_routing_feedback( + &self, + connection: Option<&ServerConnection>, + group_uid: u64, + latency: Duration, + result: &Result, + ) { + let Some(routing) = &self.location_routing else { + return; + }; + let Some(connection) = connection else { + return; + }; + if connection.is_default() { + return; + } + let address = connection.address(); + match result { + Ok(_) => { + routing.location_router.record_success(address); + if group_uid > 0 { + routing + .location_router + .record_latency(group_uid, address, latency); + } + } + Err(error) => { + let Some(code) = extract_status_code_from_error(error) else { + return; + }; + if matches!(code, Code::ResourceExhausted | Code::Unavailable) { + let server_retry_delay = extract_retry_delay_from_error(error); + routing.location_router.record_cooldown_error_with_delay( + address, + code, + server_retry_delay, + ); + if group_uid > 0 { + routing.location_router.record_error(group_uid, address); + } + } + } + } + } + /// Records an observed round-trip latency sample for an endpoint address within a paxos group. - #[allow(dead_code)] // TODO: Used for latency recording in subsequent PRs + #[allow(dead_code)] // TODO: Used for streaming RPC latency recording in subsequent PRs pub(crate) fn record_latency(&self, group_uid: u64, server_address: &str, latency: Duration) { let Some(routing) = &self.location_routing else { return; @@ -557,7 +617,7 @@ impl DatabaseClient { } /// Records an RPC error penalty for an endpoint address within a paxos group. - #[allow(dead_code)] // TODO: Used for error penalty recording in subsequent PRs + #[allow(dead_code)] // TODO: Used for streaming RPC error recording in subsequent PRs pub(crate) fn record_routing_error(&self, group_uid: u64, server_address: &str) { let Some(routing) = &self.location_routing else { return; @@ -1339,6 +1399,60 @@ impl ObserveResponse for PartitionResponse { fn observe(&self, _client: &DatabaseClient) {} } +/// Extracts the group UID from an RPC request to associate routing feedback with the covering group. +/// +/// Protobuf requests that define a `routing_hint` field ([`ExecuteSqlRequest`], [`BeginTransactionRequest`], +/// and [`CommitRequest`]) inspect their hint and return `group_uid`. Requests without a `routing_hint` +/// field in their protobuf definitions ([`ExecuteBatchDmlRequest`], [`RollbackRequest`], [`PartitionQueryRequest`], +/// and [`PartitionReadRequest`]) return 0. +trait RequestRoutingGroupUid { + /// Returns the routing group UID if this request has an attached routing hint with a non-zero group UID, + /// or 0 if unkeyed or unhinted. + fn routing_group_uid(&self) -> u64; +} + +impl RequestRoutingGroupUid for ExecuteSqlRequest { + fn routing_group_uid(&self) -> u64 { + self.routing_hint.as_ref().map_or(0, |hint| hint.group_uid) + } +} + +impl RequestRoutingGroupUid for BeginTransactionRequest { + fn routing_group_uid(&self) -> u64 { + self.routing_hint.as_ref().map_or(0, |hint| hint.group_uid) + } +} + +impl RequestRoutingGroupUid for CommitRequest { + fn routing_group_uid(&self) -> u64 { + self.routing_hint.as_ref().map_or(0, |hint| hint.group_uid) + } +} + +impl RequestRoutingGroupUid for ExecuteBatchDmlRequest { + fn routing_group_uid(&self) -> u64 { + 0 + } +} + +impl RequestRoutingGroupUid for RollbackRequest { + fn routing_group_uid(&self) -> u64 { + 0 + } +} + +impl RequestRoutingGroupUid for PartitionQueryRequest { + fn routing_group_uid(&self) -> u64 { + 0 + } +} + +impl RequestRoutingGroupUid for PartitionReadRequest { + fn routing_group_uid(&self) -> u64 { + 0 + } +} + #[cfg(test)] mod tests { use super::*; @@ -1370,6 +1484,68 @@ mod tests { use std::sync::Mutex; use tokio::sync::mpsc; + #[test] + fn request_routing_group_uid_with_and_without_hint() { + let hint = RoutingHint::new().set_group_uid(9001u64); + + let sql_with_hint = ExecuteSqlRequest::default().set_routing_hint(hint.clone()); + assert_eq!( + sql_with_hint.routing_group_uid(), + 9001, + "ExecuteSqlRequest with hint must return hint group_uid" + ); + assert_eq!( + ExecuteSqlRequest::default().routing_group_uid(), + 0, + "ExecuteSqlRequest without hint must return 0" + ); + + let begin_with_hint = BeginTransactionRequest::default().set_routing_hint(hint.clone()); + assert_eq!( + begin_with_hint.routing_group_uid(), + 9001, + "BeginTransactionRequest with hint must return hint group_uid" + ); + assert_eq!( + BeginTransactionRequest::default().routing_group_uid(), + 0, + "BeginTransactionRequest without hint must return 0" + ); + + let commit_with_hint = CommitRequest::default().set_routing_hint(hint); + assert_eq!( + commit_with_hint.routing_group_uid(), + 9001, + "CommitRequest with hint must return hint group_uid" + ); + assert_eq!( + CommitRequest::default().routing_group_uid(), + 0, + "CommitRequest without hint must return 0" + ); + + assert_eq!( + ExecuteBatchDmlRequest::default().routing_group_uid(), + 0, + "ExecuteBatchDmlRequest must return 0" + ); + assert_eq!( + RollbackRequest::default().routing_group_uid(), + 0, + "RollbackRequest must return 0" + ); + assert_eq!( + PartitionQueryRequest::default().routing_group_uid(), + 0, + "PartitionQueryRequest must return 0" + ); + assert_eq!( + PartitionReadRequest::default().routing_group_uid(), + 0, + "PartitionReadRequest must return 0" + ); + } + fn create_test_mock() -> MockSpanner { let mut mock = MockSpanner::new(); mock.expect_create_session().returning(|req| { @@ -1531,7 +1707,8 @@ mod tests { Ok(_) => panic!("Client creation should have failed"), Err(e) => assert_eq!( e.status().map(|s| s.code), - Some(google_cloud_gax::error::rpc::Code::PermissionDenied) + Some(Code::PermissionDenied), + "error status code should match expected PermissionDenied" ), } } diff --git a/src/spanner/src/retry_delay.rs b/src/spanner/src/retry_delay.rs index de5258552e..a01f6ba80e 100644 --- a/src/spanner/src/retry_delay.rs +++ b/src/spanner/src/retry_delay.rs @@ -22,7 +22,7 @@ use crate::google::rpc::Status as ProtoStatus; use base64::Engine as _; use base64::prelude::{BASE64_STANDARD, BASE64_STANDARD_NO_PAD}; use gaxi::grpc::tonic::Status as TonicStatus; -use google_cloud_gax::error::rpc::{Status, StatusDetails}; +use google_cloud_gax::error::rpc::{Code, Status, StatusDetails}; use http::HeaderMap; use prost::Message; use prost_types::Duration as ProtoDuration; @@ -43,14 +43,56 @@ pub(crate) struct ProtoRetryInfo { pub retry_delay: Option, } +/// Extracts the gRPC status code from an [`Error`], inspecting GAX status, nested [`Error`] instances, and nested [`TonicStatus`]. +pub(crate) fn extract_status_code_from_error(error: &Error) -> Option { + if let Some(status) = error.status() { + return Some(status.code); + } + + let mut current_source = error.source(); + while let Some(source) = current_source { + if let Some(inner_error) = source.downcast_ref::() + && let Some(status) = inner_error.status() + { + return Some(status.code); + } + if let Some(status) = source.downcast_ref::() { + return Some(Code::from(status.code() as i32)); + } + current_source = source.source(); + } + + None +} + /// Extracts the server-recommended retry delay from an [`Error`], if present. pub(crate) fn extract_retry_delay_from_error(error: &Error) -> Option { if let Some(delay) = error.status().and_then(extract_retry_delay_from_status) { return Some(delay); } + if let Some(delay) = error + .http_headers() + .and_then(extract_retry_delay_from_headers) + { + return Some(delay); + } let mut current_source = error.source(); while let Some(source) = current_source { + if let Some(inner_error) = source.downcast_ref::() { + if let Some(delay) = inner_error + .status() + .and_then(extract_retry_delay_from_status) + { + return Some(delay); + } + if let Some(delay) = inner_error + .http_headers() + .and_then(extract_retry_delay_from_headers) + { + return Some(delay); + } + } if let Some(delay) = source .downcast_ref::() .and_then(extract_retry_delay_from_tonic_status) @@ -60,9 +102,7 @@ pub(crate) fn extract_retry_delay_from_error(error: &Error) -> Option current_source = source.source(); } - error - .http_headers() - .and_then(extract_retry_delay_from_headers) + None } /// Extracts the server-recommended retry delay from HTTP headers, if the `google.rpc.retryinfo-bin` header is present. @@ -144,6 +184,7 @@ mod tests { use prost_types::Any as ProtoAny; use static_assertions::assert_impl_all; use std::fmt::Debug; + use wkt::Duration as WktDuration; assert_impl_all!(ProtoRetryInfo: Send, Sync, Debug, Clone); @@ -550,4 +591,99 @@ mod tests { "ProtoRetryInfo without retry_delay must yield None" ); } + + #[test] + fn extract_status_code_from_gax_status() { + let status = Status::default().set_code(Code::ResourceExhausted); + let error = Error::service(status); + assert_eq!( + extract_status_code_from_error(&error), + Some(Code::ResourceExhausted), + "must extract ResourceExhausted status code from GAX Status" + ); + } + + #[test] + fn extract_status_code_from_tonic_status_source() { + let tonic_status = TonicStatus::unavailable("endpoint temporarily unavailable"); + let error = Error::connect(tonic_status); + assert_eq!( + extract_status_code_from_error(&error), + Some(Code::Unavailable), + "must extract Unavailable status code from TonicStatus source" + ); + } + + #[test] + fn extract_status_code_from_non_status_error() { + let error = Error::exhausted("retry policy exhausted"); + assert_eq!( + extract_status_code_from_error(&error), + None, + "non-status error must yield None for status code" + ); + } + + #[test] + fn extract_status_code_from_exhausted_nested_service_error() { + let status = Status::default().set_code(Code::ResourceExhausted); + let inner_error = Error::service(status); + let exhausted = Error::exhausted(inner_error); + assert_eq!( + extract_status_code_from_error(&exhausted), + Some(Code::ResourceExhausted), + "must extract status code from nested Error inside Error::exhausted" + ); + } + + #[test] + fn extract_retry_delay_from_exhausted_nested_service_error() { + let retry_info = RetryInfo::default().set_retry_delay(WktDuration::clamp(4, 500_000_000)); + let status = Status::default() + .set_code(Code::ResourceExhausted) + .set_details(vec![StatusDetails::RetryInfo(retry_info)]); + let inner_error = Error::service(status); + let exhausted = Error::exhausted(inner_error); + assert_eq!( + extract_retry_delay_from_error(&exhausted), + Some(Duration::new(4, 500_000_000)), + "must extract retry delay from nested Error inside Error::exhausted" + ); + } + + #[test] + fn extract_retry_delay_from_exhausted_nested_tonic_error() { + let retry_info_bytes = encode_test_retry_info(3, 0); + let tonic_status = TonicStatus::with_details_and_metadata( + TonicCode::ResourceExhausted, + "overloaded", + retry_info_bytes.into(), + MetadataMap::new(), + ); + let inner_error = Error::connect(tonic_status); + let exhausted = Error::exhausted(inner_error); + assert_eq!( + extract_retry_delay_from_error(&exhausted), + Some(Duration::from_secs(3)), + "must extract retry delay from nested TonicStatus inside Error::exhausted" + ); + } + + #[test] + fn extract_retry_delay_from_exhausted_nested_http_headers() { + let retry_info_bytes = encode_test_retry_info(6, 250_000_000); + let mut headers = HeaderMap::new(); + let base64_encoded = BASE64_STANDARD.encode(&retry_info_bytes); + headers.insert( + HeaderName::from_static(RETRY_INFO_BINARY_HEADER), + HeaderValue::from_str(&base64_encoded).expect("valid header value"), + ); + let inner_http_error = Error::http(429, headers, Bytes::new()); + let exhausted = Error::exhausted(inner_http_error); + assert_eq!( + extract_retry_delay_from_error(&exhausted), + Some(Duration::new(6, 250_000_000)), + "must extract retry delay from nested Error with http_headers inside Error::exhausted" + ); + } } diff --git a/src/spanner/src/routing/mock_tests.rs b/src/spanner/src/routing/mock_tests.rs index 87c0f7b0d2..754a904d87 100644 --- a/src/spanner/src/routing/mock_tests.rs +++ b/src/spanner/src/routing/mock_tests.rs @@ -49,6 +49,7 @@ use crate::mutation::Mutation; use crate::omni::{InstanceType, TlsConfig}; use crate::read::ReadRequest; use crate::read_write_transaction::ReadWriteTransaction; +use crate::retry_delay::ProtoRetryInfo; use crate::routing::directed_read::select_eligible_tablets_for_directed_read; use crate::routing::key_range_cache::RangeMode; use crate::routing::location_router::RoutingContext; @@ -56,10 +57,13 @@ use crate::statement::Statement; use bytes::Bytes; use gaxi::grpc::tonic::transport::server::TcpIncoming; use gaxi::grpc::tonic::transport::{Certificate, Identity, Server, ServerTlsConfig}; -use gaxi::grpc::tonic::{Response, Status as TonicStatus}; +use gaxi::grpc::tonic::{Code as TonicCode, MetadataMap, Response, Status as TonicStatus}; use google_cloud_auth::credentials::anonymous::Builder as Anonymous; +use google_cloud_gax::retry_policy::NeverRetry; use google_cloud_test_macros::tokio_test_no_panics; -use prost_types::{Timestamp, Value}; +use mockall::Sequence; +use prost::Message as _; +use prost_types::{Duration as ProtoDuration, Timestamp, Value}; use spanner_grpc_mock::MockSpanner; use spanner_grpc_mock::google::rpc::Status; use spanner_grpc_mock::google::spanner::v1 as mock_v1; @@ -67,7 +71,7 @@ use spanner_grpc_mock::google::spanner::v1::spanner_server::SpannerServer; use spanner_grpc_mock::start; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; -use std::time::Duration; +use std::time::{Duration, Instant}; use tokio::sync::mpsc; #[tokio_test_no_panics] @@ -5023,3 +5027,834 @@ async fn end_to_end_routed_connection_with_omni_mtls_and_explicit_domain_overrid Ok(()) } + +fn encode_test_retry_info(seconds: i64, nanos: i32) -> Vec { + let mut bytes = Vec::new(); + ProtoRetryInfo { + retry_delay: Some(ProtoDuration { seconds, nanos }), + } + .encode(&mut bytes) + .expect("encoding ProtoRetryInfo must succeed"); + bytes +} + +#[tokio_test_no_panics] +async fn unary_rpc_feedback_success_records_latency_and_repairs_cooldown() -> anyhow::Result<()> { + let mut mock_gateway = create_base_mock(); + let mut mock_tablet = create_base_mock(); + + let tablet_called = Arc::new(AtomicBool::new(false)); + let tablet_called_clone = Arc::clone(&tablet_called); + + mock_tablet + .expect_execute_sql() + .times(3) + .returning(move |_| { + tablet_called_clone.store(true, Ordering::SeqCst); + Ok(Response::new(mock_v1::ResultSet::default())) + }); + + let (tablet_address, _tablet_server) = start("127.0.0.1:0", mock_tablet).await?; + + let tablet_address_for_update = tablet_address.clone(); + mock_gateway + .expect_execute_sql() + .times(1) + .returning(move |request| { + let operation_uid = request + .get_ref() + .routing_hint + .as_ref() + .map(|hint| hint.operation_uid) + .unwrap_or(1); + + let cache_update = sample_query_mock_cache_update( + 5555, + operation_uid, + 9001, + &tablet_address_for_update, + ); + Ok(Response::new(mock_v1::ResultSet { + cache_update: Some(cache_update), + ..Default::default() + })) + }); + + let (database_client, _spanner, _gateway_server) = + setup_mock_database_client(mock_gateway).await?; + + let router = database_client + .location_router() + .expect("location router must be present"); + + // Cold-start query populates recipe and range cache via gateway + let statement = Statement::builder("SELECT * FROM Users WHERE account_id = @account_id") + .add_param("account_id", 42i64) + .build(); + let _ = database_client + .execute_sql( + statement.clone().into_request(), + RequestOptions::default(), + 0, + ) + .await?; + + // Pre-warm tablet connection in cache + let client_config = database_client + .cache_updater() + .expect("cache updater must be present") + .client_config(); + let _ = router + .connection_cache() + .get(&tablet_address, client_config) + .await?; + + // Record an expired failure in cooldown tracker (30s in the past so cooldown is over, but failure tier is 1) + let past = Instant::now() - Duration::from_secs(30); + router + .cooldown_tracker() + .record_failure_at(&tablet_address, past); + assert!( + !router.cooldown_tracker().is_empty(), + "cooldown tracker should contain entry for tablet address before repair" + ); + + assert!( + !router + .latency_registry() + .has_score(Some(router.database_scope()), 9001, &tablet_address), + "latency registry should have no score for tablet address before direct RPC execution" + ); + + // Execute 3 successful direct unary RPCs to the tablet + for _ in 0..3 { + let _ = database_client + .execute_sql( + statement.clone().into_request(), + RequestOptions::default(), + 0, + ) + .await?; + } + + assert!( + tablet_called.load(Ordering::SeqCst), + "mock tablet must have received direct unary execute_sql calls" + ); + + // Verify latency was recorded + assert!( + router + .latency_registry() + .has_score(Some(router.database_scope()), 9001, &tablet_address), + "latency registry must have recorded latency score for tablet address after successful direct execution" + ); + let cost = router.latency_registry().get_selection_cost( + Some(router.database_scope()), + 9001, + 0, + &tablet_address, + ); + assert!( + cost > 0.0, + "selection cost for tablet address must be greater than zero after latency recording" + ); + + // Verify that 3 consecutive successes repaired the failure tier and pruned the entry + assert!( + router.cooldown_tracker().is_empty(), + "cooldown tracker must be empty after 3 successful RPCs repaired failure tier of expired entry" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn unary_rpc_feedback_resource_exhausted_places_endpoint_on_cooldown_and_penalizes_latency() +-> anyhow::Result<()> { + let mut mock_gateway = create_base_mock(); + let mut mock_tablet = create_base_mock(); + + let tablet_called = Arc::new(AtomicBool::new(false)); + let tablet_called_clone = Arc::clone(&tablet_called); + + // Encode a RetryInfo with a 2-second delay + let retry_info_bytes = encode_test_retry_info(2, 0); + + mock_tablet + .expect_execute_sql() + .times(1) + .returning(move |_| { + tablet_called_clone.store(true, Ordering::SeqCst); + Err(TonicStatus::with_details_and_metadata( + TonicCode::ResourceExhausted, + "tablet overloaded with resource exhaustion", + retry_info_bytes.clone().into(), + MetadataMap::new(), + )) + }); + + let (tablet_address, _tablet_server) = start("127.0.0.1:0", mock_tablet).await?; + + let mut gateway_sequence = Sequence::new(); + let tablet_address_for_update = tablet_address.clone(); + let fallback_called = Arc::new(AtomicBool::new(false)); + let fallback_called_clone = Arc::clone(&fallback_called); + + mock_gateway + .expect_execute_sql() + .times(1) + .in_sequence(&mut gateway_sequence) + .returning(move |request| { + // Cold-start response providing recipe and range + let operation_uid = request + .get_ref() + .routing_hint + .as_ref() + .map(|hint| hint.operation_uid) + .unwrap_or(1); + + let cache_update = sample_query_mock_cache_update( + 5555, + operation_uid, + 9001, + &tablet_address_for_update, + ); + Ok(Response::new(mock_v1::ResultSet { + cache_update: Some(cache_update), + ..Default::default() + })) + }); + + mock_gateway + .expect_execute_sql() + .times(1) + .in_sequence(&mut gateway_sequence) + .returning(move |_| { + fallback_called_clone.store(true, Ordering::SeqCst); + // Fallback execution when tablet is on cooldown + Ok(Response::new(mock_v1::ResultSet::default())) + }); + + let (database_client, _spanner, _gateway_server) = + setup_mock_database_client(mock_gateway).await?; + + let router = database_client + .location_router() + .expect("location router must be present"); + + // Cold-start query populates recipe and range cache via gateway + let statement = Statement::builder("SELECT * FROM Users WHERE account_id = @account_id") + .add_param("account_id", 42i64) + .build(); + let _ = database_client + .execute_sql( + statement.clone().into_request(), + RequestOptions::default(), + 0, + ) + .await?; + + assert!( + !fallback_called.load(Ordering::SeqCst), + "gateway fallback must not be invoked during initial cold-start query" + ); + + // Pre-warm tablet connection in cache + let client_config = database_client + .cache_updater() + .expect("cache updater must be present") + .client_config(); + let _ = router + .connection_cache() + .get(&tablet_address, client_config) + .await?; + + assert!( + !router.cooldown_tracker().is_cooling_down(&tablet_address), + "tablet address must not be cooling down initially" + ); + assert!( + !router + .latency_registry() + .has_score(Some(router.database_scope()), 9001, &tablet_address), + "latency registry should have no score for tablet address before failure" + ); + + // Second execution targets tablet directly and fails with ResourceExhausted + let result = database_client + .execute_sql( + statement.clone().into_request(), + RequestOptions::default(), + 0, + ) + .await; + assert!( + result.is_err(), + "unary execute_sql to overloaded tablet must return an error" + ); + assert!( + tablet_called.load(Ordering::SeqCst), + "mock tablet must have received direct unary execute_sql call" + ); + + // Verify tablet was placed on cooldown and latency penalty was recorded + assert!( + router.cooldown_tracker().is_cooling_down(&tablet_address), + "tablet address must be on active cooldown after ResourceExhausted error" + ); + assert!( + router + .latency_registry() + .has_score(Some(router.database_scope()), 9001, &tablet_address), + "latency registry must have recorded error penalty for tablet address after failure" + ); + + // 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) + .await; + assert!( + fallback_result.is_ok(), + "subsequent execute_sql must succeed via fallback when tablet is cooling down" + ); + assert!( + fallback_called.load(Ordering::SeqCst), + "gateway must have received the fallback query execution" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn unary_rpc_feedback_unavailable_places_endpoint_on_cooldown() -> anyhow::Result<()> { + let mut mock_gateway = create_base_mock(); + let mut mock_tablet = create_base_mock(); + + let tablet_called = Arc::new(AtomicBool::new(false)); + let tablet_called_clone = Arc::clone(&tablet_called); + + mock_tablet + .expect_execute_sql() + .times(1) + .returning(move |_| { + tablet_called_clone.store(true, Ordering::SeqCst); + Err(TonicStatus::unavailable("tablet node restarting")) + }); + + let (tablet_address, _tablet_server) = start("127.0.0.1:0", mock_tablet).await?; + + let tablet_address_for_update = tablet_address.clone(); + mock_gateway + .expect_execute_sql() + .times(1) + .returning(move |request| { + let operation_uid = request + .get_ref() + .routing_hint + .as_ref() + .map(|hint| hint.operation_uid) + .unwrap_or(1); + + let cache_update = sample_query_mock_cache_update( + 5555, + operation_uid, + 9001, + &tablet_address_for_update, + ); + Ok(Response::new(mock_v1::ResultSet { + cache_update: Some(cache_update), + ..Default::default() + })) + }); + + let (database_client, _spanner, _gateway_server) = + setup_mock_database_client(mock_gateway).await?; + + let router = database_client + .location_router() + .expect("location router must be present"); + + // Cold-start query populates recipe and range cache via gateway + let statement = Statement::builder("SELECT * FROM Users WHERE account_id = @account_id") + .add_param("account_id", 42i64) + .build(); + let _ = database_client + .execute_sql( + statement.clone().into_request(), + RequestOptions::default(), + 0, + ) + .await?; + + // Pre-warm tablet connection in cache + let client_config = database_client + .cache_updater() + .expect("cache updater must be present") + .client_config(); + let _ = router + .connection_cache() + .get(&tablet_address, client_config) + .await?; + + assert!( + !router.cooldown_tracker().is_cooling_down(&tablet_address), + "tablet address must not be cooling down initially" + ); + + // Direct execution to tablet fails with Unavailable (using NeverRetry to avoid retry backoff) + let mut options = RequestOptions::default(); + options.set_retry_policy(NeverRetry); + let result = database_client + .execute_sql(statement.into_request(), options, 0) + .await; + assert!( + result.is_err(), + "unary execute_sql to unavailable tablet must return an error" + ); + assert!( + tablet_called.load(Ordering::SeqCst), + "mock tablet must have received direct unary execute_sql call" + ); + + // Verify tablet was placed on cooldown and latency penalty was recorded + assert!( + router.cooldown_tracker().is_cooling_down(&tablet_address), + "tablet address must be on active cooldown after Unavailable error" + ); + assert!( + router + .latency_registry() + .has_score(Some(router.database_scope()), 9001, &tablet_address), + "latency registry must have recorded error penalty for tablet address after Unavailable failure" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn unary_rpc_feedback_gateway_fallback_never_placed_on_cooldown() -> anyhow::Result<()> { + let mut mock_gateway = create_base_mock(); + + let gateway_called = Arc::new(AtomicBool::new(false)); + let gateway_called_clone = Arc::clone(&gateway_called); + + mock_gateway + .expect_execute_sql() + .times(1) + .returning(move |_| { + gateway_called_clone.store(true, Ordering::SeqCst); + Err(TonicStatus::resource_exhausted("gateway overloaded")) + }); + + let (database_client, _spanner, _gateway_server) = + setup_mock_database_client(mock_gateway).await?; + + let router = database_client + .location_router() + .expect("location router must be present"); + + // Cold-start query executed against gateway fails with ResourceExhausted + let statement = Statement::builder("SELECT * FROM Users WHERE account_id = @account_id") + .add_param("account_id", 42i64) + .build(); + let result = database_client + .execute_sql(statement.into_request(), RequestOptions::default(), 0) + .await; + assert!( + result.is_err(), + "query to overloaded gateway must return an error" + ); + assert!( + gateway_called.load(Ordering::SeqCst), + "mock gateway must have received the query" + ); + + // Verify gateway fallback was never placed on cooldown + assert!( + router.cooldown_tracker().is_empty(), + "cooldown tracker must remain completely empty after gateway fallback error" + ); + assert!( + router.latency_registry().is_empty(), + "latency registry must remain empty after gateway fallback error" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn unary_rpc_feedback_non_retryable_error_does_not_place_endpoint_on_cooldown() +-> anyhow::Result<()> { + let mut mock_gateway = create_base_mock(); + let mut mock_tablet = create_base_mock(); + + let tablet_called = Arc::new(AtomicBool::new(false)); + let tablet_called_clone = Arc::clone(&tablet_called); + + mock_tablet + .expect_execute_sql() + .times(1) + .returning(move |_| { + tablet_called_clone.store(true, Ordering::SeqCst); + Err(TonicStatus::invalid_argument( + "invalid query syntax in tablet", + )) + }); + + let (tablet_address, _tablet_server) = start("127.0.0.1:0", mock_tablet).await?; + + let tablet_address_for_update = tablet_address.clone(); + mock_gateway + .expect_execute_sql() + .times(1) + .returning(move |request| { + let operation_uid = request + .get_ref() + .routing_hint + .as_ref() + .map(|hint| hint.operation_uid) + .unwrap_or(1); + + let cache_update = sample_query_mock_cache_update( + 5555, + operation_uid, + 9001, + &tablet_address_for_update, + ); + Ok(Response::new(mock_v1::ResultSet { + cache_update: Some(cache_update), + ..Default::default() + })) + }); + + let (database_client, _spanner, _gateway_server) = + setup_mock_database_client(mock_gateway).await?; + + let router = database_client + .location_router() + .expect("location router must be present"); + + // Cold-start query populates recipe and range cache via gateway + let statement = Statement::builder("SELECT * FROM Users WHERE account_id = @account_id") + .add_param("account_id", 42i64) + .build(); + let _ = database_client + .execute_sql( + statement.clone().into_request(), + RequestOptions::default(), + 0, + ) + .await?; + + // Pre-warm tablet connection in cache + let client_config = database_client + .cache_updater() + .expect("cache updater must be present") + .client_config(); + let _ = router + .connection_cache() + .get(&tablet_address, client_config) + .await?; + + assert!( + !router.cooldown_tracker().is_cooling_down(&tablet_address), + "tablet address must not be cooling down initially" + ); + + // Direct execution to tablet fails with non-retryable InvalidArgument + let result = database_client + .execute_sql(statement.into_request(), RequestOptions::default(), 0) + .await; + assert!( + result.is_err(), + "unary execute_sql with invalid argument must return an error" + ); + assert!( + tablet_called.load(Ordering::SeqCst), + "mock tablet must have received direct unary execute_sql call" + ); + + // Verify non-retryable error did NOT trigger cooldown or latency error penalty + assert!( + !router.cooldown_tracker().is_cooling_down(&tablet_address), + "tablet address must NOT be cooling down after non-retryable InvalidArgument error" + ); + assert!( + router.cooldown_tracker().is_empty(), + "cooldown tracker must remain empty after non-retryable error" + ); + assert!( + !router + .latency_registry() + .has_score(Some(router.database_scope()), 9001, &tablet_address), + "latency registry must NOT have recorded error penalty for non-retryable error" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn unary_rpc_feedback_begin_transaction_records_latency_and_repairs_cooldown() +-> anyhow::Result<()> { + let mock_gateway = create_base_mock(); + let (gateway_address, _gateway_server) = start("127.0.0.1:0", mock_gateway).await?; + + let mut mock_tablet_leader = create_base_mock(); + let leader_begin_called = Arc::new(AtomicBool::new(false)); + let leader_begin_called_clone = Arc::clone(&leader_begin_called); + + mock_tablet_leader + .expect_begin_transaction() + .times(1) + .returning(move |_| { + leader_begin_called_clone.store(true, Ordering::SeqCst); + Ok(Response::new(mock_v1::Transaction { + id: b"tx-feedback-leader-123".to_vec(), + ..Default::default() + })) + }); + let (tablet_leader_address, _tablet_leader_server) = + start("127.0.0.1:0", mock_tablet_leader).await?; + + let mock_tablet_follower = create_base_mock(); + let (tablet_follower_address, _tablet_follower_server) = + start("127.0.0.1:0", mock_tablet_follower).await?; + + let spanner = Spanner::builder() + .with_endpoint(gateway_address.clone()) + .with_instance_type(InstanceType::Omni) + .with_credentials(Anonymous::new().build()) + .build() + .await?; + + let database_client = spanner + .database_client("projects/test-project/instances/test-instance/databases/test-db") + .with_location_aware_routing(true) + .build() + .await?; + + let mut update = sample_model_cache_update( + 10102, + 8002, + &tablet_leader_address, + &tablet_follower_address, + ); + update.range = vec![ModelRange { + start_key: Bytes::from_static(b""), + limit_key: Bytes::from_static(b""), + group_uid: 8002, + split_id: 8002, + generation: Bytes::from_static(b"gen_1"), + _unknown_fields: Default::default(), + }]; + database_client.observe_cache_update(Some(update)); + + let router = database_client + .location_router() + .expect("location router must be present"); + let client_config = database_client + .cache_updater() + .expect("cache updater present") + .client_config(); + let _ = router + .connection_cache() + .get(&tablet_leader_address, client_config) + .await?; + + assert!( + !router.latency_registry().has_score( + Some(router.database_scope()), + 8002, + &tablet_leader_address + ), + "latency registry should have no score for tablet leader prior to execution" + ); + + let mutation = Mutation::new_insert_builder("Singers") + .set("SingerId") + .to(101i64) + .build(); + + let begin_request = BeginTransactionRequest::default() + .set_session( + "projects/test-project/instances/test-instance/databases/test-db/sessions/session-1", + ) + .set_options(TransactionOptions::default().set_read_write(ReadWrite::default())) + .set_mutation_key(mutation.build_proto()); + + let response = database_client + .begin_transaction(begin_request, RequestOptions::default(), 0) + .await?; + + assert!( + leader_begin_called.load(Ordering::SeqCst), + "explicit begin transaction with mutation key must route directly to leader tablet" + ); + assert_eq!( + response.id.as_ref(), + b"tx-feedback-leader-123", + "returned transaction ID must match mock response" + ); + assert!( + router.latency_registry().has_score( + Some(router.database_scope()), + 8002, + &tablet_leader_address + ), + "successful begin_transaction must record latency score for routing group 8002" + ); + + let cost = router.latency_registry().selection_cost( + Some(router.database_scope()), + 8002, + 0, + &tablet_leader_address, + ); + assert!( + cost > 0.0, + "selection cost must be positive after successful RPC" + ); + + Ok(()) +} + +#[tokio_test_no_panics] +async fn unary_rpc_feedback_direct_affinity_zero_group_uid_handles_cooldown_and_fallback() +-> anyhow::Result<()> { + let mut mock_gateway = create_base_mock(); + let mut mock_tablet = create_base_mock(); + + let mut tablet_sequence = Sequence::new(); + + mock_tablet + .expect_execute_sql() + .times(1) + .in_sequence(&mut tablet_sequence) + .returning(|_| Ok(Response::new(mock_v1::ResultSet::default()))); + + mock_tablet + .expect_execute_sql() + .times(1) + .in_sequence(&mut tablet_sequence) + .returning(|_| { + Err(TonicStatus::resource_exhausted( + "tablet node overloaded on affinity request", + )) + }); + + let (tablet_address, _tablet_server) = start("127.0.0.1:0", mock_tablet).await?; + + let fallback_called = Arc::new(AtomicBool::new(false)); + let fallback_called_clone = Arc::clone(&fallback_called); + + mock_gateway + .expect_execute_sql() + .times(1) + .returning(move |_| { + fallback_called_clone.store(true, Ordering::SeqCst); + Ok(Response::new(mock_v1::ResultSet::default())) + }); + + let (database_client, _spanner, _gateway_server) = + setup_mock_database_client(mock_gateway).await?; + + let router = database_client + .location_router() + .expect("location router must be present"); + + // Pre-warm tablet connection in cache + let client_config = database_client + .cache_updater() + .expect("cache updater must be present") + .client_config(); + let _ = router + .connection_cache() + .get(&tablet_address, client_config) + .await?; + + let transaction_id = b"tx-affinity-feedback-zero-group"; + router.record_transaction_affinity(transaction_id, &tablet_address); + + // Call 1: Direct affinity execution succeeds with group_uid == 0. + // Verifies that success repairs cooldown and does NOT record any score in latency registry under group 0. + let execute_request_1 = ExecuteSqlRequest::default() + .set_session( + "projects/test-project/instances/test-instance/databases/test-db/sessions/session-1", + ) + .set_sql("SELECT 1") + .set_transaction( + TransactionSelector::default().set_id(Bytes::copy_from_slice(transaction_id)), + ); + + let result_1 = database_client + .execute_sql(execute_request_1, RequestOptions::default(), 0) + .await; + assert!( + result_1.is_ok(), + "direct affinity execute_sql with group_uid 0 must succeed" + ); + assert!( + router.latency_registry().is_empty(), + "success with group_uid 0 must not record latency score under group 0" + ); + assert!( + !router.cooldown_tracker().is_cooling_down(&tablet_address), + "tablet address must not be cooling down after successful RPC" + ); + + // Call 2: Direct affinity execution fails with ResourceExhausted (group_uid == 0). + // Verifies that error places affinity endpoint on cooldown, but does NOT record error penalty in latency registry. + let execute_request_2 = ExecuteSqlRequest::default() + .set_session( + "projects/test-project/instances/test-instance/databases/test-db/sessions/session-1", + ) + .set_sql("SELECT 1") + .set_transaction( + TransactionSelector::default().set_id(Bytes::copy_from_slice(transaction_id)), + ); + + let result_2 = database_client + .execute_sql(execute_request_2, RequestOptions::default(), 0) + .await; + assert!( + result_2.is_err(), + "direct affinity execute_sql to overloaded tablet must return an error" + ); + assert!( + router.cooldown_tracker().is_cooling_down(&tablet_address), + "affinity tablet must be placed on cooldown after ResourceExhausted error" + ); + assert!( + router.latency_registry().is_empty(), + "error with group_uid 0 must not record error penalty in latency registry" + ); + assert!( + !fallback_called.load(Ordering::SeqCst), + "gateway fallback must not be invoked prior to post-cooldown request" + ); + + // Call 3: Subsequent request in the same transaction falls back to gateway because affinity endpoint is on cooldown. + let execute_request_3 = ExecuteSqlRequest::default() + .set_session( + "projects/test-project/instances/test-instance/databases/test-db/sessions/session-1", + ) + .set_sql("SELECT 1") + .set_transaction( + TransactionSelector::default().set_id(Bytes::copy_from_slice(transaction_id)), + ); + + let result_3 = database_client + .execute_sql(execute_request_3, RequestOptions::default(), 0) + .await; + assert!( + result_3.is_ok(), + "subsequent query must succeed via gateway fallback when affinity endpoint is cooling down" + ); + assert!( + fallback_called.load(Ordering::SeqCst), + "gateway fallback must have received the post-cooldown execution" + ); + + Ok(()) +}