chore(query): support custom query graceful shutdown timeout (#14712)

* chore(query): support custom query graceful shutdown timeout * chore(query): support custom query graceful shutdown timeout * chore(query): support custom query graceful shutdown timeout * chore(query): support custom query graceful shutdown timeout
datafuselabs · Feb 23, 2024 · 5c46a27 · 5c46a27
1 parent 11ee029
commit 5c46a27
Show file tree

Hide file tree

Showing 18 changed files with 56 additions and 41 deletions.
diff --git a/docker/query-config.toml b/docker/query-config.toml
@@ -1,6 +1,6 @@
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 flight_api_address = "0.0.0.0:9090"
 admin_api_address = "0.0.0.0:8080"

diff --git a/scripts/ci/deploy/config/databend-query-management-mode.toml b/scripts/ci/deploy/config/databend-query-management-mode.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # Databend Query http address.
 # For admin RESET API.

diff --git a/scripts/ci/deploy/config/databend-query-node-1.toml b/scripts/ci/deploy/config/databend-query-node-1.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # For flight rpc.
 flight_api_address = "0.0.0.0:9091"

diff --git a/scripts/ci/deploy/config/databend-query-node-2.toml b/scripts/ci/deploy/config/databend-query-node-2.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # For flight rpc.
 flight_api_address = "0.0.0.0:9092"

diff --git a/scripts/ci/deploy/config/databend-query-node-3.toml b/scripts/ci/deploy/config/databend-query-node-3.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # For flight rpc.
 flight_api_address = "0.0.0.0:9093"

diff --git a/scripts/ci/deploy/config/databend-query-node-hive.toml b/scripts/ci/deploy/config/databend-query-node-hive.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # For flight rpc.
 flight_api_address = "0.0.0.0:9090"

diff --git a/scripts/ci/deploy/config/databend-query-node-native.toml b/scripts/ci/deploy/config/databend-query-node-native.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # For flight rpc.
 flight_api_address = "0.0.0.0:9091"

diff --git a/scripts/ci/deploy/config/databend-query-node-share-1.toml b/scripts/ci/deploy/config/databend-query-node-share-1.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # For flight rpc.
 flight_api_address = "0.0.0.0:19091"

diff --git a/scripts/ci/deploy/config/databend-query-node-share-2.toml b/scripts/ci/deploy/config/databend-query-node-share-2.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # For flight rpc.
 flight_api_address = "0.0.0.0:29091"

diff --git a/scripts/ci/deploy/config/databend-query-node-share-3.toml b/scripts/ci/deploy/config/databend-query-node-share-3.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # For flight rpc.
 flight_api_address = "0.0.0.0:19391"

diff --git a/scripts/distribution/configs/databend-query.toml b/scripts/distribution/configs/databend-query.toml
@@ -3,7 +3,7 @@
 
 [query]
 max_active_sessions = 256
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 
 # Internal flight rpc for cluster communication.
 flight_api_address = "0.0.0.0:9091"

diff --git a/src/binaries/query/entry.rs b/src/binaries/query/entry.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::env;
+use std::time::Duration;
 
 use databend_common_base::mem_allocator::GlobalAllocator;
 use databend_common_base::runtime::GLOBAL_MEM_STAT;
@@ -348,7 +349,11 @@ pub async fn start_services(conf: &InnerConfig) -> Result<()> {
         // for one shot background service, we need to drop it manually.
         drop(shutdown_handle);
     } else {
-        shutdown_handle.wait_for_termination_request().await;
+        let graceful_shutdown_timeout =
+            Some(Duration::from_millis(conf.query.shutdown_wait_timeout_ms));
+        shutdown_handle
+            .wait_for_termination_request(graceful_shutdown_timeout)
+            .await;
     }
     info!("Shutdown server.");
     Ok(())

diff --git a/src/query/config/src/config.rs b/src/query/config/src/config.rs
@@ -1490,7 +1490,7 @@ pub struct QueryConfig {
     pub table_engine_memory_enabled: bool,
 
     #[clap(long, value_name = "VALUE", default_value = "5000")]
-    pub wait_timeout_mills: u64,
+    pub shutdown_wait_timeout_ms: u64,
 
     #[clap(long, value_name = "VALUE", default_value = "10000")]
     pub max_query_log_size: usize,
@@ -1703,7 +1703,7 @@ impl TryInto<InnerQueryConfig> for QueryConfig {
             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,
+            shutdown_wait_timeout_ms: self.shutdown_wait_timeout_ms,
             max_query_log_size: self.max_query_log_size,
             databend_enterprise_license: self.databend_enterprise_license,
             management_mode: self.management_mode,
@@ -1787,7 +1787,7 @@ impl From<InnerQueryConfig> for QueryConfig {
             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,
+            shutdown_wait_timeout_ms: inner.shutdown_wait_timeout_ms,
             max_query_log_size: inner.max_query_log_size,
             databend_enterprise_license: inner.databend_enterprise_license,
             management_mode: inner.management_mode,

diff --git a/src/query/config/src/inner.rs b/src/query/config/src/inner.rs
@@ -190,7 +190,8 @@ pub struct QueryConfig {
     pub rpc_client_timeout_secs: u64,
     /// Table engine memory enabled
     pub table_engine_memory_enabled: bool,
-    pub wait_timeout_mills: u64,
+    /// Graceful shutdown timeout
+    pub shutdown_wait_timeout_ms: u64,
     pub max_query_log_size: usize,
     pub databend_enterprise_license: Option<String>,
     /// If in management mode, only can do some meta level operations(database/table/user/stage etc.) with metasrv.
@@ -269,7 +270,7 @@ impl Default for QueryConfig {
             rpc_tls_query_service_domain_name: "localhost".to_string(),
             rpc_client_timeout_secs: 0,
             table_engine_memory_enabled: true,
-            wait_timeout_mills: 5000,
+            shutdown_wait_timeout_ms: 5000,
             max_query_log_size: 10_000,
             databend_enterprise_license: None,
             management_mode: false,

diff --git a/src/query/service/src/servers/server.rs b/src/query/service/src/servers/server.rs
@@ -16,6 +16,7 @@ use std::net::SocketAddr;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
+use std::time::Duration;
 
 use databend_common_base::base::signal_stream;
 use databend_common_base::base::DummySignalStream;
@@ -67,17 +68,17 @@ impl ShutdownHandle {
     }
 
     #[async_backtrace::framed]
-    pub async fn shutdown(&mut self, mut signal: SignalStream) {
+    pub async fn shutdown(&mut self, mut signal: SignalStream, timeout: Option<Duration>) {
         self.shutdown_services(true).await;
         ClusterDiscovery::instance()
             .unregister_to_metastore(&mut signal)
             .await;
-        self.sessions.graceful_shutdown(signal, 5).await;
+        self.sessions.graceful_shutdown(signal, timeout).await;
         self.shutdown_services(false).await;
     }
 
     #[async_backtrace::framed]
-    pub async fn wait_for_termination_request(&mut self) {
+    pub async fn wait_for_termination_request(&mut self, timeout: Option<Duration>) {
         match signal_stream() {
             Err(cause) => {
                 error!("Cannot set shutdown signal handler, {:?}", cause);
@@ -91,7 +92,7 @@ impl ShutdownHandle {
                     self.shutdown
                         .compare_exchange(false, true, Ordering::SeqCst, Ordering::Acquire)
                 {
-                    let shutdown_services = self.shutdown(stream);
+                    let shutdown_services = self.shutdown(stream, timeout);
                     shutdown_services.await;
                 }
             }
@@ -110,7 +111,7 @@ impl Drop for ShutdownHandle {
                 .compare_exchange(false, true, Ordering::SeqCst, Ordering::Acquire)
         {
             let signal_stream = DummySignalStream::create(SignalType::Exit);
-            futures::executor::block_on(self.shutdown(signal_stream));
+            futures::executor::block_on(self.shutdown(signal_stream, Some(Duration::from_secs(5))));
         }
     }
 }
diff --git a/src/query/service/src/sessions/session_mgr.rs b/src/query/service/src/sessions/session_mgr.rs
@@ -200,27 +200,35 @@ impl SessionManager {
     pub fn graceful_shutdown(
         &self,
         mut signal: SignalStream,
-        timeout_secs: i32,
+        timeout: Option<Duration>,
     ) -> impl Future<Output = ()> {
         let active_sessions = self.active_sessions.clone();
         async move {
-            info!(
-                "Waiting {} secs for connections to close. You can press Ctrl + C again to force shutdown.",
-                timeout_secs
-            );
-            let mut signal = Box::pin(signal.next());
-
-            for _index in 0..timeout_secs {
-                if SessionManager::destroy_idle_sessions(&active_sessions) {
-                    return;
-                }
+            if let Some(mut timeout) = timeout {
+                info!(
+                    "Waiting {:?} for connections to close. You can press Ctrl + C again to force shutdown.",
+                    timeout
+                );
+
+                let mut signal = Box::pin(signal.next());
 
-                let interval = Duration::from_secs(1);
-                let sleep = Box::pin(tokio::time::sleep(interval));
-                match futures::future::select(sleep, signal).await {
-                    Either::Right((_, _)) => break,
-                    Either::Left((_, reserve_signal)) => signal = reserve_signal,
-                };
+                while !timeout.is_zero() {
+                    if SessionManager::destroy_idle_sessions(&active_sessions) {
+                        return;
+                    }
+
+                    let interval = Duration::from_secs(1);
+                    let sleep = Box::pin(tokio::time::sleep(interval));
+                    match futures::future::select(sleep, signal).await {
+                        Either::Right((_, _)) => break,
+                        Either::Left((_, reserve_signal)) => signal = reserve_signal,
+                    };
+
+                    timeout = match timeout > Duration::from_secs(1) {
+                        true => timeout - Duration::from_secs(1),
+                        false => Duration::from_secs(0),
+                    };
+                }
             }
 
             info!("Will shutdown forcefully.");

diff --git a/src/query/service/tests/it/configs.rs b/src/query/service/tests/it/configs.rs
@@ -735,7 +735,7 @@ rpc_tls_server_key = ""
 rpc_tls_query_server_root_ca_cert = ""
 rpc_tls_query_service_domain_name = "localhost"
 table_engine_memory_enabled = true
-wait_timeout_mills = 5000
+shutdown_wait_timeout_ms = 5000
 max_query_log_size = 10000
 management_mode = false
 jwt_key_file = ""

diff --git a/src/query/service/tests/it/storages/testdata/configs_table_basic.txt b/src/query/service/tests/it/storages/testdata/configs_table_basic.txt
@@ -121,11 +121,11 @@ DB.Table: 'system'.'configs', Table: configs-table_id:1, ver:0, Engine: SystemCo
 | 'query'   | 'rpc_tls_server_key'                       | ''                                                             | ''       |
 | 'query'   | 'share_endpoint_address'                   | ''                                                             | ''       |
 | 'query'   | 'share_endpoint_auth_token_file'           | ''                                                             | ''       |
+| 'query'   | 'shutdown_wait_timeout_ms'                 | '5000'                                                         | ''       |
 | 'query'   | 'table_engine_memory_enabled'              | 'true'                                                         | ''       |
 | 'query'   | 'tenant_id'                                | 'test'                                                         | ''       |
 | 'query'   | 'udf_server_allow_list'                    | ''                                                             | ''       |
 | 'query'   | 'users'                                    | '{"name":"root","auth_type":"no_password","auth_string":null}' | ''       |
-| 'query'   | 'wait_timeout_mills'                       | '5000'                                                         | ''       |
 | 'storage' | 'allow_insecure'                           | 'false'                                                        | ''       |
 | 'storage' | 'azblob.account_key'                       | ''                                                             | ''       |
 | 'storage' | 'azblob.account_name'                      | ''                                                             | ''       |