adding retryable to scan

src/raw/client.rs

@@ -1,30 +1,35 @@
 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
 
+use async_recursion::async_recursion;
 use core::ops::Range;
 use std::str::FromStr;
 use std::sync::Arc;
 use std::u32;
 
-use futures::StreamExt;
 use log::debug;
+use tokio::sync::Semaphore;
+use tokio::time::sleep;
 
-use crate::backoff::DEFAULT_REGION_BACKOFF;
+use crate::backoff::{DEFAULT_REGION_BACKOFF, DEFAULT_STORE_BACKOFF};
 use crate::common::Error;
 use crate::config::Config;
 use crate::pd::PdClient;
 use crate::pd::PdRpcClient;
+use crate::proto::kvrpcpb::{RawScanRequest, RawScanResponse};
 use crate::proto::metapb;
 use crate::raw::lowering::*;
-use crate::request::Collect;
 use crate::request::CollectSingle;
 use crate::request::EncodeKeyspace;
 use crate::request::KeyMode;
 use crate::request::Keyspace;
 use crate::request::Plan;
 use crate::request::TruncateKeyspace;
+use crate::request::{plan, Collect};
+use crate::store::{HasRegionError, RegionStore};
 use crate::Backoff;
 use crate::BoundRange;
 use crate::ColumnFamily;
+use crate::Error::RegionError;
 use crate::Key;
 use crate::KvPair;
 use crate::Result;
@@ -756,57 +761,42 @@ impl<PdC: PdClient> Client<PdC> {
                 max_limit: MAX_RAW_KV_SCAN_LIMIT,
             });
         }
-
+        let backoff = DEFAULT_STORE_BACKOFF;
-        let mut cur_range = range.into().encode_keyspace(self.keyspace, KeyMode::Raw);
+        let permits = Arc::new(Semaphore::new(16));
+        let range = range.into().encode_keyspace(self.keyspace, KeyMode::Raw);
         let mut result = Vec::new();
-        let mut scan_regions = self.rpc.clone().stores_for_range(cur_range.clone()).boxed();
+        let mut current_limit = limit;
-        let mut region_store =
+        let (start_key, end_key) = range.clone().into_keys();
-            scan_regions
+        let mut current_key: Option<Key> = Some(start_key);
-                .next()
+        while current_limit > 0 {
-                .await
+            let scan_args = ScanInnerArgs {
-                .ok_or(Error::RegionForRangeNotFound {
+                start_key: current_key.clone(),
-                    range: (cur_range.clone()),
+                range: range.clone(),
-                })??;
+                limit,
-        let mut cur_limit = limit;
-
-        while cur_limit > 0 {
-            let request = new_raw_scan_request(
-                cur_range.clone(),
-                cur_limit,
                 key_only,
                 reverse,
-                self.cf.clone(),
+                permits: permits.clone(),
-            );
+                backoff: backoff.clone(),
-            let resp = crate::request::PlanBuilder::new(self.rpc.clone(), self.keyspace, request)
-                .single_region_with_store(region_store.clone())
-                .await?
-                .plan()
-                .execute()
-                .await?;
-            let mut region_scan_res = resp
-                .kvs
-                .into_iter()
-                .map(Into::into)
-                .collect::<Vec<KvPair>>();
-            let res_len = region_scan_res.len();
-            result.append(&mut region_scan_res);
-
-            // if the number of results is less than cur_limit, it means this scan range contains more than one region, so we need to scan next region
-            if res_len < cur_limit as usize {
-                region_store = match scan_regions.next().await {
-                    Some(Ok(rs)) => {
-                        cur_range = BoundRange::new(
-                            std::ops::Bound::Included(region_store.region_with_leader.range().1),
-                            cur_range.to,
-                        );
-                        rs
-                    }
-                    Some(Err(e)) => return Err(e),
-                    None => break,
             };
-                cur_limit -= res_len as u32;
+            let (res, next_key) = self.retryable_scan(scan_args).await?;
-            } else {
+
+            let mut kvs = res
+                .map(|r| r.kvs.into_iter().map(Into::into).collect::<Vec<KvPair>>())
+                .unwrap_or(Vec::new());
+
+            if !kvs.is_empty() {
+                current_limit -= kvs.len() as u32;
+                result.append(&mut kvs);
+            }
+            if end_key
+                .as_ref()
+                .map(|ek| ek <= next_key.as_ref() && !ek.is_empty())
+                .unwrap_or(false)
+                || next_key.is_empty()
+            {
                 break;
+            } else {
+                current_key = Some(next_key);
             }
         }
 
@@ -819,6 +809,61 @@ impl<PdC: PdClient> Client<PdC> {
         Ok(result)
     }
 
+    #[async_recursion]
+    async fn retryable_scan(
+        &self,
+        mut scan_args: ScanInnerArgs,
+    ) -> Result<(Option<RawScanResponse>, Key)> {
+        let start_key = match scan_args.start_key {
+            None => return Ok((None, Key::EMPTY)),
+            Some(ref sk) => sk,
+        };
+        let permit = scan_args.permits.acquire().await.unwrap();
+        let region = self.rpc.clone().region_for_key(start_key).await?;
+        let store = self.rpc.clone().store_for_id(region.id()).await?;
+        let request = new_raw_scan_request(
+            scan_args.range.clone(),
+            scan_args.limit,
+            scan_args.key_only,
+            scan_args.reverse,
+            self.cf.clone(),
+        );
+        let resp = self.do_store_scan(store.clone(), request).await;
+        drop(permit);
+        match resp {
+            Ok(mut r) => {
+                if let Some(err) = r.region_error() {
+                    let status =
+                        plan::handle_region_error(self.rpc.clone(), err.clone(), store.clone())
+                            .await?;
+                    return if status {
+                        self.retryable_scan(scan_args.clone()).await
+                    } else if let Some(duration) = scan_args.backoff.next_delay_duration() {
+                        sleep(duration).await;
+                        self.retryable_scan(scan_args.clone()).await
+                    } else {
+                        Err(RegionError(Box::new(err)))
+                    };
+                }
+                Ok((Some(r), region.end_key()))
+            }
+            Err(err) => Err(err),
+        }
+    }
+
+    async fn do_store_scan(
+        &self,
+        store: RegionStore,
+        scan_request: RawScanRequest,
+    ) -> Result<RawScanResponse> {
+        crate::request::PlanBuilder::new(self.rpc.clone(), self.keyspace, scan_request)
+            .single_region_with_store(store.clone())
+            .await?
+            .plan()
+            .execute()
+            .await
+    }
+
     async fn batch_scan_inner(
         &self,
         ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
@@ -865,6 +910,17 @@ impl<PdC: PdClient> Client<PdC> {
     }
 }
 
+#[derive(Clone)]
+struct ScanInnerArgs {
+    start_key: Option<Key>,
+    range: BoundRange,
+    limit: u32,
+    key_only: bool,
+    reverse: bool,
+    permits: Arc<Semaphore>,
+    backoff: Backoff,
+}
+
 #[cfg(test)]
 mod tests {
     use std::any::Any;

src/request/plan.rs

@@ -187,7 +187,7 @@ where
             match backoff.next_delay_duration() {
                 Some(duration) => {
                     let region_error_resolved =
-                        Self::handle_region_error(pd_client.clone(), e, region_store).await?;
+                        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;
@@ -208,11 +208,40 @@ where
         }
     }
 
+    async fn handle_grpc_error(
+        pd_client: Arc<PdC>,
+        plan: P,
+        region_store: RegionStore,
+        mut backoff: Backoff,
+        permits: Arc<Semaphore>,
+        preserve_region_results: bool,
+        e: Error,
+    ) -> Result<<Self as Plan>::Result> {
+        debug!("handle grpc error: {:?}", e);
+        let ver_id = region_store.region_with_leader.ver_id();
+        pd_client.invalidate_region_cache(ver_id).await;
+        match backoff.next_delay_duration() {
+            Some(duration) => {
+                sleep(duration).await;
+                Self::single_plan_handler(
+                    pd_client,
+                    plan,
+                    backoff,
+                    permits,
+                    preserve_region_results,
+                )
+                .await
+            }
+            None => Err(e),
+        }
+    }
+}
+
 // Returns
 // 1. Ok(true): error has been resolved, retry immediately
 // 2. Ok(false): backoff, and then retry
 // 3. Err(Error): can't be resolved, return the error to upper level
-    async fn handle_region_error(
+pub(crate) async fn handle_region_error<PdC: PdClient>(
     pd_client: Arc<PdC>,
     e: errorpb::Error,
     region_store: RegionStore,
@@ -242,12 +271,7 @@ where
         pd_client.invalidate_region_cache(ver_id).await;
         Ok(false)
     } else if e.epoch_not_match.is_some() {
-            Self::on_region_epoch_not_match(
+        on_region_epoch_not_match(pd_client.clone(), region_store, e.epoch_not_match.unwrap()).await
-                pd_client.clone(),
-                region_store,
-                e.epoch_not_match.unwrap(),
-            )
-            .await
     } else if e.stale_command.is_some() || e.region_not_found.is_some() {
         pd_client.invalidate_region_cache(ver_id).await;
         Ok(false)
@@ -268,7 +292,7 @@ where
 // 1. Ok(true): error has been resolved, retry immediately
 // 2. Ok(false): backoff, and then retry
 // 3. Err(Error): can't be resolved, return the error to upper level
-    async fn on_region_epoch_not_match(
+pub(crate) async fn on_region_epoch_not_match<PdC: PdClient>(
     pd_client: Arc<PdC>,
     region_store: RegionStore,
     error: EpochNotMatch,
@@ -304,35 +328,6 @@ where
     Ok(false)
 }
 
-    async fn handle_grpc_error(
-        pd_client: Arc<PdC>,
-        plan: P,
-        region_store: RegionStore,
-        mut backoff: Backoff,
-        permits: Arc<Semaphore>,
-        preserve_region_results: bool,
-        e: Error,
-    ) -> Result<<Self as Plan>::Result> {
-        debug!("handle grpc error: {:?}", e);
-        let ver_id = region_store.region_with_leader.ver_id();
-        pd_client.invalidate_region_cache(ver_id).await;
-        match backoff.next_delay_duration() {
-            Some(duration) => {
-                sleep(duration).await;
-                Self::single_plan_handler(
-                    pd_client,
-                    plan,
-                    backoff,
-                    permits,
-                    preserve_region_results,
-                )
-                .await
-            }
-            None => Err(e),
-        }
-    }
-}
-
 impl<P: Plan, PdC: PdClient> Clone for RetryableMultiRegion<P, PdC> {
     fn clone(&self) -> Self {
         RetryableMultiRegion {