retry with grpc errors

Author: StrikeW

src/meta/service/src/ddl_service.rs

@@ -966,15 +966,16 @@ impl DdlService for DdlServiceImpl {
             );
             for table in tables {
                 // send a request to the frontend to get the ReplaceTablePlan
+                // will retry with exponential backoff if the request fails
                 let resp = client
                     .get_table_replace_plan(GetTableReplacePlanRequest {
                         database_id: table.database_id,
                         owner: table.owner,
                         table_name: table.name,
                         table_change: Some(table_change.clone()),
                     })
-                    .await
-                    .map_err(MetaError::from)?;
+                    .await?
+                    .into_inner();
 
                 if let Some(plan) = resp.replace_plan {
                     plan.table

src/rpc_client/src/frontend_client.rs

@@ -21,21 +21,20 @@ use risingwave_common::monitor::{EndpointExt, TcpConfig};
 use risingwave_common::util::addr::HostAddr;
 use risingwave_pb::frontend_service::frontend_service_client::FrontendServiceClient;
 use risingwave_pb::frontend_service::{GetTableReplacePlanRequest, GetTableReplacePlanResponse};
+use tokio_retry::strategy::{jitter, ExponentialBackoff};
 use tonic::transport::Endpoint;
+use tonic::Response;
 
 use crate::error::Result;
 use crate::tracing::{Channel, TracingInjectedChannelExt};
-use crate::{frontend_rpc_client_method_impl, RpcClient, RpcClientPool};
+use crate::{RpcClient, RpcClientPool};
 
-#[derive(Clone)]
-pub struct FrontendClient(FrontendServiceClient<Channel>);
+const DEFAULT_RETRY_INTERVAL: u64 = 50;
+const DEFAULT_RETRY_MAX_DELAY: Duration = Duration::from_secs(5);
+const DEFAULT_RETRY_MAX_ATTEMPTS: usize = 10;
 
-#[async_trait]
-impl RpcClient for FrontendClient {
-    async fn new_client(host_addr: HostAddr) -> Result<Self> {
-        Self::new(host_addr).await
-    }
-}
+#[derive(Clone)]
+struct FrontendClient(FrontendServiceClient<Channel>);
 
 impl FrontendClient {
     async fn new(host_addr: HostAddr) -> Result<Self> {
@@ -59,17 +58,61 @@ impl FrontendClient {
 }
 
 // similar to the stream_client used in the Meta node
-pub type FrontendClientPool = RpcClientPool<FrontendClient>;
+pub type FrontendClientPool = RpcClientPool<FrontendRetryClient>;
 pub type FrontendClientPoolRef = Arc<FrontendClientPool>;
 
-macro_rules! for_all_frontend_rpc {
-    ($macro:ident) => {
-        $macro! {
-             { 0, get_table_replace_plan, GetTableReplacePlanRequest, GetTableReplacePlanResponse }
-        }
-    };
+#[async_trait]
+impl RpcClient for FrontendRetryClient {
+    async fn new_client(host_addr: HostAddr) -> Result<Self> {
+        Self::new(host_addr).await
+    }
 }
 
-impl FrontendClient {
-    for_all_frontend_rpc! { frontend_rpc_client_method_impl }
+#[derive(Clone)]
+pub struct FrontendRetryClient {
+    client: FrontendClient,
+}
+
+impl FrontendRetryClient {
+    async fn new(host_addr: HostAddr) -> Result<Self> {
+        let client = FrontendClient::new(host_addr).await?;
+        Ok(Self { client })
+    }
+
+    #[inline(always)]
+    fn get_retry_strategy() -> impl Iterator<Item = Duration> {
+        ExponentialBackoff::from_millis(DEFAULT_RETRY_INTERVAL)
+            .max_delay(DEFAULT_RETRY_MAX_DELAY)
+            .take(DEFAULT_RETRY_MAX_ATTEMPTS)
+            .map(jitter)
+    }
+
+    fn should_retry(status: &tonic::Status) -> bool {
+        if status.code() == tonic::Code::Unavailable
+            || status.code() == tonic::Code::Unknown
+            || status.code() == tonic::Code::Unauthenticated
+            || status.code() == tonic::Code::Aborted
+        {
+            return true;
+        }
+        false
+    }
+
+    pub async fn get_table_replace_plan(
+        &self,
+        request: GetTableReplacePlanRequest,
+    ) -> std::result::Result<Response<GetTableReplacePlanResponse>, tonic::Status> {
+        tokio_retry::RetryIf::spawn(
+            Self::get_retry_strategy(),
+            || async {
+                self.client
+                    .to_owned()
+                    .0
+                    .get_table_replace_plan(request.clone())
+                    .await
+            },
+            Self::should_retry,
+        )
+        .await
+    }
 }

src/rpc_client/src/lib.rs

@@ -66,7 +66,7 @@ mod tracing;
 pub use compactor_client::{CompactorClient, GrpcCompactorProxyClient};
 pub use compute_client::{ComputeClient, ComputeClientPool, ComputeClientPoolRef};
 pub use connector_client::{ConnectorClient, SinkCoordinatorStreamHandle, SinkWriterStreamHandle};
-pub use frontend_client::{FrontendClient, FrontendClientPool, FrontendClientPoolRef};
+pub use frontend_client::{FrontendClientPool, FrontendClientPoolRef};
 pub use hummock_meta_client::{CompactionEventItem, HummockMetaClient};
 pub use meta_client::{MetaClient, SinkCoordinationRpcClient};
 use rw_futures_util::await_future_with_monitor_error_stream;
@@ -210,23 +210,6 @@ macro_rules! meta_rpc_client_method_impl {
     }
 }
 
-#[macro_export]
-macro_rules! frontend_rpc_client_method_impl {
-    ($( { $client:tt, $fn_name:ident, $req:ty, $resp:ty }),*) => {
-        $(
-            pub async fn $fn_name(&self, request: $req) -> $crate::Result<$resp> {
-                Ok(self
-                    .$client
-                    .to_owned()
-                    .$fn_name(request)
-                    .await
-                    .map_err($crate::error::RpcError::from_frontend_status)?
-                    .into_inner())
-            }
-        )*
-    }
-}
-
 pub const DEFAULT_BUFFER_SIZE: usize = 16;
 
 pub struct BidiStreamSender<REQ> {