feat(query): enable workload groups for resource control over the http handler (#17994)

* feat(query): enable workload groups for resource control over the http handler

* feat(query): enable workload groups for resource control over the http handler

* feat(query): enable workload groups for resource control over the http handler

* feat(query): enable workload groups for resource control over the http handler

* feat(query): enable workload groups for resource control over the http handler

* feat(query): enable workload groups for resource control over the http handler

Author: zhang2014

src/common/base/src/mem_allocator/tracker.rs

@@ -553,7 +553,7 @@ mod tests {
     ) -> R {
         {
             let mem_stat = MemStat::create_child(
-                GlobalUniqName::unique(),
+                Some(GlobalUniqName::unique()),
                 0,
                 ParentMemStat::Normal(global.clone()),
             );

src/common/base/src/runtime/memory/mem_stat.rs

@@ -98,11 +98,15 @@ impl MemStat {
     }
 
     pub fn create(name: String) -> Arc<MemStat> {
-        MemStat::create_child(name, 0, ParentMemStat::StaticRef(&GLOBAL_MEM_STAT))
+        MemStat::create_child(Some(name), 0, ParentMemStat::StaticRef(&GLOBAL_MEM_STAT))
+    }
+
+    pub fn create_workload_group() -> Arc<MemStat> {
+        MemStat::create_child(None, 0, ParentMemStat::StaticRef(&GLOBAL_MEM_STAT))
     }
 
     pub fn create_child(
-        name: String,
+        name: Option<String>,
         priority: usize,
         parent_memory_stat: ParentMemStat,
     ) -> Arc<MemStat> {
@@ -120,7 +124,7 @@ impl MemStat {
             exceeded_mutex: Mutex::new(false),
             memory_limit: MemoryLimit::default(),
             exceeded_memory: AtomicBool::new(false),
-            memory_requester: Some(name),
+            memory_requester: name,
             queries_memory_manager: &GLOBAL_QUERIES_MANAGER,
         })
     }
@@ -138,8 +142,25 @@ impl MemStat {
                 .store(water_height, Ordering::Relaxed);
         }
 
+        let used = self.used.load(Ordering::Relaxed);
         self.memory_limit.limit.store(size, Ordering::Relaxed);
         self.memory_limit.set_limit.store(size, Ordering::Relaxed);
+
+        #[allow(clippy::collapsible_if)]
+        if used < size && self.memory_requester.is_some() {
+            if self.exceeded_memory.fetch_and(false, Ordering::SeqCst) {
+                let _guard = LimitMemGuard::enter_unlimited();
+
+                let mutex = self.exceeded_mutex.lock();
+                let mut mutex = mutex.unwrap_or_else(PoisonError::into_inner);
+
+                if *mutex {
+                    *mutex = false;
+                    self.queries_memory_manager
+                        .release_memory(self.memory_requester.as_ref());
+                }
+            }
+        }
     }
 
     /// Feed memory usage stat to MemStat and return if it exceeds the limit.
@@ -471,7 +492,7 @@ mod tests {
     fn test_multiple_level_mem_stat() -> Result<()> {
         let mem_stat = MemStat::create("TEST".to_string());
         let child_mem_stat = MemStat::create_child(
-            "TEST_CHILD".to_string(),
+            Some("TEST_CHILD".to_string()),
             0,
             ParentMemStat::Normal(mem_stat.clone()),
         );
@@ -498,7 +519,7 @@ mod tests {
         let mem_stat = MemStat::create("TEST".to_string());
         mem_stat.set_limit(MINIMUM_MEMORY_LIMIT * 2, false);
         let child_mem_stat = MemStat::create_child(
-            "TEST_CHILD".to_string(),
+            Some("TEST_CHILD".to_string()),
             0,
             ParentMemStat::Normal(mem_stat.clone()),
         );
@@ -531,7 +552,7 @@ mod tests {
         let mem_stat = MemStat::create("TEST".to_string());
         mem_stat.set_limit(MINIMUM_MEMORY_LIMIT, false);
         let child_mem_stat = MemStat::create_child(
-            "TEST_CHILD".to_string(),
+            Some("TEST_CHILD".to_string()),
             0,
             ParentMemStat::Normal(mem_stat.clone()),
         );
@@ -547,7 +568,7 @@ mod tests {
         let mem_stat = MemStat::create("TEST".to_string());
         mem_stat.set_limit(MINIMUM_MEMORY_LIMIT * 2, false);
         let child_mem_stat = MemStat::create_child(
-            "TEST_CHILD".to_string(),
+            Some("TEST_CHILD".to_string()),
             0,
             ParentMemStat::Normal(mem_stat.clone()),
         );

src/common/base/src/runtime/memory/memory_manager.rs

@@ -445,7 +445,7 @@ mod tests {
         parent_memory_stat: ParentMemStat,
         manager: &'static QueriesMemoryManager,
     ) -> Arc<MemStat> {
-        let mut mem_stat = MemStat::create_child(name, priority, parent_memory_stat);
+        let mut mem_stat = MemStat::create_child(Some(name), priority, parent_memory_stat);
         let mut_mem_stat = Arc::get_mut(&mut mem_stat).unwrap();
         mut_mem_stat.queries_memory_manager = manager;
         mem_stat

src/common/base/src/runtime/memory/stat_buffer_mem_stat.rs

@@ -246,7 +246,7 @@ mod tests {
         let mut buffer = MemStatBuffer::empty(&TEST_GLOBAL);
 
         let mem_stat = MemStat::create_child(
-            String::from("test"),
+            Some(String::from("test")),
             0,
             ParentMemStat::StaticRef(&TEST_GLOBAL),
         );
@@ -275,12 +275,12 @@ mod tests {
         let mut buffer = MemStatBuffer::empty(&TEST_GLOBAL);
 
         let mem_stat_1 = MemStat::create_child(
-            String::from("test"),
+            Some(String::from("test")),
             0,
             ParentMemStat::StaticRef(&TEST_GLOBAL),
         );
         let mem_stat_2 = MemStat::create_child(
-            String::from("test"),
+            Some(String::from("test")),
             0,
             ParentMemStat::StaticRef(&TEST_GLOBAL),
         );
@@ -308,7 +308,7 @@ mod tests {
         let mut buffer = MemStatBuffer::empty(&TEST_GLOBAL);
 
         let mem_stat = MemStat::create_child(
-            String::from("test"),
+            Some(String::from("test")),
             0,
             ParentMemStat::StaticRef(&TEST_GLOBAL),
         );
@@ -339,12 +339,12 @@ mod tests {
         let mut buffer = MemStatBuffer::empty(&TEST_GLOBAL);
 
         let mem_stat_1 = MemStat::create_child(
-            String::from("test"),
+            Some(String::from("test")),
             0,
             ParentMemStat::StaticRef(&TEST_GLOBAL),
         );
         let mem_stat_2 = MemStat::create_child(
-            String::from("test"),
+            Some(String::from("test")),
             0,
             ParentMemStat::StaticRef(&TEST_GLOBAL),
         );

src/common/base/src/runtime/mod.rs

@@ -25,6 +25,7 @@ mod runtime;
 mod runtime_tracker;
 mod thread;
 mod time_series;
+pub mod workload_group;
 
 pub use backtrace::dump_backtrace;
 pub use backtrace::get_all_tasks;

src/common/base/src/runtime/workload_group.rs

@@ -0,0 +1,56 @@
+// 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::collections::HashMap;
+use std::fmt::Display;
+use std::fmt::Formatter;
+use std::time::Duration;
+
+pub const CPU_QUOTA_KEY: &str = "cpu_quota";
+pub const MEMORY_QUOTA_KEY: &str = "memory_quota";
+pub const QUERY_TIMEOUT_QUOTA_KEY: &str = "query_timeout";
+pub const MAX_CONCURRENCY_QUOTA_KEY: &str = "max_concurrency";
+pub const QUERY_QUEUED_TIMEOUT_QUOTA_KEY: &str = "query_queued_timeout";
+
+#[derive(serde::Serialize, serde::Deserialize, Clone, Eq, PartialEq, Debug)]
+pub enum QuotaValue {
+    Duration(Duration),
+    Percentage(usize),
+    Bytes(usize),
+    Number(usize),
+}
+
+impl Display for QuotaValue {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            QuotaValue::Percentage(v) => write!(f, "{}%", v),
+            QuotaValue::Duration(v) => write!(f, "{:?}", v),
+            QuotaValue::Bytes(v) => write!(f, "{}", v),
+            QuotaValue::Number(v) => write!(f, "{}", v),
+        }
+    }
+}
+
+#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)]
+pub struct WorkloadGroup {
+    pub id: String,
+    pub name: String,
+    pub quotas: HashMap<String, QuotaValue>,
+}
+
+impl WorkloadGroup {
+    pub fn get_quota(&self, key: &'static str) -> Option<QuotaValue> {
+        self.quotas.get(key).cloned()
+    }
+}

src/query/management/src/workload/mod.rs

@@ -14,13 +14,9 @@
 
 mod workload_api;
 mod workload_mgr;
+mod workload_resource;
 
-pub use workload_api::QuotaValue;
 pub use workload_api::WorkloadApi;
-pub use workload_api::WorkloadGroup;
-pub use workload_api::CPU_QUOTA_KEY;
-pub use workload_api::MAX_CONCURRENCY_QUOTA_KEY;
-pub use workload_api::MEMORY_QUOTA_KEY;
-pub use workload_api::QUERY_QUEUED_TIMEOUT_QUOTA_KEY;
-pub use workload_api::QUERY_TIMEOUT_QUOTA_KEY;
 pub use workload_mgr::WorkloadMgr;
+pub use workload_resource::WorkloadGroupResource;
+pub use workload_resource::WorkloadGroupResourceManager;

src/query/management/src/workload/workload_api.rs

@@ -13,50 +13,11 @@
 // limitations under the License.
 
 use std::collections::HashMap;
-use std::fmt::Display;
-use std::fmt::Formatter;
-use std::time::Duration;
 
+use databend_common_base::runtime::workload_group::QuotaValue;
+use databend_common_base::runtime::workload_group::WorkloadGroup;
 use databend_common_exception::Result;
 
-pub const CPU_QUOTA_KEY: &str = "cpu_quota";
-pub const MEMORY_QUOTA_KEY: &str = "memory_quota";
-pub const QUERY_TIMEOUT_QUOTA_KEY: &str = "query_timeout";
-pub const MAX_CONCURRENCY_QUOTA_KEY: &str = "max_concurrency";
-pub const QUERY_QUEUED_TIMEOUT_QUOTA_KEY: &str = "query_queued_timeout";
-
-#[derive(serde::Serialize, serde::Deserialize, Clone, Eq, PartialEq, Debug)]
-pub enum QuotaValue {
-    Duration(Duration),
-    Percentage(usize),
-    Bytes(usize),
-    Number(usize),
-}
-
-impl Display for QuotaValue {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        match self {
-            QuotaValue::Percentage(v) => write!(f, "{}%", v),
-            QuotaValue::Duration(v) => write!(f, "{:?}", v),
-            QuotaValue::Bytes(v) => write!(f, "{}", v),
-            QuotaValue::Number(v) => write!(f, "{}", v),
-        }
-    }
-}
-
-#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)]
-pub struct WorkloadGroup {
-    pub id: String,
-    pub name: String,
-    pub quotas: HashMap<String, QuotaValue>,
-}
-
-impl WorkloadGroup {
-    pub fn get_quota(&self, key: &'static str) -> Option<QuotaValue> {
-        self.quotas.get(key).cloned()
-    }
-}
-
 #[async_trait::async_trait]
 pub trait WorkloadApi: Sync + Send {
     async fn create(&self, group: WorkloadGroup) -> Result<WorkloadGroup>;

src/query/management/src/workload/workload_mgr.rs

@@ -16,6 +16,8 @@ use std::collections::HashMap;
 
 use databend_common_base::base::escape_for_key;
 use databend_common_base::base::GlobalUniqName;
+use databend_common_base::runtime::workload_group::QuotaValue;
+use databend_common_base::runtime::workload_group::WorkloadGroup;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_meta_kvapi::kvapi::KVApi;
@@ -26,9 +28,7 @@ use databend_common_meta_types::TxnCondition;
 use databend_common_meta_types::TxnOp;
 use databend_common_meta_types::TxnRequest;
 
-use crate::workload::workload_api::QuotaValue;
 use crate::workload::workload_api::WorkloadApi;
-use crate::workload::workload_api::WorkloadGroup;
 pub static WORKLOAD_META_KEY_PREFIX: &str = "__fd_workloads";
 
 pub struct WorkloadMgr {