refactor(query): refactor runtime part 4 (#14728)

* refactor(query): extract stat buffer

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

* refactor(query): extract mem stat

src/binaries/query/ee_main.rs

@@ -19,6 +19,7 @@ mod entry;
 
 use databend_common_base::mem_allocator::GlobalAllocator;
 use databend_common_base::runtime::Runtime;
+use databend_common_base::runtime::ThreadTracker;
 use databend_common_config::InnerConfig;
 use databend_common_exception::Result;
 use databend_enterprise_query::enterprise_services::EnterpriseServices;
@@ -31,6 +32,8 @@ use crate::entry::start_services;
 pub static GLOBAL_ALLOCATOR: GlobalAllocator = GlobalAllocator;
 
 fn main() {
+    ThreadTracker::init();
+
     match Runtime::with_default_worker_threads() {
         Err(cause) => {
             eprintln!("Databend Query start failure, cause: {:?}", cause);

src/binaries/query/entry.rs

@@ -16,8 +16,8 @@ use std::env;
 use std::time::Duration;
 
 use databend_common_base::mem_allocator::GlobalAllocator;
+use databend_common_base::runtime::set_alloc_error_hook;
 use databend_common_base::runtime::GLOBAL_MEM_STAT;
-use databend_common_base::set_alloc_error_hook;
 use databend_common_config::Commands;
 use databend_common_config::InnerConfig;
 use databend_common_config::DATABEND_COMMIT_VERSION;

src/binaries/query/oss_main.rs

@@ -19,6 +19,7 @@ mod entry;
 
 use databend_common_base::mem_allocator::GlobalAllocator;
 use databend_common_base::runtime::Runtime;
+use databend_common_base::runtime::ThreadTracker;
 use databend_common_config::InnerConfig;
 use databend_common_exception::Result;
 use databend_common_license::license_manager::LicenseManager;
@@ -32,6 +33,8 @@ use crate::entry::start_services;
 pub static GLOBAL_ALLOCATOR: GlobalAllocator = GlobalAllocator;
 
 fn main() {
+    ThreadTracker::init();
+
     match Runtime::with_default_worker_threads() {
         Err(cause) => {
             eprintln!("Databend Query start failure, cause: {:?}", cause);

src/common/base/Cargo.toml

@@ -28,7 +28,6 @@ databend-common-exception = { path = "../exception" }
 
 # Crates.io dependencies
 async-backtrace = { workspace = true }
-async-channel = "2"
 async-trait = { workspace = true }
 bytesize = "1.1.0"
 ctrlc = { version = "3.2.3", features = ["termination"] }

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

@@ -0,0 +1,30 @@
+// 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 crate::runtime::LimitMemGuard;
+use crate::runtime::ThreadTracker;
+
+pub fn set_alloc_error_hook() {
+    std::alloc::set_alloc_error_hook(|layout| {
+        let _guard = LimitMemGuard::enter_unlimited();
+
+        let out_of_limit_desc = ThreadTracker::replace_error_message(None);
+
+        panic!(
+            "{}",
+            out_of_limit_desc
+                .unwrap_or_else(|| format!("memory allocation of {} bytes failed", layout.size()))
+        );
+    })
+}