feat(cluster): allow specify a client timeout on creating flight clie…

…nt (databendlabs#12529)

* add timeout for create_client

* feat: add timeout for rpc client

* chore: refactor rpc client initialization

* chore: add log about ready time for connections

* fix: ut

src/binaries/query/entry.rs

@@ -134,6 +134,7 @@ pub async fn start_services(conf: &InnerConfig) -> Result<()> {
     precheck_services(conf).await?;
 
     let mut shutdown_handle = ShutdownHandle::create()?;
+    let start_time = std::time::Instant::now();
 
     info!("Databend Query start with config: {:?}", conf);
 
@@ -346,7 +347,11 @@ pub async fn start_services(conf: &InnerConfig) -> Result<()> {
         println!("    {}={}", k, v);
     }
 
-    info!("Ready for connections.");
+    info!(
+        "Ready for connections after {}s.",
+        start_time.elapsed().as_secs_f32()
+    );
+
     if conf.background.enable {
         println!("Start background service");
         get_background_service_handler().start().await?;

src/query/config/src/config.rs

@@ -1364,6 +1364,9 @@ pub struct QueryConfig {
     #[clap(long, default_value = "localhost")]
     pub rpc_tls_query_service_domain_name: String,
 
+    #[clap(long, default_value = "0")]
+    pub rpc_client_timeout_secs: u64,
+
     /// Table engine memory enabled
     #[clap(long, parse(try_from_str), default_value = "true")]
     pub table_engine_memory_enabled: bool,
@@ -1548,6 +1551,7 @@ impl TryInto<InnerQueryConfig> for QueryConfig {
             rpc_tls_server_key: self.rpc_tls_server_key,
             rpc_tls_query_server_root_ca_cert: self.rpc_tls_query_server_root_ca_cert,
             rpc_tls_query_service_domain_name: self.rpc_tls_query_service_domain_name,
+            rpc_client_timeout_secs: self.rpc_client_timeout_secs,
             table_engine_memory_enabled: self.table_engine_memory_enabled,
             wait_timeout_mills: self.wait_timeout_mills,
             max_query_log_size: self.max_query_log_size,
@@ -1621,6 +1625,7 @@ impl From<InnerQueryConfig> for QueryConfig {
             rpc_tls_server_key: inner.rpc_tls_server_key,
             rpc_tls_query_server_root_ca_cert: inner.rpc_tls_query_server_root_ca_cert,
             rpc_tls_query_service_domain_name: inner.rpc_tls_query_service_domain_name,
+            rpc_client_timeout_secs: inner.rpc_client_timeout_secs,
             table_engine_memory_enabled: inner.table_engine_memory_enabled,
             wait_timeout_mills: inner.wait_timeout_mills,
             max_query_log_size: inner.max_query_log_size,

src/query/config/src/inner.rs

@@ -179,6 +179,7 @@ pub struct QueryConfig {
     /// Certificate for client to identify query rpc server
     pub rpc_tls_query_server_root_ca_cert: String,
     pub rpc_tls_query_service_domain_name: String,
+    pub rpc_client_timeout_secs: u64,
     /// Table engine memory enabled
     pub table_engine_memory_enabled: bool,
     pub wait_timeout_mills: u64,
@@ -247,6 +248,7 @@ impl Default for QueryConfig {
             rpc_tls_server_key: "".to_string(),
             rpc_tls_query_server_root_ca_cert: "".to_string(),
             rpc_tls_query_service_domain_name: "localhost".to_string(),
+            rpc_client_timeout_secs: 0,
             table_engine_memory_enabled: true,
             wait_timeout_mills: 5000,
             max_query_log_size: 10_000,

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

@@ -19,6 +19,7 @@ use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::time::Duration;
+use std::time::Instant;
 
 use common_arrow::arrow_format::flight::service::flight_service_client::FlightServiceClient;
 use common_base::base::tokio;
@@ -235,10 +236,13 @@ impl ClusterDiscovery {
                 let mut res = Vec::with_capacity(cluster_nodes.len());
                 for node in &cluster_nodes {
                     if node.id != self.local_id {
+                        let start_at = Instant::now();
                         if let Err(cause) = create_client(config, &node.flight_address).await {
                             warn!(
-                                "Cannot connect node [{:?}], remove it in query. cause: {:?}",
-                                node.flight_address, cause
+                                "Cannot connect node [{:?}] after {:?}s, remove it in query. cause: {:?}",
+                                node.flight_address,
+                                start_at.elapsed().as_secs_f32(),
+                                cause
                             );
 
                             continue;
@@ -519,17 +523,19 @@ impl ClusterHeartbeat {
 
 #[async_backtrace::framed]
 pub async fn create_client(config: &InnerConfig, address: &str) -> Result<FlightClient> {
-    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()),
-            )
-            .await?,
-        ))),
-        false => Ok(FlightClient::new(FlightServiceClient::new(
-            ConnectionFactory::create_rpc_channel(address.to_owned(), None, None).await?,
-        ))),
-    }
+    let timeout = if config.query.rpc_client_timeout_secs > 0 {
+        Some(Duration::from_secs(config.query.rpc_client_timeout_secs))
+    } else {
+        None
+    };
+
+    let rpc_tls_config = if config.tls_query_cli_enabled() {
+        Some(config.query.to_rpc_client_tls_config())
+    } else {
+        None
+    };
+
+    Ok(FlightClient::new(FlightServiceClient::new(
+        ConnectionFactory::create_rpc_channel(address.to_owned(), timeout, rpc_tls_config).await?,
+    )))
 }

src/query/service/tests/it/storages/testdata/configs_table_basic.txt

@@ -94,6 +94,7 @@ DB.Table: 'system'.'configs', Table: configs-table_id:1, ver:0, Engine: SystemCo
 | 'query'   | 'openai_api_version'                       | ''                                                             | ''       |
 | 'query'   | 'parquet_fast_read_bytes'                  | 'null'                                                         | ''       |
 | 'query'   | 'quota'                                    | 'null'                                                         | ''       |
+| 'query'   | 'rpc_client_timeout_secs'                  | '0'                                                            | ''       |
 | 'query'   | 'rpc_tls_query_server_root_ca_cert'        | ''                                                             | ''       |
 | 'query'   | 'rpc_tls_query_service_domain_name'        | 'localhost'                                                    | ''       |
 | 'query'   | 'rpc_tls_server_cert'                      | ''                                                             | ''       |