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