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/future_throttle_count_mode.rs

@@ -0,0 +1,100 @@
+// use core::panicking::panic;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use futures::FutureExt;
+use tokio::time::Sleep;
+
+use crate::runtime_throttle_count_mode::RuntimeThrottleShareWithFuture;
+
+enum State {
+    Common,
+    Backoff(Pin<Box<Sleep>>),
+}
+
+impl State {
+    fn unwrap_backoff(&mut self) -> &mut Pin<Box<Sleep>> {
+        match self {
+            State::Backoff(sleep) => sleep,
+            _ => panic!("unwrap_backoff failed"),
+        }
+    }
+}
+
+#[pin_project::pin_project]
+pub struct ThrottleFuture<F: Future + Send + 'static> {
+    #[pin]
+    future: F,
+    /// priority of this future
+    handle: Arc<RuntimeThrottleShareWithFuture>,
+    /// count of pendings, track the pending count for test
+    pub pend_cnt: u32,
+    state: State,
+}
+
+impl<F> ThrottleFuture<F>
+where
+    F: Future + Send + 'static,
+    F::Output: Send + 'static,
+{
+    pub fn new(handle: Arc<RuntimeThrottleShareWithFuture>, future: F) -> Self {
+        Self {
+            future,
+            handle,
+            pend_cnt: 0,
+            state: State::Common,
+        }
+    }
+}
+
+impl<F> Future for ThrottleFuture<F>
+where
+    F: Future + Send + 'static,
+    F::Output: Send + 'static,
+{
+    type Output = F::Output;
+
+    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+
+        match this.state {
+            State::Common => {}
+            State::Backoff(ref mut sleep) => match sleep.poll_unpin(cx) {
+                Poll::Ready(_) => {
+                    *this.state = State::Common;
+                }
+                Poll::Pending => return Poll::Pending,
+            },
+        };
+
+        if let Some(ratelimiter) = &this.handle.ratelimiter {
+            *this.pend_cnt += 1;
+
+            if let Err(wait) = ratelimiter.try_wait() {
+                *this.state = State::Backoff(Box::pin(tokio::time::sleep(wait)));
+                match this.state.unwrap_backoff().poll_unpin(cx) {
+                    Poll::Ready(_) => {
+                        *this.state = State::Common;
+                        cx.waker().clone().wake();
+                        return Poll::Pending;
+                    }
+                    Poll::Pending => {
+                        return Poll::Pending;
+                    }
+                }
+            }
+        }
+
+        let poll_res = this.future.poll(cx);
+
+        match poll_res {
+            Poll::Ready(r) => Poll::Ready(r),
+            Poll::Pending => {
+                *this.pend_cnt += 1;
+                Poll::Pending
+            }
+        }
+    }
+}

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

@@ -13,10 +13,16 @@
 // limitations under the License.
 
 pub mod error;
+mod future_throttle_count_mode;
 pub mod global;
 mod metrics;
 mod repeated_task;
 pub mod runtime;
+mod runtime_default;
+mod runtime_throttle_count_mode;
+#[cfg(test)]
+mod test_metrics;
+mod throttle_priority;
 
 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.