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"
 itertools = "0.10.5"
 jsonb = "0.4.1"
 jwt-simple = "0.11.0"

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,7 +81,7 @@ 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>,
@@ -116,7 +118,7 @@ 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>,
@@ -145,12 +147,33 @@ 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;
+                    let max_attempts = 2;
+
+                    loop {
+                        let mut conn = create_client(&config, &flight_address).await?;
+                        match conn
+                            .do_action::<_, Res>(
+                                path,
+                                node_secret.clone(),
+                                message.clone(),
+                                timeout,
+                            )
+                            .await
+                        {
+                            Ok(result) => return Ok((id, result)),
+                            Err(e)
+                                if e.code() == ErrorCode::CANNOT_CONNECT_NODE
+                                    && attempt < max_attempts =>
+                            {
+                                // only retry when error is network problem
+                                info!("retry do_action, attempt: {}", attempt);
+                                attempt += 1;
+                                sleep(Duration::from_secs(1)).await;
+                            }
+                            Err(e) => return Err(e),
+                        }
+                    }
                 }
             });
         }

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/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)
+    }
+}