client: fall back to followers when leader is unreachable

src/common/security.rs

@@ -82,6 +82,31 @@ impl SecurityManager {
         addr: &str,
         factory: Factory,
     ) -> Result<Client>
+    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<Factory, Client>(
+        &self,
+        addr: &str,
+        timeout: Duration,
+        factory: Factory,
+    ) -> Result<Client>
+    where
+        Factory: FnOnce(Channel) -> Client,
+    {
+        self.connect_inner(addr, Some(timeout), factory).await
+    }
+
+    async fn connect_inner<Factory, Client>(
+        &self,
+        addr: &str,
+        timeout: Option<Duration>,
+        factory: Factory,
+    ) -> Result<Client>
     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))

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<Arc<dyn Fn(RegionWithLeader) -> Result<RegionStore> + 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<Arc<dyn Fn(&Key) -> Result<RegionWithLeader> + Send + Sync + 'static>>,
+    /// Optional observer/override for leader cache updates.
+    #[new(default)]
+    update_leader_hook:
+        Option<Arc<dyn Fn(RegionVerId, metapb::Peer) -> Result<()> + Send + Sync + 'static>>,
+    /// Optional observer for region cache invalidations.
+    #[new(default)]
+    invalidate_region_hook: Option<Arc<dyn Fn(RegionVerId) + Send + Sync + 'static>>,
 }
 
 #[async_trait]
@@ -110,6 +122,14 @@ impl MockPdClient {
         MockPdClient::new(MockKvClient::default())
     }
 
+    pub fn with_map_region_to_store_hook<F>(mut self, hook: F) -> MockPdClient
+    where
+        F: Fn(RegionWithLeader) -> Result<RegionStore> + 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<F>(mut self, hook: F) -> MockPdClient
@@ -120,6 +140,22 @@ impl MockPdClient {
         self
     }
 
+    pub fn with_update_leader_hook<F>(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<F>(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<Self>, region: RegionWithLeader) -> Result<RegionStore> {
+        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) {}