feat(query): support flight keepalive settings

src/common/grpc/src/dns_resolver.rs

@@ -140,19 +140,38 @@ impl Future for DNSServiceFuture {
     }
 }
 
+#[derive(Clone, Copy, Debug, Default)]
+pub struct TcpKeepAliveConfig {
+    pub time: Option<Duration>,
+    pub interval: Option<Duration>,
+    pub retries: Option<u32>,
+}
+
 pub struct ConnectionFactory;
 
 impl ConnectionFactory {
     pub async fn create_rpc_channel(
         addr: impl ToString,
         timeout: Option<Duration>,
         rpc_client_config: Option<RpcClientTlsConfig>,
+        keep_alive: Option<TcpKeepAliveConfig>,
     ) -> std::result::Result<Channel, GrpcConnectionError> {
         let endpoint = Self::create_rpc_endpoint(addr, timeout, rpc_client_config)?;
 
         let mut inner_connector = HttpConnector::new_with_resolver(DNSService);
         inner_connector.set_nodelay(true);
-        inner_connector.set_keepalive(None);
+        match keep_alive {
+            Some(config) => {
+                inner_connector.set_keepalive(config.time);
+                inner_connector.set_keepalive_interval(config.interval);
+                inner_connector.set_keepalive_retries(config.retries);
+            }
+            None => {
+                inner_connector.set_keepalive(None);
+                inner_connector.set_keepalive_interval(None);
+                inner_connector.set_keepalive_retries(None);
+            }
+        }
         inner_connector.enforce_http(false);
         inner_connector.set_connect_timeout(timeout);
 

src/common/grpc/src/lib.rs

@@ -20,6 +20,7 @@ pub use dns_resolver::ConnectionFactory;
 pub use dns_resolver::DNSResolver;
 pub use dns_resolver::DNSService;
 pub use dns_resolver::GrpcConnectionError;
+pub use dns_resolver::TcpKeepAliveConfig;
 pub use grpc_token::GrpcClaim;
 pub use grpc_token::GrpcToken;
 

src/meta/client/src/channel_manager.rs

@@ -142,19 +142,24 @@ impl MetaChannelManager {
     async fn build_channel(&self, addr: &String) -> Result<Channel, MetaNetworkError> {
         info!("MetaChannelManager::build_channel to {}", addr);
 
-        let ch = ConnectionFactory::create_rpc_channel(addr, self.timeout, self.tls_config.clone())
-            .await
-            .map_err(|e| match e {
-                GrpcConnectionError::InvalidUri { .. } => MetaNetworkError::BadAddressFormat(
-                    AnyError::new(&e).add_context(|| "while creating rpc channel"),
-                ),
-                GrpcConnectionError::TLSConfigError { .. } => MetaNetworkError::TLSConfigError(
-                    AnyError::new(&e).add_context(|| "while creating rpc channel"),
-                ),
-                GrpcConnectionError::CannotConnect { .. } => MetaNetworkError::ConnectionError(
-                    ConnectionError::new(e, "while creating rpc channel"),
-                ),
-            })?;
+        let ch = ConnectionFactory::create_rpc_channel(
+            addr,
+            self.timeout,
+            self.tls_config.clone(),
+            None,
+        )
+        .await
+        .map_err(|e| match e {
+            GrpcConnectionError::InvalidUri { .. } => MetaNetworkError::BadAddressFormat(
+                AnyError::new(&e).add_context(|| "while creating rpc channel"),
+            ),
+            GrpcConnectionError::TLSConfigError { .. } => MetaNetworkError::TLSConfigError(
+                AnyError::new(&e).add_context(|| "while creating rpc channel"),
+            ),
+            GrpcConnectionError::CannotConnect { .. } => MetaNetworkError::ConnectionError(
+                ConnectionError::new(e, "while creating rpc channel"),
+            ),
+        })?;
         Ok(ch)
     }
 }

src/meta/client/tests/it/grpc_client.rs

@@ -68,7 +68,7 @@ async fn test_grpc_client_handshake_timeout() {
         // our mock grpc server's handshake impl will sleep 2secs.
         // see: GrpcServiceForTestImp.handshake
         let timeout = Duration::from_secs(1);
-        let c = ConnectionFactory::create_rpc_channel(srv_addr.clone(), Some(timeout), None)
+        let c = ConnectionFactory::create_rpc_channel(srv_addr.clone(), Some(timeout), None, None)
             .await
             .unwrap();
 
@@ -86,7 +86,7 @@ async fn test_grpc_client_handshake_timeout() {
     // handshake success
     {
         let timeout = Duration::from_secs(3);
-        let c = ConnectionFactory::create_rpc_channel(srv_addr, Some(timeout), None)
+        let c = ConnectionFactory::create_rpc_channel(srv_addr, Some(timeout), None, None)
             .await
             .unwrap();
 

src/meta/service/src/meta_node/meta_node.rs

@@ -937,7 +937,7 @@ impl MetaNode {
             addr, timeout
         );
 
-        let chan_res = ConnectionFactory::create_rpc_channel(addr, timeout, None).await;
+        let chan_res = ConnectionFactory::create_rpc_channel(addr, timeout, None, None).await;
         let chan = match chan_res {
             Ok(c) => c,
             Err(e) => {

src/meta/service/tests/it/grpc/metasrv_grpc_handshake.rs

@@ -52,8 +52,9 @@ async fn test_metasrv_handshake() -> anyhow::Result<()> {
 
     let (_tc, addr) = start_metasrv().await?;
 
-    let c = ConnectionFactory::create_rpc_channel(addr, Some(Duration::from_millis(1000)), None)
-        .await?;
+    let c =
+        ConnectionFactory::create_rpc_channel(addr, Some(Duration::from_millis(1000)), None, None)
+            .await?;
     let (mut client, _once) = MetaChannelManager::new_real_client(c);
 
     info!("--- client has smaller ver than S.min_cli_ver");

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

@@ -54,6 +54,7 @@ use databend_common_meta_store::MetaStoreProvider;
 use databend_common_meta_types::NodeInfo;
 use databend_common_meta_types::SeqV;
 use databend_common_metrics::cluster::*;
+use databend_common_settings::FlightKeepAliveParams;
 use databend_common_settings::Settings;
 use databend_common_telemetry::report_node_telemetry;
 use databend_common_version::DATABEND_TELEMETRY_API_KEY;
@@ -74,6 +75,7 @@ use serde::Serialize;
 use tokio::time::sleep;
 
 use crate::servers::flight::FlightClient;
+use crate::servers::flight::keep_alive::build_keep_alive_config;
 
 pub struct ClusterDiscovery {
     local_id: String,
@@ -188,7 +190,9 @@ impl ClusterHelper for Cluster {
                     let mut attempt = 0;
 
                     loop {
-                        let mut conn = create_client(&config, &flight_address).await?;
+                        let mut conn =
+                            create_client(&config, &flight_address, flight_params.keep_alive)
+                                .await?;
                         match conn
                             .do_action::<_, Res>(
                                 path,
@@ -312,7 +316,13 @@ impl ClusterDiscovery {
 
                     if config.query.check_connection_before_schedule && node.id != self.local_id {
                         let start_at = Instant::now();
-                        if let Err(cause) = create_client(config, &node.flight_address).await {
+                        if let Err(cause) = create_client(
+                            config,
+                            &node.flight_address,
+                            FlightKeepAliveParams::default(),
+                        )
+                        .await
+                        {
                             warn!(
                                 "Cannot connect node [{:?}] after {:?}s, remove it in query. cause: {:?}",
                                 node.flight_address,
@@ -964,7 +974,11 @@ impl ClusterHeartbeat {
 }
 
 #[async_backtrace::framed]
-pub async fn create_client(config: &InnerConfig, address: &str) -> Result<FlightClient> {
+pub async fn create_client(
+    config: &InnerConfig,
+    address: &str,
+    keep_alive: FlightKeepAliveParams,
+) -> Result<FlightClient> {
     let timeout = if config.query.rpc_client_timeout_secs > 0 {
         Some(Duration::from_secs(config.query.rpc_client_timeout_secs))
     } else {
@@ -976,9 +990,16 @@ pub async fn create_client(config: &InnerConfig, address: &str) -> Result<Flight
     } else {
         None
     };
+    let keep_alive_config = build_keep_alive_config(keep_alive);
 
     Ok(FlightClient::new(FlightServiceClient::new(
-        ConnectionFactory::create_rpc_channel(address.to_owned(), timeout, rpc_tls_config).await?,
+        ConnectionFactory::create_rpc_channel(
+            address.to_owned(),
+            timeout,
+            rpc_tls_config,
+            keep_alive_config,
+        )
+        .await?,
     )))
 }
 
@@ -987,6 +1008,7 @@ pub struct FlightParams {
     pub(crate) timeout: u64,
     pub(crate) retry_times: u64,
     pub(crate) retry_interval: u64,
+    pub(crate) keep_alive: FlightKeepAliveParams,
 }
 
 #[derive(Clone)]

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

@@ -60,6 +60,7 @@ impl KillInterpreter {
             timeout: settings.get_flight_client_timeout()?,
             retry_times: settings.get_flight_max_retry_times()?,
             retry_interval: settings.get_flight_retry_interval()?,
+            keep_alive: settings.get_flight_keep_alive_params()?,
         };
 
         let mut message = HashMap::with_capacity(warehouse.nodes.len());

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

@@ -66,6 +66,7 @@ impl SetPriorityInterpreter {
             timeout: settings.get_flight_client_timeout()?,
             retry_times: settings.get_flight_max_retry_times()?,
             retry_interval: settings.get_flight_retry_interval()?,
+            keep_alive: settings.get_flight_keep_alive_params()?,
         };
         let res = warehouse
             .do_action::<_, bool>(SET_PRIORITY, message, flight_params)

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

@@ -81,6 +81,7 @@ impl Interpreter for SystemActionInterpreter {
                 timeout: settings.get_flight_client_timeout()?,
                 retry_times: settings.get_flight_max_retry_times()?,
                 retry_interval: settings.get_flight_retry_interval()?,
+                keep_alive: settings.get_flight_keep_alive_params()?,
             };
             warehouse
                 .do_action::<_, ()>(SYSTEM_ACTION, message, flight_params)

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

@@ -100,6 +100,7 @@ impl Interpreter for TruncateTableInterpreter {
                 timeout: settings.get_flight_client_timeout()?,
                 retry_times: settings.get_flight_max_retry_times()?,
                 retry_interval: settings.get_flight_retry_interval()?,
+                keep_alive: settings.get_flight_keep_alive_params()?,
             };
             warehouse
                 .do_action::<_, ()>(TRUNCATE_TABLE, message, flight_params)

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

@@ -16,6 +16,7 @@ use std::collections::HashMap;
 
 use databend_common_config::GlobalConfig;
 use databend_common_exception::Result;
+use databend_common_settings::FlightKeepAliveParams;
 use http::StatusCode;
 use poem::IntoResponse;
 use poem::web::Json;
@@ -65,6 +66,7 @@ async fn get_running_query_dump(query_id: &str) -> Result<HashMap<String, String
         timeout: 60,
         retry_times: 3,
         retry_interval: 3,
+        keep_alive: FlightKeepAliveParams::default(),
     };
     cluster
         .do_action::<_, String>(GET_RUNNING_QUERY_DUMP, message, flight_params)

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

@@ -0,0 +1,28 @@
+// 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 databend_common_grpc::TcpKeepAliveConfig;
+use databend_common_settings::FlightKeepAliveParams;
+
+pub fn build_keep_alive_config(params: FlightKeepAliveParams) -> Option<TcpKeepAliveConfig> {
+    if params.is_disabled() {
+        None
+    } else {
+        Some(TcpKeepAliveConfig {
+            time: params.time,
+            interval: params.interval,
+            retries: params.retries,
+        })
+    }
+}

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

@@ -14,6 +14,7 @@
 
 mod flight_client;
 mod flight_service;
+pub(crate) mod keep_alive;
 mod request_builder;
 pub mod v1;
 

src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs

@@ -39,6 +39,7 @@ use databend_common_grpc::ConnectionFactory;
 use databend_common_pipeline::core::ExecutionInfo;
 use databend_common_pipeline::core::always_callback;
 use databend_common_pipeline::core::basic_callback;
+use databend_common_settings::FlightKeepAliveParams;
 use fastrace::prelude::*;
 use log::warn;
 use parking_lot::Mutex;
@@ -67,6 +68,7 @@ use crate::servers::flight::FlightClient;
 use crate::servers::flight::FlightExchange;
 use crate::servers::flight::FlightReceiver;
 use crate::servers::flight::FlightSender;
+use crate::servers::flight::keep_alive::build_keep_alive_config;
 use crate::servers::flight::v1::actions::INIT_QUERY_FRAGMENTS;
 use crate::servers::flight::v1::actions::START_PREPARED_QUERY;
 use crate::servers::flight::v1::actions::init_query_fragments;
@@ -184,6 +186,11 @@ impl DataExchangeManager {
 
         let config = GlobalConfig::instance();
         let with_cur_rt = env.create_rpc_clint_with_current_rt;
+        let settings = match ctx {
+            Some(ref ctx) => ctx.get_settings(),
+            None => env.settings.clone(),
+        };
+        let keep_alive = settings.get_flight_keep_alive_params()?;
 
         let mut request_exchanges = HashMap::new();
         let mut targets_exchanges = HashMap::<String, Vec<FlightExchange>>::new();
@@ -203,8 +210,10 @@ impl DataExchangeManager {
                     let query_id = env.query_id.clone();
                     let address = source.flight_address.clone();
 
+                    let keep_alive_params = keep_alive;
                     flight_exchanges.push(async move {
-                        let mut flight_client = Self::create_client(&address, with_cur_rt).await?;
+                        let mut flight_client =
+                            Self::create_client(&address, with_cur_rt, keep_alive_params).await?;
 
                         Ok::<QueryExchange, ErrorCode>(match edge {
                             Edge::Fragment(channel) => QueryExchange::Fragment {
@@ -326,21 +335,33 @@ impl DataExchangeManager {
     }
 
     #[async_backtrace::framed]
-    pub async fn create_client(address: &str, use_current_rt: bool) -> Result<FlightClient> {
+    pub async fn create_client(
+        address: &str,
+        use_current_rt: bool,
+        keep_alive: FlightKeepAliveParams,
+    ) -> Result<FlightClient> {
         let config = GlobalConfig::instance();
         let address = address.to_string();
+        let keep_alive_config = build_keep_alive_config(keep_alive);
         let task = async move {
             match config.tls_query_cli_enabled() {
                 true => Ok(FlightClient::new(FlightServiceClient::new(
                     ConnectionFactory::create_rpc_channel(
                         address.to_owned(),
                         None,
                         Some(config.query.to_rpc_client_tls_config()),
+                        keep_alive_config,
                     )
                     .await?,
                 ))),
                 false => Ok(FlightClient::new(FlightServiceClient::new(
-                    ConnectionFactory::create_rpc_channel(address.to_owned(), None, None).await?,
+                    ConnectionFactory::create_rpc_channel(
+                        address.to_owned(),
+                        None,
+                        None,
+                        keep_alive_config,
+                    )
+                    .await?,
                 ))),
             }
         };
@@ -474,6 +495,7 @@ impl DataExchangeManager {
             timeout: settings.get_flight_client_timeout()?,
             retry_times: settings.get_flight_max_retry_times()?,
             retry_interval: settings.get_flight_retry_interval()?,
+            keep_alive: settings.get_flight_keep_alive_params()?,
         };
         let mut root_fragment_ids = actions.get_root_fragment_ids()?;
         let conf = GlobalConfig::instance();