From c394a7f61d3c102082ae2796431cdb86a7a16da6 Mon Sep 17 00:00:00 2001 From: Evan Zhou Date: Thu, 3 Sep 2026 18:26:32 +0800 Subject: [PATCH] feat: fall back to followers when leader is unreachable Signed-off-by: Evan Zhou --- src/common/security.rs | 29 ++ src/mock.rs | 54 ++- src/request/mod.rs | 699 +++++++++++++++++++++++++++++++-- src/request/plan.rs | 863 +++++++++++++++++++++++++++++++++-------- src/store/client.rs | 2 +- 5 files changed, 1455 insertions(+), 192 deletions(-) diff --git a/src/common/security.rs b/src/common/security.rs index 89d60060..d8d013ec 100644 --- a/src/common/security.rs +++ b/src/common/security.rs @@ -82,6 +82,31 @@ impl SecurityManager { addr: &str, factory: Factory, ) -> Result + where + Factory: FnOnce(Channel) -> Client, + { + self.connect_inner(addr, None, factory).await + } + + /// Connect with an explicit TCP connection timeout. + pub(crate) async fn connect_with_timeout( + &self, + addr: &str, + timeout: Duration, + factory: Factory, + ) -> Result + where + Factory: FnOnce(Channel) -> Client, + { + self.connect_inner(addr, Some(timeout), factory).await + } + + async fn connect_inner( + &self, + addr: &str, + timeout: Option, + factory: Factory, + ) -> Result where Factory: FnOnce(Channel) -> Client, { @@ -91,6 +116,10 @@ impl SecurityManager { } else { self.default_channel(addr).await? }; + let channel = match timeout { + Some(timeout) => channel.connect_timeout(timeout), + None => channel, + }; let ch = channel.connect().await?; Ok(factory(ch)) diff --git a/src/mock.rs b/src/mock.rs index 3279bad2..0a1c1b77 100644 --- a/src/mock.rs +++ b/src/mock.rs @@ -18,6 +18,7 @@ use crate::proto::keyspacepb; use crate::proto::metapb::RegionEpoch; use crate::proto::metapb::{self}; use crate::region::RegionId; +use crate::region::RegionVerId; use crate::region::RegionWithLeader; use crate::store::KvConnect; use crate::store::RegionStore; @@ -76,11 +77,22 @@ pub struct MockCluster; #[derive(new)] pub struct MockPdClient { client: MockKvClient, + /// Optional override for `map_region_to_store`. + #[new(default)] + map_region_to_store_hook: + Option Result + Send + Sync + 'static>>, /// Optional override for `region_for_key`, e.g. to simulate PD failing to /// locate a region for a key. #[new(default)] region_for_key_hook: Option Result + Send + Sync + 'static>>, + /// Optional observer/override for leader cache updates. + #[new(default)] + update_leader_hook: + Option Result<()> + Send + Sync + 'static>>, + /// Optional observer for region cache invalidations. + #[new(default)] + invalidate_region_hook: Option>, } #[async_trait] @@ -110,6 +122,14 @@ impl MockPdClient { MockPdClient::new(MockKvClient::default()) } + pub fn with_map_region_to_store_hook(mut self, hook: F) -> MockPdClient + where + F: Fn(RegionWithLeader) -> Result + Send + Sync + 'static, + { + self.map_region_to_store_hook = Some(Arc::new(hook)); + self + } + /// Override `region_for_key` with a custom hook, leaving the rest of the /// mock's behavior untouched. pub fn with_region_for_key_hook(mut self, hook: F) -> MockPdClient @@ -120,6 +140,22 @@ impl MockPdClient { self } + pub fn with_update_leader_hook(mut self, hook: F) -> MockPdClient + where + F: Fn(RegionVerId, metapb::Peer) -> Result<()> + Send + Sync + 'static, + { + self.update_leader_hook = Some(Arc::new(hook)); + self + } + + pub fn with_invalidate_region_hook(mut self, hook: F) -> MockPdClient + where + F: Fn(RegionVerId) + Send + Sync + 'static, + { + self.invalidate_region_hook = Some(Arc::new(hook)); + self + } + pub fn region1() -> RegionWithLeader { let mut region = RegionWithLeader::default(); region.region.id = 1; @@ -183,6 +219,9 @@ impl PdClient for MockPdClient { type KvClient = MockKvClient; async fn map_region_to_store(self: Arc, region: RegionWithLeader) -> Result { + if let Some(hook) = &self.map_region_to_store_hook { + return hook(region); + } Ok(RegionStore::new(region, Arc::new(self.client.clone()))) } @@ -225,13 +264,20 @@ impl PdClient for MockPdClient { async fn update_leader( &self, - _ver_id: crate::region::RegionVerId, - _leader: metapb::Peer, + ver_id: crate::region::RegionVerId, + leader: metapb::Peer, ) -> Result<()> { - todo!() + match &self.update_leader_hook { + Some(hook) => hook(ver_id, leader), + None => todo!(), + } } - async fn invalidate_region_cache(&self, _ver_id: crate::region::RegionVerId) {} + async fn invalidate_region_cache(&self, ver_id: crate::region::RegionVerId) { + if let Some(hook) = &self.invalidate_region_hook { + hook(ver_id); + } + } async fn invalidate_store_cache(&self, _store_id: crate::region::StoreId) {} diff --git a/src/request/mod.rs b/src/request/mod.rs index 4b8415c8..4cac6c50 100644 --- a/src/request/mod.rs +++ b/src/request/mod.rs @@ -90,7 +90,7 @@ impl RetryOptions { mod test { use std::any::Any; use std::iter; - use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -211,7 +211,7 @@ mod test { } #[tokio::test] - async fn test_region_store_mapping_retry() { + async fn test_fallback_mapping_error_tries_remaining_voter() { #[derive(Debug, Clone)] struct MockOkResponse; @@ -231,12 +231,13 @@ mod test { struct FlakyStoreMappingPdClient { client: MockKvClient, - invalidated: AtomicBool, - invalidation_count: AtomicUsize, + mapped_store_ids: std::sync::Mutex>, + invalidate_region_count: AtomicUsize, + invalidated_store_ids: std::sync::Mutex>, } impl FlakyStoreMappingPdClient { - fn region(store_id: StoreId) -> RegionWithLeader { + fn region() -> RegionWithLeader { let mut region = RegionWithLeader::default(); region.region.id = 1; region.region.start_key = vec![]; @@ -245,10 +246,23 @@ mod test { conf_ver: 0, version: 0, }); - region.leader = Some(metapb::Peer { - store_id, + let leader = metapb::Peer { + id: 1, + store_id: 41, ..Default::default() - }); + }; + let follower = metapb::Peer { + id: 2, + store_id: 42, + ..Default::default() + }; + let second_follower = metapb::Peer { + id: 3, + store_id: 43, + ..Default::default() + }; + region.region.peers = vec![leader.clone(), follower, second_follower]; + region.leader = Some(leader); region } } @@ -261,21 +275,22 @@ mod test { self: Arc, region: RegionWithLeader, ) -> Result { - match region.get_store_id()? { - 41 => Err(Error::InternalError { - message: "invalid store ID 41, not found".to_owned(), - }), - _ => Ok(RegionStore::new(region, Arc::new(self.client.clone()))), + let store_id = region.get_store_id()?; + { + let mut mapped_store_ids = self.mapped_store_ids.lock().unwrap(); + mapped_store_ids.push(store_id); + } + if store_id == 42 { + Err(Error::GrpcAPI(tonic::Status::unavailable( + "follower mapping failure", + ))) + } else { + Ok(RegionStore::new(region, Arc::new(self.client.clone()))) } } async fn region_for_key(&self, _: &Key) -> Result { - let store_id = if self.invalidated.load(Ordering::SeqCst) { - 42 - } else { - 41 - }; - Ok(Self::region(store_id)) + Ok(Self::region()) } async fn region_for_id(&self, id: RegionId) -> Result { @@ -310,11 +325,12 @@ mod test { } async fn invalidate_region_cache(&self, _ver_id: RegionVerId) { - self.invalidated.store(true, Ordering::SeqCst); - self.invalidation_count.fetch_add(1, Ordering::SeqCst); + self.invalidate_region_count.fetch_add(1, Ordering::SeqCst); } - async fn invalidate_store_cache(&self, _store_id: StoreId) {} + async fn invalidate_store_cache(&self, store_id: StoreId) { + self.invalidated_store_ids.lock().unwrap().push(store_id); + } } #[derive(Clone)] @@ -378,11 +394,17 @@ mod test { let pd_client = Arc::new(FlakyStoreMappingPdClient { client: MockKvClient::with_dispatch_hook(move |_: &dyn Any| { - dispatch_count_clone.fetch_add(1, Ordering::SeqCst); - Ok(Box::new(MockOkResponse) as Box) + if dispatch_count_clone.fetch_add(1, Ordering::SeqCst) == 0 { + Err(Error::GrpcAPI(tonic::Status::unavailable( + "leader unavailable", + ))) + } else { + Ok(Box::new(MockOkResponse) as Box) + } }), - invalidated: AtomicBool::new(false), - invalidation_count: AtomicUsize::new(0), + mapped_store_ids: std::sync::Mutex::new(Vec::new()), + invalidate_region_count: AtomicUsize::new(0), + invalidated_store_ids: std::sync::Mutex::new(Vec::new()), }); let request = MockKvRequest { @@ -395,9 +417,76 @@ mod test { let response = plan.execute().await; assert!(response.is_ok()); - assert_eq!(dispatch_count.load(Ordering::SeqCst), 1); - assert_eq!(shard_invoking_count.load(Ordering::SeqCst), 2); - assert_eq!(pd_client.invalidation_count.load(Ordering::SeqCst), 1); + assert_eq!(dispatch_count.load(Ordering::SeqCst), 2); + assert_eq!(shard_invoking_count.load(Ordering::SeqCst), 1); + assert_eq!( + pd_client.mapped_store_ids.lock().unwrap().as_slice(), + &[41, 42, 43], + "a failed follower mapping must not hide the remaining healthy voter" + ); + assert_eq!(pd_client.invalidate_region_count.load(Ordering::SeqCst), 0); + assert_eq!( + pd_client.invalidated_store_ids.lock().unwrap().as_slice(), + &[41, 42], + "both the failed leader RPC and failed follower mapping invalidate their Store entries" + ); + } + + #[tokio::test] + async fn test_all_mapping_errors_reload_stale_region() { + let old_region = fallback_region(&[44]); + let mut new_region = fallback_region(&[]); + let new_leader = metapb::Peer { + id: 3, + store_id: 45, + ..Default::default() + }; + new_region.region.peers = vec![new_leader.clone()]; + new_region.leader = Some(new_leader); + + let region_invalidated = Arc::new(AtomicBool::new(false)); + let locate_invalidated = region_invalidated.clone(); + let invalidate_observer = region_invalidated.clone(); + let mapping_attempts = Arc::new(std::sync::Mutex::new(Vec::new())); + let tracked_mapping_attempts = mapping_attempts.clone(); + let client = MockKvClient::with_dispatch_hook(|_: &dyn Any| { + Ok(Box::new(kvrpcpb::GetResponse::default()) as Box) + }); + let mapping_client = client.clone(); + let pd_client = Arc::new( + MockPdClient::new(client) + .with_region_for_key_hook(move |_| { + if locate_invalidated.load(Ordering::SeqCst) { + Ok(new_region.clone()) + } else { + Ok(old_region.clone()) + } + }) + .with_map_region_to_store_hook(move |region| { + let store_id = region.get_store_id()?; + tracked_mapping_attempts.lock().unwrap().push(store_id); + if store_id == 41 || store_id == 44 { + Err(Error::GrpcAPI(tonic::Status::unavailable( + "stale peer mapping failed", + ))) + } else { + Ok(RegionStore::new(region, Arc::new(mapping_client.clone()))) + } + }) + .with_invalidate_region_hook(move |_| { + invalidate_observer.store(true, Ordering::SeqCst); + }), + ); + + let response = execute_fallback_get(pd_client, Backoff::no_jitter_backoff(1, 1, 1)).await; + + assert!(response.is_ok()); + assert!(region_invalidated.load(Ordering::SeqCst)); + assert_eq!( + mapping_attempts.lock().unwrap().as_slice(), + &[41, 44, 45], + "after every stale peer fails to map, retry must load the replacement Region" + ); } #[tokio::test] @@ -603,4 +692,556 @@ mod test { assert_eq!(pd_client.invalidate_region_count.load(Ordering::SeqCst), 1); assert_eq!(pd_client.invalidate_store_count.load(Ordering::SeqCst), 1); } + + #[tokio::test] + async fn test_unreachable_leader_falls_back_to_voter_without_replica_read() { + let mut region = RegionWithLeader::default(); + region.region.id = 1; + region.region.region_epoch = Some(RegionEpoch { + conf_ver: 1, + version: 1, + }); + let leader = metapb::Peer { + id: 1, + store_id: 41, + ..Default::default() + }; + let learner = metapb::Peer { + id: 2, + store_id: 42, + role: metapb::PeerRole::Learner as i32, + ..Default::default() + }; + let witness = metapb::Peer { + id: 3, + store_id: 43, + is_witness: true, + ..Default::default() + }; + let follower = metapb::Peer { + id: 4, + store_id: 44, + ..Default::default() + }; + region.region.peers = vec![leader.clone(), learner, witness, follower.clone()]; + region.leader = Some(leader); + + let accesses = Arc::new(std::sync::Mutex::new(Vec::new())); + let dispatch_accesses = accesses.clone(); + let client = MockKvClient::with_dispatch_hook(move |request: &dyn Any| { + let context = fallback_get_context(request); + let store_id = context.peer.as_ref().expect("target peer").store_id; + dispatch_accesses + .lock() + .unwrap() + .push((store_id, context.replica_read)); + if store_id == 41 { + Err(Error::GrpcAPI(tonic::Status::unavailable( + "leader unavailable", + ))) + } else { + Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + not_leader: Some(crate::proto::errorpb::NotLeader { + region_id: 1, + leader: None, + }), + ..Default::default() + }), + ..Default::default() + }) as Box) + } + }); + let located_region = region.clone(); + let pd_client = Arc::new( + MockPdClient::new(client).with_region_for_key_hook(move |_| Ok(located_region.clone())), + ); + let response = execute_fallback_get(pd_client, Backoff::no_backoff()).await; + + assert!(response.is_err()); + assert_eq!( + accesses.lock().unwrap().as_slice(), + &[(41, false), (44, false)], + "leader is tried first; learner and witness are skipped; fallback stays a normal request" + ); + } + + fn fallback_region(follower_store_ids: &[StoreId]) -> RegionWithLeader { + let mut region = RegionWithLeader::default(); + region.region.id = 1; + region.region.region_epoch = Some(RegionEpoch { + conf_ver: 1, + version: 1, + }); + let leader = metapb::Peer { + id: 1, + store_id: 41, + ..Default::default() + }; + region.region.peers.push(leader.clone()); + region + .region + .peers + .extend( + follower_store_ids + .iter() + .enumerate() + .map(|(index, store_id)| metapb::Peer { + id: index as u64 + 2, + store_id: *store_id, + ..Default::default() + }), + ); + region.leader = Some(leader); + region + } + + fn fallback_get_context(request: &dyn Any) -> &kvrpcpb::Context { + request + .downcast_ref::() + .expect("get request") + .context + .as_ref() + .expect("request context") + } + + async fn execute_fallback_get( + pd_client: Arc, + backoff: Backoff, + ) -> Result>> { + PlanBuilder::new( + pd_client, + Keyspace::Disable, + kvrpcpb::GetRequest { + key: b"key".to_vec(), + ..Default::default() + }, + ) + .retry_multi_region(backoff) + .plan() + .execute() + .await + } + + #[tokio::test] + async fn test_not_leader_follower_does_not_hide_remaining_voters() { + let accesses = Arc::new(std::sync::Mutex::new(Vec::new())); + let dispatch_accesses = accesses.clone(); + let client = MockKvClient::with_dispatch_hook(move |request: &dyn Any| { + let context = fallback_get_context(request); + let store_id = context.peer.as_ref().expect("target peer").store_id; + dispatch_accesses + .lock() + .unwrap() + .push((store_id, context.replica_read)); + match store_id { + 41 => Err(Error::GrpcAPI(tonic::Status::unavailable( + "leader unavailable", + ))), + 44 => Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + not_leader: Some(crate::proto::errorpb::NotLeader { + region_id: 1, + leader: None, + }), + ..Default::default() + }), + ..Default::default() + }) as Box), + 45 => Ok(Box::new(kvrpcpb::GetResponse::default()) as Box), + _ => unreachable!(), + } + }); + let region = fallback_region(&[44, 45]); + let located_region = region.clone(); + let updated_leader_store_id = Arc::new(AtomicU64::new(0)); + let tracked_leader_store_id = updated_leader_store_id.clone(); + let pd_client = Arc::new( + MockPdClient::new(client) + .with_region_for_key_hook(move |_| Ok(located_region.clone())) + .with_update_leader_hook(move |_, leader| { + tracked_leader_store_id.store(leader.store_id, Ordering::SeqCst); + Ok(()) + }), + ); + + let response = execute_fallback_get(pd_client, Backoff::no_backoff()).await; + + assert!(response.is_ok()); + assert_eq!( + accesses.lock().unwrap().as_slice(), + &[(41, false), (44, false), (45, false)] + ); + assert_eq!(updated_leader_store_id.load(Ordering::SeqCst), 45); + } + + #[tokio::test] + async fn test_follower_leader_hint_is_ignored() { + let accesses = Arc::new(std::sync::Mutex::new(Vec::new())); + let dispatch_accesses = accesses.clone(); + let client = MockKvClient::with_dispatch_hook(move |request: &dyn Any| { + let store_id = fallback_get_context(request) + .peer + .as_ref() + .expect("target peer") + .store_id; + dispatch_accesses.lock().unwrap().push(store_id); + if store_id == 41 { + return Err(Error::GrpcAPI(tonic::Status::unavailable( + "leader unavailable", + ))); + } + + let leader = (store_id == 44).then_some(metapb::Peer { + id: 3, + store_id: 45, + ..Default::default() + }); + Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + not_leader: Some(crate::proto::errorpb::NotLeader { + region_id: 1, + // Store 44 points at store 45, but fallback deliberately + // ignores follower hints. Store 45 is tried naturally + // and independently reports that it is not the leader. + leader, + }), + ..Default::default() + }), + ..Default::default() + }) as Box) + }); + let region = fallback_region(&[44, 45]); + let located_region = region.clone(); + let update_leader_count = Arc::new(AtomicUsize::new(0)); + let tracked_update_leader_count = update_leader_count.clone(); + let pd_client = Arc::new( + MockPdClient::new(client) + .with_region_for_key_hook(move |_| Ok(located_region.clone())) + .with_update_leader_hook(move |_, _| { + tracked_update_leader_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + }), + ); + + let response = execute_fallback_get(pd_client, Backoff::no_jitter_backoff(1, 1, 1)).await; + + assert!(matches!( + response, + Err(Error::GrpcAPI(status)) if status.code() == tonic::Code::Unavailable + )); + assert_eq!( + accesses.lock().unwrap().as_slice(), + &[41, 44, 45, 41, 44, 45] + ); + assert_eq!( + update_leader_count.load(Ordering::SeqCst), + 0, + "a fallback follower's leader hint must not update the cache" + ); + } + + #[tokio::test] + async fn test_invalidating_routing_error_outweighs_fallback_server_busy() { + let accesses = Arc::new(std::sync::Mutex::new(Vec::new())); + let dispatch_accesses = accesses.clone(); + let client = MockKvClient::with_dispatch_hook(move |request: &dyn Any| { + let store_id = fallback_get_context(request) + .peer + .as_ref() + .expect("target peer") + .store_id; + dispatch_accesses.lock().unwrap().push(store_id); + if store_id == 41 { + Err(Error::GrpcAPI(tonic::Status::unavailable( + "cached leader unavailable", + ))) + } else { + Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + server_is_busy: Some(crate::proto::errorpb::ServerIsBusy::default()), + ..Default::default() + }), + ..Default::default() + }) as Box) + } + }); + let region = fallback_region(&[44]); + let located_region = region.clone(); + let pd_client = Arc::new( + MockPdClient::new(client).with_region_for_key_hook(move |_| Ok(located_region.clone())), + ); + + let response = execute_fallback_get(pd_client, Backoff::no_backoff()).await; + + assert!(matches!( + response, + Err(Error::GrpcAPI(status)) if status.code() == tonic::Code::Unavailable + )); + assert_eq!(accesses.lock().unwrap().as_slice(), &[41, 44]); + } + + #[tokio::test] + async fn test_fallback_not_leader_invalidates_stale_region() { + let accesses = Arc::new(std::sync::Mutex::new(Vec::new())); + let dispatch_accesses = accesses.clone(); + let client = MockKvClient::with_dispatch_hook(move |request: &dyn Any| { + let store_id = fallback_get_context(request) + .peer + .as_ref() + .expect("target peer") + .store_id; + dispatch_accesses.lock().unwrap().push(store_id); + if store_id == 41 { + Err(Error::GrpcAPI(tonic::Status::deadline_exceeded( + "cached leader timed out", + ))) + } else { + Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + not_leader: Some(crate::proto::errorpb::NotLeader { + region_id: 1, + leader: None, + }), + ..Default::default() + }), + ..Default::default() + }) as Box) + } + }); + let region = fallback_region(&[44]); + let located_region = region.clone(); + let invalidate_count = Arc::new(AtomicUsize::new(0)); + let tracked_invalidate_count = invalidate_count.clone(); + let pd_client = Arc::new( + MockPdClient::new(client) + .with_region_for_key_hook(move |_| Ok(located_region.clone())) + .with_invalidate_region_hook(move |_| { + tracked_invalidate_count.fetch_add(1, Ordering::SeqCst); + }), + ); + + let response = execute_fallback_get(pd_client, Backoff::no_backoff()).await; + + assert!(matches!(response, Err(Error::LeaderNotFound { .. }))); + assert_eq!(accesses.lock().unwrap().as_slice(), &[41, 44]); + assert_eq!(invalidate_count.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_nonzero_server_busy_reloads_cached_leader() { + let accesses = Arc::new(std::sync::Mutex::new(Vec::new())); + let dispatch_accesses = accesses.clone(); + let leader_busy = Arc::new(AtomicBool::new(false)); + let dispatch_leader_busy = leader_busy.clone(); + let client = MockKvClient::with_dispatch_hook(move |request: &dyn Any| { + let context = fallback_get_context(request); + let store_id = context.peer.as_ref().expect("target peer").store_id; + dispatch_accesses + .lock() + .unwrap() + .push((store_id, context.replica_read)); + if store_id == 41 { + dispatch_leader_busy.store(true, Ordering::SeqCst); + Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + server_is_busy: Some(crate::proto::errorpb::ServerIsBusy { + estimated_wait_ms: 500, + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }) as Box) + } else { + Ok(Box::new(kvrpcpb::GetResponse::default()) as Box) + } + }); + let initial_region = fallback_region(&[44]); + let mut updated_region = initial_region.clone(); + updated_region.leader = updated_region + .region + .peers + .iter() + .find(|peer| peer.store_id == 44) + .cloned(); + let pd_client = Arc::new( + MockPdClient::new(client).with_region_for_key_hook(move |_| { + if leader_busy.load(Ordering::SeqCst) { + Ok(updated_region.clone()) + } else { + Ok(initial_region.clone()) + } + }), + ); + + let response = execute_fallback_get(pd_client, Backoff::no_jitter_backoff(1, 1, 1)).await; + + assert!(response.is_ok()); + assert_eq!( + accesses.lock().unwrap().as_slice(), + &[(41, false), (44, false)], + "a nonzero busy response must re-read the shared leader cache" + ); + } + + #[tokio::test] + async fn test_second_server_busy_zero_probes_followers_without_replica_read() { + let accesses = Arc::new(std::sync::Mutex::new(Vec::new())); + let dispatch_accesses = accesses.clone(); + let client = MockKvClient::with_dispatch_hook(move |request: &dyn Any| { + let context = fallback_get_context(request); + let store_id = context.peer.as_ref().expect("target peer").store_id; + dispatch_accesses + .lock() + .unwrap() + .push((store_id, context.replica_read)); + match store_id { + 41 => Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + server_is_busy: Some(crate::proto::errorpb::ServerIsBusy::default()), + ..Default::default() + }), + ..Default::default() + }) as Box), + 44 => Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + not_leader: Some(crate::proto::errorpb::NotLeader { + region_id: 1, + leader: None, + }), + ..Default::default() + }), + ..Default::default() + }) as Box), + 45 => Ok(Box::new(kvrpcpb::GetResponse::default()) as Box), + _ => unreachable!(), + } + }); + let region = fallback_region(&[44, 45]); + let located_region = region.clone(); + let updated_leader_store_id = Arc::new(AtomicU64::new(0)); + let tracked_leader_store_id = updated_leader_store_id.clone(); + let pd_client = Arc::new( + MockPdClient::new(client) + .with_region_for_key_hook(move |_| Ok(located_region.clone())) + .with_update_leader_hook(move |_, leader| { + tracked_leader_store_id.store(leader.store_id, Ordering::SeqCst); + Ok(()) + }), + ); + + let response = execute_fallback_get(pd_client, Backoff::no_jitter_backoff(1, 1, 2)).await; + + assert!(response.is_ok()); + assert_eq!( + accesses.lock().unwrap().as_slice(), + &[(41, false), (41, false), (44, false), (45, false)], + "the same leader is retried once, then followers are probed with leader-read semantics" + ); + assert_eq!(updated_leader_store_id.load(Ordering::SeqCst), 45); + } + + #[tokio::test] + async fn test_server_busy_follower_probe_is_one_shot() { + let accesses = Arc::new(std::sync::Mutex::new(Vec::new())); + let dispatch_accesses = accesses.clone(); + let leader_attempts = Arc::new(AtomicUsize::new(0)); + let dispatch_leader_attempts = leader_attempts.clone(); + let client = MockKvClient::with_dispatch_hook(move |request: &dyn Any| { + let context = fallback_get_context(request); + let store_id = context.peer.as_ref().expect("target peer").store_id; + dispatch_accesses + .lock() + .unwrap() + .push((store_id, context.replica_read)); + if store_id == 41 { + if dispatch_leader_attempts.fetch_add(1, Ordering::SeqCst) < 3 { + return Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + server_is_busy: Some(crate::proto::errorpb::ServerIsBusy::default()), + ..Default::default() + }), + ..Default::default() + }) as Box); + } + return Ok(Box::new(kvrpcpb::GetResponse::default()) as Box); + } + + Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(crate::proto::errorpb::Error { + not_leader: Some(crate::proto::errorpb::NotLeader { + region_id: 1, + leader: None, + }), + ..Default::default() + }), + ..Default::default() + }) as Box) + }); + let region = fallback_region(&[44, 45]); + let located_region = region.clone(); + let pd_client = Arc::new( + MockPdClient::new(client).with_region_for_key_hook(move |_| Ok(located_region.clone())), + ); + + let response = execute_fallback_get(pd_client, Backoff::no_jitter_backoff(1, 1, 3)).await; + + assert!(response.is_ok()); + assert_eq!( + accesses.lock().unwrap().as_slice(), + &[ + (41, false), + (41, false), + (44, false), + (45, false), + (41, false), + (41, false), + ], + "a fruitless probe restores the cached leader and is not fired again" + ); + } + + #[tokio::test] + async fn test_fallback_key_error_updates_leader_cache() { + let client = MockKvClient::with_dispatch_hook(move |request: &dyn Any| { + let store_id = fallback_get_context(request) + .peer + .as_ref() + .expect("target peer") + .store_id; + if store_id == 41 { + Err(Error::GrpcAPI(tonic::Status::unavailable( + "leader unavailable", + ))) + } else { + Ok(Box::new(kvrpcpb::GetResponse { + error: Some(kvrpcpb::KeyError::default()), + ..Default::default() + }) as Box) + } + }); + let region = fallback_region(&[44]); + let located_region = region.clone(); + let updated_leader_store_id = Arc::new(AtomicU64::new(0)); + let tracked_leader_store_id = updated_leader_store_id.clone(); + let pd_client = Arc::new( + MockPdClient::new(client) + .with_region_for_key_hook(move |_| Ok(located_region.clone())) + .with_update_leader_hook(move |_, leader| { + tracked_leader_store_id.store(leader.store_id, Ordering::SeqCst); + Ok(()) + }), + ); + + let response = execute_fallback_get(pd_client, Backoff::no_backoff()) + .await + .expect("key errors stay in the per-region result"); + + assert!(matches!( + response.as_slice(), + [Err(Error::MultipleKeyErrors(_))] + )); + assert_eq!(updated_leader_store_id.load(Ordering::SeqCst), 44); + } } diff --git a/src/request/plan.rs b/src/request/plan.rs index f832f608..e1a1eb5b 100644 --- a/src/request/plan.rs +++ b/src/request/plan.rs @@ -13,12 +13,14 @@ use log::warn; use tokio::sync::Semaphore; use tokio::task::JoinSet; use tokio::time::sleep; +use tonic::Code; use crate::backoff::Backoff; use crate::pd::PdClient; use crate::proto::errorpb; use crate::proto::errorpb::EpochNotMatch; use crate::proto::kvrpcpb; +use crate::proto::metapb; use crate::proto::pdpb::Timestamp; use crate::region::StoreId; use crate::region::{RegionVerId, RegionWithLeader}; @@ -175,6 +177,182 @@ pub struct RetryableMultiRegion { pub terminal_on_dispatch_error: bool, } +struct CandidateResponse { + response: R, + key_errors: Option>, + region_error: Option, + region_store: RegionStore, + used_fallback: bool, +} + +impl CandidateResponse +where + R: HasKeyErrors + HasRegionError, +{ + fn new(mut response: R, region_store: RegionStore, used_fallback: bool) -> Self { + let key_errors = response.key_errors(); + let region_error = response.region_error(); + Self { + response, + key_errors, + region_error, + region_store, + used_fallback, + } + } + + fn is_fallback_not_leader(&self) -> bool { + self.used_fallback + && self.key_errors.is_none() + && self + .region_error + .as_ref() + .is_some_and(|error| error.not_leader.is_some()) + } + + fn is_fallback_server_busy(&self) -> bool { + self.used_fallback + && self.key_errors.is_none() + && self + .region_error + .as_ref() + .is_some_and(|error| error.server_is_busy.is_some()) + } + + fn accepted_fallback_leader(&self) -> Option { + if !self.used_fallback || self.region_error.is_some() { + return None; + } + self.region_store.region_with_leader.leader.clone() + } +} + +enum CandidateRoundResult { + Response(Box>), + MapRegionToStoreError(Error), + RoutingError { + error: Error, + invalidate_region: bool, + }, + OtherError(Error), +} + +#[derive(Default)] +struct CandidateState { + busy_leader_peer_id: Option, + busy_count: u8, + follower_probe_done: bool, +} + +impl CandidateState { + const LEADER_BUSY_PROBE_THRESHOLD: u8 = 2; + + fn record_leader_busy(&mut self, peer_id: u64, estimated_wait_ms: u32) { + if estimated_wait_ms != 0 || self.follower_probe_done { + return; + } + if self.busy_leader_peer_id != Some(peer_id) { + self.busy_leader_peer_id = Some(peer_id); + self.busy_count = 0; + } + self.busy_count = self.busy_count.saturating_add(1); + } + + fn take_follower_probe(&mut self) -> bool { + if !self.follower_probe_done && self.busy_count >= Self::LEADER_BUSY_PROBE_THRESHOLD { + self.follower_probe_done = true; + true + } else { + false + } + } +} + +struct RetryContext { + pd_client: Arc, + permits: Arc, + preserve_region_results: bool, + terminal_on_undetermined: bool, + terminal_on_dispatch_error: bool, +} + +impl Clone for RetryContext { + fn clone(&self) -> Self { + Self { + pd_client: self.pd_client.clone(), + permits: self.permits.clone(), + preserve_region_results: self.preserve_region_results, + terminal_on_undetermined: self.terminal_on_undetermined, + terminal_on_dispatch_error: self.terminal_on_dispatch_error, + } + } +} + +enum GrpcErrorAction { + TryNextPeer, + TryNextPeerAndInvalidate, + Return, +} + +fn grpc_error_action(error: &Error) -> GrpcErrorAction { + match error { + // A transport setup error means this Store cannot currently be used. + Error::Grpc(_) => GrpcErrorAction::TryNextPeerAndInvalidate, + Error::GrpcAPI(status) if std::error::Error::source(status).is_some() => { + // Tonic maps lower-level I/O and HTTP/2 failures to Status, and + // keeps the transport error as its source. Its code is not always + // Unavailable (for example, an HTTP/2 failure can be Internal or + // ResourceExhausted), so inspect the source before the code. + GrpcErrorAction::TryNextPeerAndInvalidate + } + Error::GrpcAPI(status) => match status.code() { + // client-go retries a remote/keepalive cancellation after replacing + // the connection. Dropping a Rust request future cancels this whole + // task, so a Canceled status observed here came from the RPC side. + // Tonic can synthesize source-less Internal/Unknown statuses when + // an HTTP/2 stream ends before a complete response is decoded + // (for example "Missing response message"). Those are + // indistinguishable here from an application status, but treating + // them as terminal would miss the cold-region recovery path when + // the cached leader disconnects mid-response. + Code::Unavailable | Code::Cancelled | Code::Internal | Code::Unknown => { + GrpcErrorAction::TryNextPeerAndInvalidate + } + // A request deadline does not prove that the Store or Region route + // is stale. Try another replica, but preserve the shared caches. + Code::DeadlineExceeded => GrpcErrorAction::TryNextPeer, + // Application-level statuses (invalid arguments, auth failures, + // unsupported APIs, explicit server resource limits, etc.) cannot + // be healed by replaying the request on every replica. client-go + // uses a separate Store health RPC to distinguish this case; Rust + // has no equivalent liveness subsystem, so an explicit status with + // no transport source is the strongest available signal. + _ => GrpcErrorAction::Return, + }, + _ => GrpcErrorAction::Return, + } +} + +fn region_request_candidates( + region: &RegionWithLeader, + followers_only: bool, +) -> impl Iterator + '_ { + let leader_store_id = region.leader.as_ref().map(|leader| leader.store_id); + // A request to an unreachable cached leader cannot wake a cold Raft + // group. Lazily try each data-bearing voter after the cached leader + // before invalidating it and going back to PD. Learners and witnesses + // cannot campaign and are not useful for cold-region recovery. + region + .leader + .iter() + .filter(move |_| !followers_only) + .chain(region.region.peers.iter().filter(move |peer| { + !peer.is_witness + && peer.role != metapb::PeerRole::Learner as i32 + && leader_store_id != Some(peer.store_id) + })) +} + impl RetryableMultiRegion where P::Result: HasKeyErrors + HasRegionError, @@ -182,15 +360,14 @@ where // A plan may involve multiple shards #[async_recursion] async fn single_plan_handler( - pd_client: Arc, + context: RetryContext, current_plan: P, backoff: Backoff, - permits: Arc, - preserve_region_results: bool, - terminal_on_undetermined: bool, - terminal_on_dispatch_error: bool, ) -> Result<::Result> { - let shards = current_plan.shards(&pd_client).collect::>().await; + let shards = current_plan + .shards(&context.pd_client) + .collect::>() + .await; let shards_len = shards.len(); debug!("single_plan_handler, shards: {}", shards_len); let mut join_set = JoinSet::new(); @@ -203,21 +380,17 @@ where } }; let clone = current_plan.clone_then_apply_shard(shard); - let pd_client = pd_client.clone(); + let shard_context = context.clone(); let backoff = backoff.clone(); - let permits = permits.clone(); join_set.spawn(async move { ( idx, Self::single_shard_handler( - pd_client, + shard_context, clone, region, backoff, - permits, - preserve_region_results, - terminal_on_undetermined, - terminal_on_dispatch_error, + CandidateState::default(), ) .await, ) @@ -226,7 +399,7 @@ where let results = collect_join_set_results(join_set, shards_len, "single_plan_handler").await?; - if preserve_region_results { + if context.preserve_region_results { Ok(results .into_iter() .flat_map_ok(|x| x) @@ -263,17 +436,239 @@ where } } - #[allow(clippy::too_many_arguments)] + async fn execute_candidate( + pd_client: &Arc, + plan: &mut P, + region: &RegionWithLeader, + peer: &metapb::Peer, + original_leader_store_id: Option, + permits: &Semaphore, + terminal_on_dispatch_error: bool, + ) -> CandidateRoundResult { + let is_fallback = original_leader_store_id != Some(peer.store_id); + let mut candidate_region = region.clone(); + candidate_region.leader = Some(peer.clone()); + + let region_store = match pd_client + .clone() + .map_region_to_store(candidate_region) + .await + { + Ok(region_store) => region_store, + Err(err) => { + debug!( + "single_shard_handler::map_store, fallback: {}, error: {:?}", + is_fallback, err + ); + // Mapping includes opening the TiKV channel. Drop the Store + // entry so a later attempt can reload a changed address from + // PD instead of reconnecting to stale metadata indefinitely. + pd_client.invalidate_store_cache(peer.store_id).await; + return CandidateRoundResult::MapRegionToStoreError(err); + } + }; + if let Err(error) = plan.apply_store(®ion_store) { + // Applying a successfully mapped Store mutates the request. This is + // not a PD/connection lookup failure and trying another peer cannot + // repair an invalid request. + return CandidateRoundResult::OtherError(error); + } + + // Fallback attempts are sequential inside one shard, so at most one + // concurrency permit is held at a time. + let permit = permits.acquire().await.unwrap(); + let result = plan.execute().await; + drop(permit); + + match result { + Ok(response) => CandidateRoundResult::Response(Box::new(CandidateResponse::new( + response, + region_store, + is_fallback, + ))), + Err(error) if is_grpc_error(&error) => { + debug!( + "single_shard_handler:execute: grpc error, fallback: {}, error: {:?}", + is_fallback, error + ); + // A non-idempotent request (raw CAS) may have reached the + // server and must not be replayed, including on a follower. + if terminal_on_dispatch_error { + pd_client.invalidate_region_cache(region.ver_id()).await; + pd_client.invalidate_store_cache(peer.store_id).await; + return CandidateRoundResult::OtherError(error); + } + + match grpc_error_action(&error) { + GrpcErrorAction::TryNextPeerAndInvalidate => { + pd_client.invalidate_store_cache(peer.store_id).await; + CandidateRoundResult::RoutingError { + error, + invalidate_region: true, + } + } + GrpcErrorAction::TryNextPeer => CandidateRoundResult::RoutingError { + error, + invalidate_region: false, + }, + GrpcErrorAction::Return => CandidateRoundResult::OtherError(error), + } + } + Err(error) => { + debug!("single_shard_handler:execute: error: {:?}", error); + CandidateRoundResult::OtherError(error) + } + } + } + + async fn execute_candidate_round( + pd_client: &Arc, + plan: &mut P, + region: &RegionWithLeader, + permits: &Semaphore, + terminal_on_dispatch_error: bool, + followers_only: bool, + ) -> CandidateRoundResult { + let original_leader_store_id = region.leader.as_ref().map(|peer| peer.store_id); + let mut last_map_region_to_store_error = None; + let mut last_routing_error = None; + let mut invalidate_region = false; + let mut saw_fallback_not_leader = false; + let mut last_server_busy_response = None; + + for peer in region_request_candidates(region, followers_only) { + match Self::execute_candidate( + pd_client, + plan, + region, + peer, + original_leader_store_id, + permits, + terminal_on_dispatch_error, + ) + .await + { + CandidateRoundResult::MapRegionToStoreError(error) => { + // Mapping includes Store lookup and opening a TiKV channel. + // One unavailable candidate must not hide a later healthy + // voter, especially while waking a cold Region. + last_map_region_to_store_error = Some(error); + } + CandidateRoundResult::RoutingError { + error, + invalidate_region: candidate_invalidates_region, + } => { + // Preserve the fact that at least one mapped candidate had + // a transport routing failure. A later mapping failure must + // not hide the need to reload an exhausted Region route. + invalidate_region |= candidate_invalidates_region; + last_routing_error = Some(error); + } + CandidateRoundResult::Response(response) if response.is_fallback_not_leader() => { + // A follower hint is unnecessary for the three-replica + // cold-region path: keep walking the Region's voters. A + // real new leader will accept its normal (non-replica-read) + // request later in this same round. + saw_fallback_not_leader = true; + } + CandidateRoundResult::Response(response) if response.is_fallback_server_busy() => { + // A follower probe can itself be rejected at the read-pool + // entrance. Keep trying other voters before restoring the + // cached leader. + last_server_busy_response = Some(response); + } + CandidateRoundResult::OtherError(error) => { + return CandidateRoundResult::OtherError(error); + } + response @ CandidateRoundResult::Response(_) => return response, + } + } + + if invalidate_region { + // A transport failure means the cached Region route may be stale. + // Do not let a later fallback ServerIsBusy response preserve that + // route: the follower may have rejected the request before + // validating Region membership. + if let Some(error) = last_routing_error.take() { + return CandidateRoundResult::RoutingError { + error, + invalidate_region: true, + }; + } + } + + if saw_fallback_not_leader && !followers_only { + // The cached peer list may be stale even though follower hints are + // deliberately ignored. Reload it after exhausting every voter. + return CandidateRoundResult::RoutingError { + error: Error::LeaderNotFound { + region: region.ver_id(), + }, + invalidate_region: true, + }; + } + + if let Some(response) = last_server_busy_response { + CandidateRoundResult::Response(response) + } else if let Some(error) = last_routing_error { + CandidateRoundResult::RoutingError { + error, + invalidate_region, + } + } else if let Some(error) = last_map_region_to_store_error { + CandidateRoundResult::MapRegionToStoreError(error) + } else { + CandidateRoundResult::RoutingError { + error: Error::LeaderNotFound { + region: region.ver_id(), + }, + invalidate_region: !followers_only, + } + } + } + + async fn update_leader_after_fallback( + pd_client: &Arc, + region_ver_id: &RegionVerId, + leader: Option, + ) { + // A response without a Region error proves that a normal request was + // accepted by the fallback peer as leader. Update the cache before + // returning key-level errors, which are interpreted by higher layers. + if let Some(peer) = leader { + if let Err(error) = pd_client.update_leader(region_ver_id.clone(), peer).await { + debug!( + "failed to update leader after follower fallback: {:?}", + error + ); + } + } + } + + async fn retry_same_region( + context: RetryContext, + plan: P, + region: RegionWithLeader, + mut backoff: Backoff, + candidate_state: CandidateState, + error: Error, + ) -> Result<::Result> { + match backoff.next_delay_duration() { + Some(duration) => { + sleep(duration).await; + Self::single_shard_handler(context, plan, region, backoff, candidate_state).await + } + None => Err(error), + } + } + #[async_recursion] async fn single_shard_handler( - pd_client: Arc, + context: RetryContext, mut plan: P, region: RegionWithLeader, - mut backoff: Backoff, - permits: Arc, - preserve_region_results: bool, - terminal_on_undetermined: bool, - terminal_on_dispatch_error: bool, + backoff: Backoff, + mut candidate_state: CandidateState, ) -> Result<::Result> { let region_ver_id = region.ver_id(); let store_id = region.get_store_id().ok(); @@ -281,159 +676,226 @@ where "single_shard_handler, region: {:?}, store: {:?}", region_ver_id, store_id ); - let region_store = match pd_client - .clone() - .map_region_to_store(region) - .await - .and_then(|region_store| { - plan.apply_store(®ion_store)?; - Ok(region_store) - }) { - Ok(region_store) => region_store, - Err(err) => { - debug!("single_shard_handler::sharding, error: {:?}", err); - return Self::handle_other_error( - pd_client, + + let followers_only = candidate_state.take_follower_probe(); + let response = match Self::execute_candidate_round( + &context.pd_client, + &mut plan, + ®ion, + &context.permits, + context.terminal_on_dispatch_error, + followers_only, + ) + .await + { + CandidateRoundResult::Response(response) => *response, + CandidateRoundResult::MapRegionToStoreError(error) if followers_only => { + return Self::retry_same_region( + context, + plan, + region, + backoff, + candidate_state, + error, + ) + .await; + } + CandidateRoundResult::MapRegionToStoreError(error) => { + // Every candidate failed before dispatch while mapping its Store. + // The cached peer list may have been replaced completely, so + // reload the Region instead of retrying the same stale peers. + return Self::retry_after_routing_error( + context, + plan, + region_ver_id, + backoff, + error, + ) + .await; + } + CandidateRoundResult::OtherError(error) => return Err(error), + CandidateRoundResult::RoutingError { + error: Error::LeaderNotFound { .. }, + .. + } if followers_only => { + // Every probed follower answered NotLeader. Ignore all hints + // and restore the cached leader without charging another + // backoff: the one-shot probe cannot fire again in this + // request, so the next leader error follows the normal path. + return Self::single_shard_handler(context, plan, region, backoff, candidate_state) + .await; + } + CandidateRoundResult::RoutingError { error, .. } if followers_only => { + return Self::retry_same_region( + context, + plan, + region, + backoff, + candidate_state, + error, + ) + .await; + } + CandidateRoundResult::RoutingError { + error, + invalidate_region: false, + } => { + return Self::retry_after_error(context, plan, backoff, error).await; + } + CandidateRoundResult::RoutingError { + error, + invalidate_region: true, + } => { + return Self::retry_after_routing_error( + context, plan, region_ver_id, - store_id, backoff, - permits, - preserve_region_results, - terminal_on_undetermined, - terminal_on_dispatch_error, - err, + error, ) .await; } }; - // limit concurrent requests - let permit = permits.acquire().await.unwrap(); - let res = plan.execute().await; - drop(permit); + Self::handle_candidate_response( + context, + plan, + region, + backoff, + candidate_state, + followers_only, + response, + ) + .await + } - let mut resp = match res { - Ok(resp) => resp, - Err(e) if is_grpc_error(&e) => { - debug!("single_shard_handler:execute: grpc error: {:?}", e); - // See `terminal_on_dispatch_error`: a non-idempotent request (raw - // CAS) that may have reached the server must not be replayed — the - // error surfaces UNCHANGED for the caller to handle. Only the REPLAY - // is skipped: routing is refreshed exactly as `handle_other_error` - // would, because the failure may be a dead cached leader/store and - // later requests must not keep selecting it. - if terminal_on_dispatch_error { - pd_client - .invalidate_region_cache(region_store.region_with_leader.ver_id()) - .await; - if let Ok(store_id) = region_store.region_with_leader.get_store_id() { - pd_client.invalidate_store_cache(store_id).await; - } - return Err(e); + async fn handle_candidate_response( + context: RetryContext, + plan: P, + region: RegionWithLeader, + backoff: Backoff, + candidate_state: CandidateState, + followers_only: bool, + response: CandidateResponse, + ) -> Result<::Result> { + let region_ver_id = region.ver_id(); + let fallback_leader = response.accepted_fallback_leader(); + Self::update_leader_after_fallback(&context.pd_client, ®ion_ver_id, fallback_leader) + .await; + let CandidateResponse { + response, + key_errors, + region_error, + region_store, + used_fallback, + } = response; + + if let Some(error) = key_errors { + debug!("single_shard_handler:execute: key errors: {:?}", error); + return Ok(vec![Err(Error::MultipleKeyErrors(error))]); + } + let Some(error) = region_error else { + return Ok(vec![Ok(response)]); + }; + + if let Some(server_is_busy) = error.server_is_busy.as_ref() { + if !used_fallback && server_is_busy.estimated_wait_ms == 0 { + let mut candidate_state = candidate_state; + if let Some(peer) = region_store.region_with_leader.leader.as_ref() { + candidate_state.record_leader_busy(peer.id, server_is_busy.estimated_wait_ms); } - return Self::handle_other_error( - pd_client, + return Self::retry_same_region( + context, plan, - region_store.region_with_leader.ver_id(), - region_store.region_with_leader.get_store_id().ok(), + region, backoff, - permits, - preserve_region_results, - terminal_on_undetermined, - terminal_on_dispatch_error, - e, + candidate_state, + Error::RegionError(Box::new(error)), ) .await; } - Err(e) => { - debug!("single_shard_handler:execute: error: {:?}", e); - return Err(e); - } - }; - if let Some(e) = resp.key_errors() { - debug!("single_shard_handler:execute: key errors: {:?}", e); - Ok(vec![Err(Error::MultipleKeyErrors(e))]) - } else if let Some(e) = resp.region_error() { - debug!( - "single_shard_handler:execute: region error: {:?}, region: {:?}", - e, region_ver_id - ); - // See `terminal_on_undetermined`: for CAS and commit points, an unknown - // apply outcome must surface on FIRST sight — a replay could contradict - // its own effect, and a later different error must not overwrite the - // uncertainty. The caller classifies (commit maps it to UndeterminedError). - if terminal_on_undetermined && e.undetermined_result.is_some() { - return Err(Error::RegionError(Box::new(e))); + if followers_only { + // The one-shot probe was inconclusive. Restore the cached + // leader and resume the normal ServerIsBusy backoff path. + return Self::retry_same_region( + context, + plan, + region, + backoff, + candidate_state, + Error::RegionError(Box::new(error)), + ) + .await; } - match backoff.next_delay_duration() { - Some(duration) => { - let region_error_resolved = - handle_region_error(pd_client.clone(), e, region_store).await?; - // don't sleep if we have resolved the region error - if !region_error_resolved { - sleep(duration).await; - } - Self::single_plan_handler( - pd_client, - plan, - backoff, - permits, - preserve_region_results, - terminal_on_undetermined, - terminal_on_dispatch_error, - ) - .await - } - None => { - warn!( - "giving up after exhausting retries on region error, region: {:?}", - region_ver_id - ); - Err(Error::RegionError(Box::new(e))) + } + + Self::handle_region_response(context, plan, region.ver_id(), region_store, backoff, error) + .await + } + + async fn handle_region_response( + context: RetryContext, + plan: P, + region_ver_id: RegionVerId, + region_store: RegionStore, + mut backoff: Backoff, + error: errorpb::Error, + ) -> Result<::Result> { + debug!( + "single_shard_handler:execute: region error: {:?}, region: {:?}", + error, region_ver_id + ); + // For CAS and commit points, an unknown apply outcome must surface on + // first sight: replaying could contradict its own effect, and a later + // different error must not overwrite the uncertainty. + if context.terminal_on_undetermined && error.undetermined_result.is_some() { + return Err(Error::RegionError(Box::new(error))); + } + + match backoff.next_delay_duration() { + Some(duration) => { + let region_error_resolved = + handle_region_error(context.pd_client.clone(), error, region_store).await?; + if !region_error_resolved { + sleep(duration).await; } + Self::single_plan_handler(context, plan, backoff).await + } + None => { + warn!( + "giving up after exhausting retries on region error, region: {:?}", + region_ver_id + ); + Err(Error::RegionError(Box::new(error))) } - } else { - Ok(vec![Ok(resp)]) } } - #[allow(clippy::too_many_arguments)] - async fn handle_other_error( - pd_client: Arc, + async fn retry_after_routing_error( + context: RetryContext, plan: P, region: RegionVerId, - store: Option, + backoff: Backoff, + error: Error, + ) -> Result<::Result> { + debug!("retry_after_routing_error: {:?}", error); + context.pd_client.invalidate_region_cache(region).await; + Self::retry_after_error(context, plan, backoff, error).await + } + + async fn retry_after_error( + context: RetryContext, + plan: P, mut backoff: Backoff, - permits: Arc, - preserve_region_results: bool, - terminal_on_undetermined: bool, - terminal_on_dispatch_error: bool, - e: Error, + error: Error, ) -> Result<::Result> { - debug!("handle_other_error: {:?}", e); - pd_client.invalidate_region_cache(region).await; - if is_grpc_error(&e) { - if let Some(store_id) = store { - pd_client.invalidate_store_cache(store_id).await; - } - } match backoff.next_delay_duration() { Some(duration) => { sleep(duration).await; - Self::single_plan_handler( - pd_client, - plan, - backoff, - permits, - preserve_region_results, - terminal_on_undetermined, - terminal_on_dispatch_error, - ) - .await + Self::single_plan_handler(context, plan, backoff).await } - None => Err(e), + None => Err(error), } } } @@ -490,10 +952,12 @@ pub(crate) async fn handle_region_error( // paths classify it via `is_undetermined_region_error`. Plans for which a // replay is unsafe never reach this arm — see `terminal_on_undetermined`. Ok(false) - } else if e.server_is_busy.is_some() - || e.raft_entry_too_large.is_some() - || e.max_timestamp_not_synced.is_some() - { + } else if e.server_is_busy.is_some() { + // ServerIsBusy is a definitive rejection, so retrying is safe. The + // cached leader remains valid; the caller applies the normal Region + // backoff and may run the one-shot follower probe for ServerIsBusy(0). + Ok(false) + } else if e.raft_entry_too_large.is_some() || e.max_timestamp_not_synced.is_some() { Err(Error::RegionError(Box::new(e))) } else { debug!( @@ -573,16 +1037,14 @@ where // too many concurrent requests, TiKV is more likely to return a "TiKV // is busy" error let concurrency_permits = Arc::new(Semaphore::new(MULTI_REGION_CONCURRENCY)); - Self::single_plan_handler( - self.pd_client.clone(), - self.inner.clone(), - self.backoff.clone(), - concurrency_permits.clone(), - self.preserve_region_results, - self.terminal_on_undetermined, - self.terminal_on_dispatch_error, - ) - .await + let context = RetryContext { + pd_client: self.pd_client.clone(), + permits: concurrency_permits, + preserve_region_results: self.preserve_region_results, + terminal_on_undetermined: self.terminal_on_undetermined, + terminal_on_dispatch_error: self.terminal_on_dispatch_error, + }; + Self::single_plan_handler(context, self.inner.clone(), self.backoff.clone()).await } } @@ -1108,6 +1570,91 @@ mod test { use crate::mock::MockPdClient; use crate::proto::kvrpcpb::BatchGetResponse; + #[test] + fn grpc_statuses_are_classified_without_cache_churn_for_deadlines() { + #[derive(Debug)] + struct StatusWithSource(tonic::Status); + + impl std::fmt::Display for StatusWithSource { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(formatter) + } + } + + impl std::error::Error for StatusWithSource { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(&self.0) + } + } + + let unavailable = Error::GrpcAPI(tonic::Status::unavailable("down")); + assert!(matches!( + grpc_error_action(&unavailable), + GrpcErrorAction::TryNextPeerAndInvalidate + )); + + let deadline = Error::GrpcAPI(tonic::Status::deadline_exceeded("slow")); + assert!(matches!( + grpc_error_action(&deadline), + GrpcErrorAction::TryNextPeer + )); + + let invalid = Error::GrpcAPI(tonic::Status::invalid_argument("bad request")); + assert!(matches!( + grpc_error_action(&invalid), + GrpcErrorAction::Return + )); + + let explicit_unknown = Error::GrpcAPI(tonic::Status::unknown("stream closed")); + assert!(matches!( + grpc_error_action(&explicit_unknown), + GrpcErrorAction::TryNextPeerAndInvalidate + )); + + let source_less_internal = + Error::GrpcAPI(tonic::Status::internal("Missing response message.")); + assert!(matches!( + grpc_error_action(&source_less_internal), + GrpcErrorAction::TryNextPeerAndInvalidate + )); + + let explicit_resource_limit = Error::GrpcAPI(tonic::Status::resource_exhausted("limit")); + assert!(matches!( + grpc_error_action(&explicit_resource_limit), + GrpcErrorAction::Return + )); + + let transport_status = tonic::Status::from_error(Box::new(StatusWithSource( + tonic::Status::resource_exhausted("http2 overload"), + ))); + assert_eq!(transport_status.code(), Code::ResourceExhausted); + let transport_error = Error::GrpcAPI(transport_status); + assert!(matches!( + grpc_error_action(&transport_error), + GrpcErrorAction::TryNextPeerAndInvalidate + )); + } + + #[test] + fn server_busy_probe_counts_only_zero_wait_for_the_same_leader() { + let mut state = CandidateState::default(); + state.record_leader_busy(1, 0); + state.record_leader_busy(1, 10); + assert!(!state.take_follower_probe()); + + state.record_leader_busy(1, 0); + assert!(state.take_follower_probe()); + assert!(!state.take_follower_probe(), "the probe is one-shot"); + + let mut state = CandidateState::default(); + state.record_leader_busy(1, 0); + state.record_leader_busy(2, 0); + assert!( + !state.take_follower_probe(), + "a new leader must not inherit the old leader's busy count" + ); + } + #[derive(Clone)] struct ErrPlan; diff --git a/src/store/client.rs b/src/store/client.rs index 1c873285..fea9080a 100644 --- a/src/store/client.rs +++ b/src/store/client.rs @@ -35,7 +35,7 @@ impl KvConnect for TikvConnect { async fn connect(&self, address: &str) -> Result { self.security_mgr - .connect(address, move |channel| { + .connect_with_timeout(address, self.timeout, move |channel| { TikvClient::new(channel) .max_decoding_message_size(self.grpc_max_decoding_message_size) .accept_compressed(CompressionEncoding::Gzip)