feat(query): add statement_queued_timeout_in_seconds setting for queries queue (#14945)

* feat(query): add timeout for queries queue

* feat(query): add timeout for queries queue

Author: zhang2014

src/common/exception/src/exception_code.rs

@@ -149,6 +149,7 @@ build_exceptions! {
     UnknownCatalog(1119),
     UnknownCatalogType(1120),
     UnmatchMaskPolicyReturnType(1121),
+    Timeout(1122),
 
     // Data Related Errors
 

src/query/config/src/config.rs

@@ -1399,7 +1399,7 @@ pub struct QueryConfig {
     #[clap(long, value_name = "VALUE", default_value = "256")]
     pub max_active_sessions: u64,
 
-    #[clap(long, value_name = "VALUE", default_value = "0")]
+    #[clap(long, value_name = "VALUE", default_value = "8")]
     pub max_running_queries: u64,
 
     /// The max total memory in bytes that can be used by this process.

src/query/config/src/inner.rs

@@ -247,7 +247,7 @@ impl Default for QueryConfig {
             mysql_tls_server_cert: "".to_string(),
             mysql_tls_server_key: "".to_string(),
             max_active_sessions: 256,
-            max_running_queries: 0,
+            max_running_queries: 8,
             max_server_memory_usage: 0,
             max_memory_limit_enabled: false,
             clickhouse_http_handler_host: "127.0.0.1".to_string(),

src/query/service/src/sessions/queue_mgr.rs

@@ -22,6 +22,7 @@ use std::sync::Arc;
 use std::task::Context;
 use std::task::Poll;
 use std::task::Waker;
+use std::time::Duration;
 use std::time::SystemTime;
 
 use databend_common_base::base::GlobalInstance;
@@ -35,6 +36,7 @@ use pin_project_lite::pin_project;
 use tokio::sync::AcquireError;
 use tokio::sync::OwnedSemaphorePermit;
 use tokio::sync::Semaphore;
+use tokio::time::error::Elapsed;
 
 use crate::sessions::QueryContext;
 
@@ -44,6 +46,8 @@ pub trait QueueData: Send + Sync + 'static {
     fn get_key(&self) -> Self::Key;
 
     fn remove_error_message(key: Option<Self::Key>) -> ErrorCode;
+
+    fn timeout(&self) -> Duration;
 }
 
 pub(crate) struct Inner<Data: QueueData> {
@@ -96,9 +100,10 @@ impl<Data: QueueData> QueueManager<Data> {
     }
 
     pub async fn acquire(self: &Arc<Self>, data: Data) -> Result<AcquireQueueGuard> {
+        let timeout = data.timeout();
         let future = AcquireQueueFuture::create(
             Arc::new(data),
-            self.semaphore.clone().acquire_owned(),
+            tokio::time::timeout(timeout, self.semaphore.clone().acquire_owned()),
             self.clone(),
         );
 
@@ -131,7 +136,7 @@ impl AcquireQueueGuard {
 
 pin_project! {
     pub struct AcquireQueueFuture<Data: QueueData, T>
-where T: Future<Output = Result<OwnedSemaphorePermit, AcquireError>>
+where T: Future<Output = Result<Result<OwnedSemaphorePermit, AcquireError>, Elapsed>>
 {
     #[pin]
     inner: T,
@@ -146,7 +151,7 @@ where T: Future<Output = Result<OwnedSemaphorePermit, AcquireError>>
 }
 
 impl<Data: QueueData, T> AcquireQueueFuture<Data, T>
-where T: Future<Output = Result<OwnedSemaphorePermit, AcquireError>>
+where T: Future<Output = Result<Result<OwnedSemaphorePermit, AcquireError>, Elapsed>>
 {
     pub fn create(data: Arc<Data>, inner: T, mgr: Arc<QueueManager<Data>>) -> Self {
         AcquireQueueFuture {
@@ -161,7 +166,7 @@ where T: Future<Output = Result<OwnedSemaphorePermit, AcquireError>>
 }
 
 impl<Data: QueueData, T> Future for AcquireQueueFuture<Data, T>
-where T: Future<Output = Result<OwnedSemaphorePermit, AcquireError>>
+where T: Future<Output = Result<Result<OwnedSemaphorePermit, AcquireError>, Elapsed>>
 {
     type Output = Result<AcquireQueueGuard>;
 
@@ -181,8 +186,9 @@ where T: Future<Output = Result<OwnedSemaphorePermit, AcquireError>>
                 }
 
                 Poll::Ready(match res {
-                    Ok(v) => Ok(AcquireQueueGuard::create(v)),
-                    Err(_) => Err(ErrorCode::TokioError("acquire queue failure.")),
+                    Ok(Ok(v)) => Ok(AcquireQueueGuard::create(v)),
+                    Ok(Err(_)) => Err(ErrorCode::TokioError("acquire queue failure.")),
+                    Err(_elapsed) => Err(ErrorCode::Timeout("query queuing timeout")),
                 })
             }
             Poll::Pending => {
@@ -209,14 +215,20 @@ pub struct QueryEntry {
     pub query_id: String,
     pub create_time: SystemTime,
     pub user_info: UserInfo,
+    pub timeout: Duration,
 }
 
 impl QueryEntry {
     pub fn create(ctx: &Arc<QueryContext>) -> Result<QueryEntry> {
+        let settings = ctx.get_settings();
         Ok(QueryEntry {
             query_id: ctx.get_id(),
             create_time: ctx.get_created_time(),
             user_info: ctx.get_current_user()?,
+            timeout: match settings.get_statement_queued_timeout()? {
+                0 => Duration::from_secs(60 * 60 * 24 * 365 * 35),
+                timeout => Duration::from_secs(timeout),
+            },
         })
     }
 }
@@ -237,6 +249,10 @@ impl QueueData for QueryEntry {
             )),
         }
     }
+
+    fn timeout(&self) -> Duration {
+        self.timeout
+    }
 }
 
 pub type QueriesQueueManager = QueueManager<QueryEntry>;

src/query/service/tests/it/sessions/queue_mgr.rs

@@ -36,6 +36,10 @@ impl QueueData for TestData {
     fn remove_error_message(key: Option<Self::Key>) -> ErrorCode {
         ErrorCode::Internal(format!("{:?}", key))
     }
+
+    fn timeout(&self) -> Duration {
+        Duration::from_secs(1000)
+    }
 }
 
 #[tokio::test(flavor = "multi_thread")]

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

@@ -99,7 +99,7 @@ DB.Table: 'system'.'configs', Table: configs-table_id:1, ver:0, Engine: SystemCo
 | 'query'   | 'max_active_sessions'                      | '256'                                                          | ''       |
 | 'query'   | 'max_memory_limit_enabled'                 | 'false'                                                        | ''       |
 | 'query'   | 'max_query_log_size'                       | '10000'                                                        | ''       |
-| 'query'   | 'max_running_queries'                      | '0'                                                            | ''       |
+| 'query'   | 'max_running_queries'                      | '8'                                                            | ''       |
 | 'query'   | 'max_server_memory_usage'                  | '0'                                                            | ''       |
 | 'query'   | 'max_storage_io_requests'                  | 'null'                                                         | ''       |
 | 'query'   | 'metric_api_address'                       | '127.0.0.1:7070'                                               | ''       |

src/query/settings/src/settings_default.rs

@@ -709,6 +709,12 @@ impl DefaultSettings {
                     mode: SettingMode::Both,
                     range: None,
                 }),
+                ("statement_queued_timeout_in_seconds", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(0),
+                    desc: "The maximum waiting seconds in the queue. The default value is 0(no limit).",
+                    mode: SettingMode::Both,
+                    range: Some(SettingRange::Numeric(0..=u64::MAX)),
+                })
             ]);
 
             Ok(Arc::new(DefaultSettings {

src/query/settings/src/settings_getter_setter.rs

@@ -626,4 +626,8 @@ impl Settings {
     pub fn get_enable_experimental_queries_executor(&self) -> Result<bool> {
         Ok(self.try_get_u64("enable_experimental_queries_executor")? == 1)
     }
+
+    pub fn get_statement_queued_timeout(&self) -> Result<u64> {
+        self.try_get_u64("statement_queued_timeout_in_seconds")
+    }
 }