resolved_ts: use smaller timeout when do check_leader (#16000) (#16943)

components/backup-stream/src/endpoint.rs

@@ -1158,9 +1158,14 @@ where
     RT: CdcHandle<EK> + 'static,
     EK: KvEngine,
 {
-    pub async fn resolve(&mut self, regions: Vec<u64>, min_ts: TimeStamp) -> Vec<u64> {
+    pub async fn resolve(
+        &mut self,
+        regions: Vec<u64>,
+        min_ts: TimeStamp,
+        timeout: Option<Duration>,
+    ) -> Vec<u64> {
         match self {
-            BackupStreamResolver::V1(x) => x.resolve(regions, min_ts).await,
+            BackupStreamResolver::V1(x) => x.resolve(regions, min_ts, timeout).await,
             BackupStreamResolver::V2(x, _) => {
                 let x = x.clone();
                 resolve_by_raft(regions, min_ts, x).await

components/backup-stream/src/subscription_manager.rs

@@ -454,7 +454,9 @@ where
                         warn!("waiting for initial scanning done timed out, forcing progress!";
                             "take" => ?now.saturating_elapsed(), "timedout" => %timedout);
                     }
-                    let regions = resolver.resolve(self.subs.current_regions(), min_ts).await;
+                    let regions = resolver
+                        .resolve(self.subs.current_regions(), min_ts, None)
+                        .await;
                     let cps = self.subs.resolve_with(min_ts, regions);
                     let min_region = cps.iter().min_by_key(|rs| rs.checkpoint);
                     // If there isn't any region observed, the `min_ts` can be used as resolved ts

components/cdc/src/endpoint.rs

@@ -1199,7 +1199,7 @@ impl<T: 'static + CdcHandle<E>, E: KvEngine, S: StoreRegionMeta> Endpoint<T, E,
             let regions =
                 if hibernate_regions_compatible && gate.can_enable(FEATURE_RESOLVED_TS_STORE) {
                     CDC_RESOLVED_TS_ADVANCE_METHOD.set(1);
-                    leader_resolver.resolve(regions, min_ts).await
+                    leader_resolver.resolve(regions, min_ts, None).await
                 } else {
                     CDC_RESOLVED_TS_ADVANCE_METHOD.set(0);
                     resolve_by_raft(regions, min_ts, cdc_handle).await

components/resolved_ts/src/advance.rs

@@ -125,7 +125,9 @@ impl AdvanceTsWorker {
                 }
             }
 
-            let regions = leader_resolver.resolve(regions, min_ts).await;
+            let regions = leader_resolver
+                .resolve(regions, min_ts, Some(advance_ts_interval))
+                .await;
             if !regions.is_empty() {
                 if let Err(e) = scheduler.schedule(Task::ResolvedTsAdvanced {
                     regions,
@@ -225,7 +227,12 @@ impl LeadershipResolver {
     // This function broadcasts a special message to all stores, gets the leader id
     // of them to confirm whether current peer has a quorum which accepts its
     // leadership.
-    pub async fn resolve(&mut self, regions: Vec<u64>, min_ts: TimeStamp) -> Vec<u64> {
+    pub async fn resolve(
+        &mut self,
+        regions: Vec<u64>,
+        min_ts: TimeStamp,
+        timeout: Option<Duration>,
+    ) -> Vec<u64> {
         if regions.is_empty() {
             return regions;
         }
@@ -309,6 +316,8 @@ impl LeadershipResolver {
             .find(|req| !req.regions.is_empty())
             .map_or(0, |req| req.regions[0].compute_size());
         let mut check_leader_rpcs = Vec::with_capacity(store_req_map.len());
+        let timeout = get_min_timeout(timeout, DEFAULT_CHECK_LEADER_TIMEOUT_DURATION);
+
         for (store_id, req) in store_req_map {
             if req.regions.is_empty() {
                 continue;
@@ -323,9 +332,16 @@ impl LeadershipResolver {
             let rpc = async move {
                 PENDING_CHECK_LEADER_REQ_COUNT.inc();
                 defer!(PENDING_CHECK_LEADER_REQ_COUNT.dec());
-                let client = get_tikv_client(to_store, pd_client, security_mgr, env, tikv_clients)
-                    .await
-                    .map_err(|e| (to_store, e.retryable(), format!("[get tikv client] {}", e)))?;
+                let client = get_tikv_client(
+                    to_store,
+                    pd_client,
+                    security_mgr,
+                    env,
+                    tikv_clients,
+                    timeout,
+                )
+                .await
+                .map_err(|e| (to_store, e.retryable(), format!("[get tikv client] {}", e)))?;
 
                 // Set min_ts in the request.
                 req.set_ts(min_ts.into_inner());
@@ -356,7 +372,6 @@ impl LeadershipResolver {
 
                 PENDING_CHECK_LEADER_REQ_SENT_COUNT.inc();
                 defer!(PENDING_CHECK_LEADER_REQ_SENT_COUNT.dec());
-                let timeout = DEFAULT_CHECK_LEADER_TIMEOUT_DURATION;
                 let resp = tokio::time::timeout(timeout, rpc)
                     .map_err(|e| (to_store, true, format!("[timeout] {}", e)))
                     .await?
@@ -451,6 +466,11 @@ where
     resps.into_iter().flatten().collect::<Vec<u64>>()
 }
 
+#[inline]
+fn get_min_timeout(timeout: Option<Duration>, default: Duration) -> Duration {
+    timeout.unwrap_or(default).min(default)
+}
+
 fn region_has_quorum(peers: &[Peer], stores: &[u64]) -> bool {
     let mut voters = 0;
     let mut incoming_voters = 0;
@@ -507,14 +527,14 @@ async fn get_tikv_client(
     security_mgr: &SecurityManager,
     env: Arc<Environment>,
     tikv_clients: &Mutex<HashMap<u64, TikvClient>>,
+    timeout: Duration,
 ) -> pd_client::Result<TikvClient> {
     {
         let clients = tikv_clients.lock().await;
         if let Some(client) = clients.get(&store_id).cloned() {
             return Ok(client);
         }
     }
-    let timeout = DEFAULT_CHECK_LEADER_TIMEOUT_DURATION;
     let store = tokio::time::timeout(timeout, pd_client.get_store_async(store_id))
         .await
         .map_err(|e| pd_client::Error::Other(Box::new(e)))
@@ -653,19 +673,45 @@ mod tests {
             .region_read_progress
             .insert(2, Arc::new(progress2));
 
-        leader_resolver.resolve(vec![1, 2], TimeStamp::new(1)).await;
+        leader_resolver
+            .resolve(vec![1, 2], TimeStamp::new(1), None)
+            .await;
         let req = rx.recv_timeout(Duration::from_secs(1)).unwrap();
         assert_eq!(req.regions.len(), 2);
 
         // Checking one region only send 1 region in request.
-        leader_resolver.resolve(vec![1], TimeStamp::new(1)).await;
+        leader_resolver
+            .resolve(vec![1], TimeStamp::new(1), None)
+            .await;
         let req = rx.recv_timeout(Duration::from_secs(1)).unwrap();
         assert_eq!(req.regions.len(), 1);
 
         // Checking zero region does not send request.
-        leader_resolver.resolve(vec![], TimeStamp::new(1)).await;
+        leader_resolver
+            .resolve(vec![], TimeStamp::new(1), None)
+            .await;
         rx.recv_timeout(Duration::from_secs(1)).unwrap_err();
 
         let _ = server.shutdown().await;
     }
+
+    #[test]
+    fn test_get_min_timeout() {
+        assert_eq!(
+            get_min_timeout(None, Duration::from_secs(5)),
+            Duration::from_secs(5)
+        );
+        assert_eq!(
+            get_min_timeout(None, Duration::from_secs(2)),
+            Duration::from_secs(2)
+        );
+        assert_eq!(
+            get_min_timeout(Some(Duration::from_secs(1)), Duration::from_secs(5)),
+            Duration::from_secs(1)
+        );
+        assert_eq!(
+            get_min_timeout(Some(Duration::from_secs(20)), Duration::from_secs(5)),
+            Duration::from_secs(5)
+        );
+    }
 }