refactor: add retry for flight service

Cargo.toml

@@ -244,6 +244,7 @@ geos = { version = "8.3", features = ["static", "geo", "geo-types"] }
 geozero = { version = "0.13.0", features = ["default", "with-wkb", "with-geos", "with-geojson"] }
 hashbrown = { version = "0.14.3", default-features = false }
 http = "1"
+hyper = "0.14.20"
 iceberg = { version = "0.3.0", git = "https://github.com/apache/iceberg-rust/", rev = "c3549836796f93aa3ad22276af788aa3d92533a1" }
 iceberg-catalog-hms = { version = "0.3.0", git = "https://github.com/apache/iceberg-rust/", rev = "c3549836796f93aa3ad22276af788aa3d92533a1" }
 iceberg-catalog-rest = { version = "0.3.0", git = "https://github.com/apache/iceberg-rust/", rev = "c3549836796f93aa3ad22276af788aa3d92533a1" }

scripts/ci/ci-run-stateful-tests-cluster-minio.sh

@@ -20,6 +20,8 @@ export STORAGE_ALLOW_INSECURE=true
 
 echo "Install dependence"
 python3 -m pip install --quiet mysql-connector-python requests
+sudo apt-get update -yq
+sudo apt-get install -yq dsniff net-tools
 
 echo "calling test suite"
 echo "Starting Cluster databend-query"

scripts/ci/ci-run-stateful-tests-standalone-minio.sh

@@ -20,6 +20,8 @@ export STORAGE_ALLOW_INSECURE=true
 
 echo "Install dependence"
 python3 -m pip install --quiet mysql-connector-python requests
+sudo apt-get update -yq
+sudo apt-get install -yq dsniff net-tools
 
 echo "calling test suite"
 echo "Starting standalone DatabendQuery(debug)"

src/common/exception/Cargo.toml

@@ -21,6 +21,7 @@ bincode = { workspace = true }
 geos = { workspace = true }
 geozero = { workspace = true }
 http = { workspace = true }
+hyper = { workspace = true }
 opendal = { workspace = true }
 parquet = { workspace = true }
 paste = { workspace = true }

src/common/exception/src/exception_into.rs

@@ -375,6 +375,13 @@ impl From<tonic::Status> for ErrorCode {
             tonic::Code::Unknown => {
                 let details = status.details();
                 if details.is_empty() {
+                    if status.source().map_or(false, |e| e.is::<hyper::Error>()) {
+                        return ErrorCode::CannotConnectNode(format!(
+                            "{}, source: {:?}",
+                            status.message(),
+                            status.source()
+                        ));
+                    }
                     return ErrorCode::UnknownException(format!(
                         "{}, source: {:?}",
                         status.message(),

src/query/service/src/clusters/cluster.rs

@@ -50,11 +50,13 @@ use futures::future::Either;
 use futures::Future;
 use futures::StreamExt;
 use log::error;
+use log::info;
 use log::warn;
 use rand::thread_rng;
 use rand::Rng;
 use serde::Deserialize;
 use serde::Serialize;
+use tokio::time::sleep;
 
 use crate::servers::flight::FlightClient;
 
@@ -79,11 +81,11 @@ pub trait ClusterHelper {
 
     fn get_nodes(&self) -> Vec<Arc<NodeInfo>>;
 
-    async fn do_action<T: Serialize + Send, Res: for<'de> Deserialize<'de> + Send>(
+    async fn do_action<T: Serialize + Send + Clone, Res: for<'de> Deserialize<'de> + Send>(
         &self,
         path: &str,
         message: HashMap<String, T>,
-        timeout: u64,
+        flight_params: FlightParams,
     ) -> Result<HashMap<String, Res>>;
 }
 
@@ -116,11 +118,11 @@ impl ClusterHelper for Cluster {
         self.nodes.to_vec()
     }
 
-    async fn do_action<T: Serialize + Send, Res: for<'de> Deserialize<'de> + Send>(
+    async fn do_action<T: Serialize + Send + Clone, Res: for<'de> Deserialize<'de> + Send>(
         &self,
         path: &str,
         message: HashMap<String, T>,
-        timeout: u64,
+        flight_params: FlightParams,
     ) -> Result<HashMap<String, Res>> {
         fn get_node<'a>(nodes: &'a [Arc<NodeInfo>], id: &str) -> Result<&'a Arc<NodeInfo>> {
             for node in nodes {
@@ -145,12 +147,32 @@ impl ClusterHelper for Cluster {
                 let node_secret = node.secret.clone();
 
                 async move {
-                    let mut conn = create_client(&config, &flight_address).await?;
-                    Ok::<_, ErrorCode>((
-                        id,
-                        conn.do_action::<_, Res>(path, node_secret, message, timeout)
-                            .await?,
-                    ))
+                    let mut attempt = 0;
+
+                    loop {
+                        let mut conn = create_client(&config, &flight_address).await?;
+                        match conn
+                            .do_action::<_, Res>(
+                                path,
+                                node_secret.clone(),
+                                message.clone(),
+                                flight_params.timeout,
+                            )
+                            .await
+                        {
+                            Ok(result) => return Ok((id, result)),
+                            Err(e)
+                                if e.code() == ErrorCode::CANNOT_CONNECT_NODE
+                                    && attempt < flight_params.retry_times =>
+                            {
+                                // only retry when error is network problem
+                                info!("retry do_action, attempt: {}", attempt);
+                                attempt += 1;
+                                sleep(Duration::from_secs(flight_params.retry_interval)).await;
+                            }
+                            Err(e) => return Err(e),
+                        }
+                    }
                 }
             });
         }
@@ -505,3 +527,10 @@ pub async fn create_client(config: &InnerConfig, address: &str) -> Result<Flight
         ConnectionFactory::create_rpc_channel(address.to_owned(), timeout, rpc_tls_config).await?,
     )))
 }
+
+#[derive(Clone, Copy, Debug)]
+pub struct FlightParams {
+    pub(crate) timeout: u64,
+    pub(crate) retry_times: u64,
+    pub(crate) retry_interval: u64,
+}

src/query/service/src/clusters/mod.rs

@@ -17,3 +17,4 @@ mod cluster;
 pub use cluster::Cluster;
 pub use cluster::ClusterDiscovery;
 pub use cluster::ClusterHelper;
+pub use cluster::FlightParams;

src/query/service/src/interpreters/interpreter_kill.rs

@@ -21,6 +21,7 @@ use databend_common_exception::Result;
 use databend_common_sql::plans::KillPlan;
 
 use crate::clusters::ClusterHelper;
+use crate::clusters::FlightParams;
 use crate::interpreters::Interpreter;
 use crate::pipelines::PipelineBuildResult;
 use crate::servers::flight::v1::actions::KILL_QUERY;
@@ -54,7 +55,6 @@ impl KillInterpreter {
     async fn kill_cluster_query(&self) -> Result<PipelineBuildResult> {
         let cluster = self.ctx.get_cluster();
         let settings = self.ctx.get_settings();
-        let timeout = settings.get_flight_client_timeout()?;
 
         let mut message = HashMap::with_capacity(cluster.nodes.len());
 
@@ -63,9 +63,14 @@ impl KillInterpreter {
                 message.insert(node_info.id.clone(), self.plan.clone());
             }
         }
+        let flight_params = FlightParams {
+            timeout: settings.get_flight_client_timeout()?,
+            retry_times: settings.get_max_flight_retry_times()?,
+            retry_interval: settings.get_flight_retry_interval()?,
+        };
 
         let res = cluster
-            .do_action::<_, bool>(KILL_QUERY, message, timeout)
+            .do_action::<_, bool>(KILL_QUERY, message, flight_params)
             .await?;
 
         match res.values().any(|x| *x) {

src/query/service/src/interpreters/interpreter_set_priority.rs

@@ -21,6 +21,7 @@ use databend_common_exception::Result;
 use databend_common_sql::plans::SetPriorityPlan;
 
 use crate::clusters::ClusterHelper;
+use crate::clusters::FlightParams;
 use crate::interpreters::Interpreter;
 use crate::pipelines::PipelineBuildResult;
 use crate::servers::flight::v1::actions::SET_PRIORITY;
@@ -61,9 +62,13 @@ impl SetPriorityInterpreter {
         }
 
         let settings = self.ctx.get_settings();
-        let timeout = settings.get_flight_client_timeout()?;
+        let flight_params = FlightParams {
+            timeout: settings.get_flight_client_timeout()?,
+            retry_times: settings.get_max_flight_retry_times()?,
+            retry_interval: settings.get_flight_retry_interval()?,
+        };
         let res = cluster
-            .do_action::<_, bool>(SET_PRIORITY, message, timeout)
+            .do_action::<_, bool>(SET_PRIORITY, message, flight_params)
             .await?;
 
         match res.values().any(|x| *x) {

src/query/service/src/interpreters/interpreter_system_action.rs

@@ -22,6 +22,7 @@ use databend_common_sql::plans::SystemAction;
 use databend_common_sql::plans::SystemPlan;
 
 use crate::clusters::ClusterHelper;
+use crate::clusters::FlightParams;
 use crate::interpreters::Interpreter;
 use crate::pipelines::PipelineBuildResult;
 use crate::servers::flight::v1::actions::SYSTEM_ACTION;
@@ -74,9 +75,13 @@ impl Interpreter for SystemActionInterpreter {
             }
 
             let settings = self.ctx.get_settings();
-            let timeout = settings.get_flight_client_timeout()?;
+            let flight_params = FlightParams {
+                timeout: settings.get_flight_client_timeout()?,
+                retry_times: settings.get_max_flight_retry_times()?,
+                retry_interval: settings.get_flight_retry_interval()?,
+            };
             cluster
-                .do_action::<_, ()>(SYSTEM_ACTION, message, timeout)
+                .do_action::<_, ()>(SYSTEM_ACTION, message, flight_params)
                 .await?;
         }
 

src/query/service/src/interpreters/interpreter_table_truncate.rs

@@ -21,6 +21,7 @@ use databend_common_exception::Result;
 use databend_common_sql::plans::TruncateTablePlan;
 
 use crate::clusters::ClusterHelper;
+use crate::clusters::FlightParams;
 use crate::interpreters::Interpreter;
 use crate::pipelines::PipelineBuildResult;
 use crate::servers::flight::v1::actions::TRUNCATE_TABLE;
@@ -95,9 +96,13 @@ impl Interpreter for TruncateTableInterpreter {
             }
 
             let settings = self.ctx.get_settings();
-            let timeout = settings.get_flight_client_timeout()?;
+            let flight_params = FlightParams {
+                timeout: settings.get_flight_client_timeout()?,
+                retry_times: settings.get_max_flight_retry_times()?,
+                retry_interval: settings.get_flight_retry_interval()?,
+            };
             cluster
-                .do_action::<_, ()>(TRUNCATE_TABLE, message, timeout)
+                .do_action::<_, ()>(TRUNCATE_TABLE, message, flight_params)
                 .await?;
         }
 

src/query/service/src/pipelines/processors/transforms/aggregator/serde/transform_deserializer.rs

@@ -225,6 +225,7 @@ where
             DataPacket::MutationStatus { .. } => unreachable!(),
             DataPacket::DataCacheMetrics(_) => unreachable!(),
             DataPacket::FragmentData(v) => Ok(vec![self.recv_data(meta.packet, v)?]),
+            DataPacket::FlightControl(_) => unreachable!(),
         }
     }
 }

src/query/service/src/servers/admin/v1/query_profiling.rs

@@ -30,6 +30,7 @@ use poem::IntoResponse;
 
 use crate::clusters::ClusterDiscovery;
 use crate::clusters::ClusterHelper;
+use crate::clusters::FlightParams;
 use crate::servers::flight::v1::actions::GET_PROFILE;
 use crate::sessions::SessionManager;
 
@@ -103,9 +104,13 @@ async fn get_cluster_profile(query_id: &str) -> Result<Vec<PlanProfile>, ErrorCo
             message.insert(node_info.id.clone(), query_id.to_owned());
         }
     }
-
+    let flight_params = FlightParams {
+        timeout: 60,
+        retry_times: 3,
+        retry_interval: 3,
+    };
     let res = cluster
-        .do_action::<_, Option<Vec<PlanProfile>>>(GET_PROFILE, message, 60)
+        .do_action::<_, Option<Vec<PlanProfile>>>(GET_PROFILE, message, flight_params)
         .await?;
 
     match res.into_values().find(Option::is_some) {

src/query/service/src/servers/flight/codec.rs

@@ -0,0 +1,72 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+use prost::Message;
+use tonic::codec::Codec;
+use tonic::codec::DecodeBuf;
+use tonic::codec::Decoder;
+use tonic::codec::EncodeBuf;
+use tonic::codec::Encoder;
+use tonic::Status;
+
+#[derive(Default)]
+pub struct MessageCodec<E, D>(PhantomData<(E, D)>);
+
+impl<E: Message + 'static, D: Message + Default + 'static> Codec for MessageCodec<E, D> {
+    type Encode = Arc<E>;
+    type Decode = D;
+    type Encoder = ArcEncoder<E>;
+    type Decoder = DefaultDecoder<D>;
+
+    fn encoder(&mut self) -> Self::Encoder {
+        ArcEncoder(PhantomData)
+    }
+
+    fn decoder(&mut self) -> Self::Decoder {
+        DefaultDecoder(PhantomData)
+    }
+}
+
+pub struct ArcEncoder<E>(PhantomData<E>);
+
+impl<T: Message> Encoder for ArcEncoder<T> {
+    type Item = Arc<T>;
+
+    type Error = Status;
+
+    fn encode(&mut self, item: Self::Item, dst: &mut EncodeBuf<'_>) -> Result<(), Self::Error> {
+        item.as_ref()
+            .encode(dst)
+            .map_err(|e| Status::internal(e.to_string()))
+    }
+}
+
+pub struct DefaultDecoder<D>(PhantomData<D>);
+
+impl<T: Message + Default> Decoder for DefaultDecoder<T> {
+    type Item = T;
+
+    type Error = Status;
+
+    fn decode(&mut self, buf: &mut DecodeBuf<'_>) -> Result<Option<Self::Item>, Self::Error> {
+        let item = Message::decode(buf)
+            .map(Some)
+            .map_err(|e| Status::internal(e.to_string()))?;
+
+        Ok(item)
+    }
+}