From 5ac72f2c32139f820a74bc8ca785677e3a3f1e81 Mon Sep 17 00:00:00 2001
From: Ping Yu <yuping@pingcap.com>
Date: Thu, 7 Sep 2023 15:16:02 +0800
Subject: [PATCH] handle grpc error (#419)

Signed-off-by: Ping Yu <yuping@pingcap.com>
---
 src/request/plan.rs | 48 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/request/plan.rs b/src/request/plan.rs
index 7cee192e..1808078a 100644
--- a/src/request/plan.rs
+++ b/src/request/plan.rs
@@ -147,9 +147,27 @@ where
     ) -> Result<<Self as Plan>::Result> {
         // limit concurrent requests
         let permit = permits.acquire().await.unwrap();
-        let mut resp = plan.execute().await?;
+        let res = plan.execute().await;
         drop(permit);
 
+        let is_grpc_error = |e: &Error| matches!(e, Error::GrpcAPI(_) | Error::Grpc(_));
+        let mut resp = match res {
+            Ok(resp) => resp,
+            Err(e) if is_grpc_error(&e) => {
+                return Self::handle_grpc_error(
+                    pd_client,
+                    plan,
+                    region_store,
+                    backoff,
+                    permits,
+                    preserve_region_results,
+                    e,
+                )
+                .await;
+            }
+            Err(e) => return Err(e),
+        };
+
         if let Some(e) = resp.key_errors() {
             Ok(vec![Err(Error::MultipleKeyErrors(e))])
         } else if let Some(e) = resp.region_error() {
@@ -272,6 +290,34 @@ where
         pd_client.invalidate_region_cache(ver_id).await;
         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> {