feat(query): enable memory percentage quota for the workload group (#18328)

Author: zhang2014

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

@@ -163,6 +163,10 @@ impl MemStat {
         }
     }
 
+    pub fn get_limit(&self) -> i64 {
+        self.memory_limit.limit.load(Ordering::Relaxed)
+    }
+
     /// Feed memory usage stat to MemStat and return if it exceeds the limit.
     ///
     /// It feeds `state` to the this tracker and all of its ancestors, including GLOBAL_TRACKER.

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

@@ -15,6 +15,7 @@
 use std::collections::HashMap;
 use std::fmt::Display;
 use std::fmt::Formatter;
+use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -62,14 +63,20 @@ pub struct WorkloadGroupResource {
     pub meta: WorkloadGroup,
     pub queue_key: String,
     pub mem_stat: Arc<MemStat>,
+    pub max_memory_usage: Arc<AtomicUsize>,
     #[allow(clippy::type_complexity)]
-    pub destroy_fn: Option<Box<dyn FnOnce(&str) + Send + Sync + 'static>>,
+    pub destroy_fn: Option<Box<dyn FnOnce(&str, Option<usize>) + Send + Sync + 'static>>,
 }
 
 impl Drop for WorkloadGroupResource {
     fn drop(&mut self) {
         if let Some(destroy_fn) = self.destroy_fn.take() {
-            destroy_fn(&self.meta.id);
+            let mut mem_percentage = None;
+            if let Some(QuotaValue::Percentage(v)) = self.meta.quotas.get(MEMORY_QUOTA_KEY) {
+                mem_percentage = Some(*v);
+            }
+
+            destroy_fn(&self.meta.id, mem_percentage);
         }
     }
 }

src/query/management/src/lib.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#![feature(let_chains)]
+#![feature(get_mut_unchecked)]
 #![allow(clippy::uninlined_format_args)]
 
 mod connection;