feat: add throttle runtime (#3685)

Cargo.toml

@@ -104,6 +104,7 @@ chrono = { version = "0.4", features = ["serde"] }
 clap = { version = "4.4", features = ["derive"] }
 config = "0.13.0"
 crossbeam-utils = "0.8"
+ratelimit = "0.9"
 dashmap = "5.4"
 datafusion = { git = "https://github.com/waynexia/arrow-datafusion.git", rev = "7823ef2f63663907edab46af0d51359900f608d6" }
 datafusion-common = { git = "https://github.com/waynexia/arrow-datafusion.git", rev = "7823ef2f63663907edab46af0d51359900f608d6" }
@@ -154,6 +155,7 @@ reqwest = { version = "0.12", default-features = false, features = [
     "stream",
     "multipart",
 ] }
+parking_lot = "0.12"
 rskafka = { git = "https://github.com/influxdata/rskafka.git", rev = "75535b5ad9bae4a5dbb582c82e44dfd81ec10105", features = [
     "transport-tls",
 ] }
@@ -245,6 +247,7 @@ store-api = { path = "src/store-api" }
 substrait = { path = "src/common/substrait" }
 table = { path = "src/table" }
 
+
 [patch.crates-io]
 # change all rustls dependencies to use our fork to default to `ring` to make it "just work"
 hyper-rustls = { git = "https://github.com/GreptimeTeam/hyper-rustls" }

src/client/Cargo.toml

@@ -28,7 +28,7 @@ enum_dispatch = "0.3"
 futures-util.workspace = true
 lazy_static.workspace = true
 moka = { workspace = true, features = ["future"] }
-parking_lot = "0.12"
+parking_lot.workspace = true
 prometheus.workspace = true
 prost.workspace = true
 query.workspace = true

src/common/runtime/.gitignore

@@ -0,0 +1,3 @@
+test_generate_dir
+priority_workload_cpu_usage.json
+scripts

src/common/runtime/Cargo.toml

@@ -23,6 +23,12 @@ tokio.workspace = true
 tokio-metrics = "0.3"
 tokio-metrics-collector = { git = "https://github.com/MichaelScofield/tokio-metrics-collector.git", rev = "89d692d5753d28564a7aac73c6ac5aba22243ba0" }
 tokio-util.workspace = true
+serde_json.workspace = true
+parking_lot.workspace = true
+ratelimit.workspace = true
+rand.workspace = true
+pin-project.workspace = true
+futures.workspace = true
 
 [dev-dependencies]
 tokio-test = "0.4"

src/common/runtime/README

@@ -0,0 +1,5 @@
+# Run performance test for different priority & workload type
+
+```
+cargo test --package common-runtime --lib -- test_metrics::test_all_cpu_usage --nocapture
+```

src/common/runtime/src/global.rs

@@ -21,6 +21,7 @@ use once_cell::sync::Lazy;
 use paste::paste;
 use serde::{Deserialize, Serialize};
 
+use crate::runtime::{BuilderBuild, RuntimeTrait};
 use crate::{Builder, JoinHandle, Runtime};
 
 const GLOBAL_WORKERS: usize = 8;

src/common/runtime/src/lib.rs

@@ -17,6 +17,8 @@ pub mod global;
 mod metrics;
 mod repeated_task;
 pub mod runtime;
+mod runtime_default;
+mod runtime_throttle_count_mode;
 
 pub use global::{
     block_on_compact, block_on_global, compact_runtime, create_runtime, global_runtime,

src/common/runtime/src/repeated_task.rs

@@ -23,6 +23,7 @@ use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
 
 use crate::error::{IllegalStateSnafu, Result, WaitGcTaskStopSnafu};
+use crate::runtime::RuntimeTrait;
 use crate::Runtime;
 
 /// Task to execute repeatedly.

src/common/runtime/src/runtime.rs

@@ -19,23 +19,19 @@ use std::thread;
 use std::time::Duration;
 
 use snafu::ResultExt;
-use tokio::runtime::{Builder as RuntimeBuilder, Handle};
+use tokio::runtime::Builder as RuntimeBuilder;
 use tokio::sync::oneshot;
 pub use tokio::task::{JoinError, JoinHandle};
 
 use crate::error::*;
 use crate::metrics::*;
+use crate::runtime_default::DefaultRuntime;
+use crate::runtime_throttle_count_mode::Priority;
 
-static RUNTIME_ID: AtomicUsize = AtomicUsize::new(0);
+// configurations
+pub type Runtime = DefaultRuntime;
 
-/// A runtime to run future tasks
-#[derive(Clone, Debug)]
-pub struct Runtime {
-    name: String,
-    handle: Handle,
-    // Used to receive a drop signal when dropper is dropped, inspired by databend
-    _dropper: Arc<Dropper>,
-}
+static RUNTIME_ID: AtomicUsize = AtomicUsize::new(0);
 
 /// Dropping the dropper will cause runtime to shutdown.
 #[derive(Debug)]
@@ -50,45 +46,42 @@ impl Drop for Dropper {
     }
 }
 
-impl Runtime {
-    pub fn builder() -> Builder {
+pub trait RuntimeTrait {
+    /// Get a runtime builder
+    fn builder() -> Builder {
         Builder::default()
     }
 
     /// Spawn a future and execute it in this thread pool
     ///
     /// Similar to tokio::runtime::Runtime::spawn()
-    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
+    fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
     where
         F: Future + Send + 'static,
-        F::Output: Send + 'static,
-    {
-        self.handle.spawn(future)
-    }
+        F::Output: Send + 'static;
 
     /// Run the provided function on an executor dedicated to blocking
     /// operations.
-    pub fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R>
+    fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R>
     where
         F: FnOnce() -> R + Send + 'static,
-        R: Send + 'static,
-    {
-        self.handle.spawn_blocking(func)
-    }
+        R: Send + 'static;
 
     /// Run a future to complete, this is the runtime's entry point
-    pub fn block_on<F: Future>(&self, future: F) -> F::Output {
-        self.handle.block_on(future)
-    }
+    fn block_on<F: Future>(&self, future: F) -> F::Output;
 
-    pub fn name(&self) -> &str {
-        &self.name
-    }
+    /// Get the name of the runtime
+    fn name(&self) -> &str;
+}
+
+pub trait BuilderBuild<R: RuntimeTrait> {
+    fn build(&mut self) -> Result<R>;
 }
 
 pub struct Builder {
     runtime_name: String,
     thread_name: String,
+    priority: Priority,
     builder: RuntimeBuilder,
 }
 
@@ -98,11 +91,17 @@ impl Default for Builder {
             runtime_name: format!("runtime-{}", RUNTIME_ID.fetch_add(1, Ordering::Relaxed)),
             thread_name: "default-worker".to_string(),
             builder: RuntimeBuilder::new_multi_thread(),
+            priority: Priority::VeryHigh,
         }
     }
 }
 
 impl Builder {
+    pub fn priority(&mut self, priority: Priority) -> &mut Self {
+        self.priority = priority;
+        self
+    }
+
     /// Sets the number of worker threads the Runtime will use.
     ///
     /// This can be any number above 0. The default value is the number of cores available to the system.
@@ -139,8 +138,10 @@ impl Builder {
         self.thread_name = val.into();
         self
     }
+}
 
-    pub fn build(&mut self) -> Result<Runtime> {
+impl BuilderBuild<DefaultRuntime> for Builder {
+    fn build(&mut self) -> Result<DefaultRuntime> {
         let runtime = self
             .builder
             .enable_all()
@@ -163,13 +164,13 @@ impl Builder {
         #[cfg(tokio_unstable)]
         register_collector(name.clone(), &handle);
 
-        Ok(Runtime {
-            name,
+        Ok(DefaultRuntime::new(
+            &name,
             handle,
-            _dropper: Arc::new(Dropper {
+            Arc::new(Dropper {
                 close: Some(send_stop),
             }),
-        })
+        ))
     }
 }
 
@@ -215,6 +216,7 @@ fn on_thread_unpark(thread_name: String) -> impl Fn() + 'static {
 
 #[cfg(test)]
 mod tests {
+
     use std::sync::Arc;
     use std::thread;
     use std::time::Duration;
@@ -235,12 +237,12 @@ mod tests {
 
     #[test]
     fn test_metric() {
-        let runtime = Builder::default()
+        let runtime: Runtime = Builder::default()
             .worker_threads(5)
             .thread_name("test_runtime_metric")
             .build()
             .unwrap();
-        // wait threads created
+        // wait threads create
         thread::sleep(Duration::from_millis(50));
 
         let _handle = runtime.spawn(async {

src/common/runtime/src/runtime_default.rs

@@ -0,0 +1,77 @@
+// Copyright 2023 Greptime Team
+//
+// 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::future::Future;
+use std::sync::Arc;
+
+use tokio::runtime::Handle;
+pub use tokio::task::JoinHandle;
+
+use crate::runtime::{Dropper, RuntimeTrait};
+use crate::Builder;
+
+/// A runtime to run future tasks
+#[derive(Clone, Debug)]
+pub struct DefaultRuntime {
+    name: String,
+    handle: Handle,
+    // Used to receive a drop signal when dropper is dropped, inspired by databend
+    _dropper: Arc<Dropper>,
+}
+
+impl DefaultRuntime {
+    pub fn new(name: &str, handle: Handle, dropper: Arc<Dropper>) -> Self {
+        Self {
+            name: name.to_string(),
+            handle,
+            _dropper: dropper,
+        }
+    }
+}
+
+impl RuntimeTrait for DefaultRuntime {
+    fn builder() -> Builder {
+        Builder::default()
+    }
+
+    /// Spawn a future and execute it in this thread pool
+    ///
+    /// Similar to tokio::runtime::Runtime::spawn()
+    fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        self.handle.spawn(future)
+    }
+
+    /// Run the provided function on an executor dedicated to blocking
+    /// operations.
+    fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R>
+    where
+        F: FnOnce() -> R + Send + 'static,
+        R: Send + 'static,
+    {
+        self.handle.spawn_blocking(func)
+    }
+
+    /// Run a future to complete, this is the runtime's entry point
+    fn block_on<F: Future>(&self, future: F) -> F::Output {
+        self.handle.block_on(future)
+    }
+
+    fn name(&self) -> &str {
+        &self.name
+    }
+}