feat: Limit CPU in 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"
+clocksource = "0.8.1"
 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,2 @@
+test_generate_dir
+priority_workload_cpu_usage.json

src/common/runtime/Cargo.toml

@@ -23,6 +23,13 @@ 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
+sysinfo.workspace = true
+parking_lot.workspace = true
+clocksource.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/scripts/draw.py

@@ -0,0 +1,57 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import json
+import os
+os.chdir(os.path.dirname(__file__))
+
+# read json from ../priority_workload_cpu_usage.json
+with open('../priority_workload_cpu_usage.json', 'r') as f:
+    data = json.load(f)
+
+
+# 解构数据
+priorities = [0, 1, 2, 3, 4]
+load_types = [0, 1, 2, 3, 4]
+values = np.zeros((len(priorities), len(load_types)))
+
+for key, value in data.items():
+    p, l = map(int, key.split(','))
+    values[p, l] = value
+
+# 绘图配置
+bar_width = 0.15
+index = np.arange(len(load_types))
+
+# 创建柱状图
+fig, ax = plt.subplots()
+bars = []
+colors = ['b', 'g', 'r', 'c', 'm']
+
+for i in range(len(priorities)):
+    bar = ax.bar(index + i * bar_width, values[i], bar_width, label=f'Priority {i}', color=colors[i])
+    bars.append(bar)
+
+# 添加值标签
+def add_labels(rects):
+    for rect in rects:
+        height = rect.get_height()
+        ax.annotate('{}'.format(height),
+                    xy=(rect.get_x() + rect.get_width() / 2, height),
+                    xytext=(0, 3),  # 3 points vertical offset
+                    textcoords="offset points",
+                    rotation=90,
+                    ha='center', va='bottom')
+
+for bar_group in bars:
+    add_labels(bar_group)
+
+# 设置图表标题和坐标轴标签
+ax.set_title('CPU Usage by Priority and Load Type')
+ax.set_xlabel('Workload Type')
+ax.set_ylabel('CPU Usage')
+ax.set_xticks(index + bar_width * 2)
+ax.set_xticklabels(load_types)
+ax.legend()
+
+# 显示图表
+plt.show()

src/common/runtime/scripts/priority_workload_cpu_usage_count.json

@@ -0,0 +1 @@
+{"0,0":262.3619,"0,1":256.2622,"0,2":516.4156,"0,3":522.557,"1,0":342.16656,"1,1":389.12814,"1,2":523.02625,"1,3":551.30365,"2,0":437.40067,"2,1":449.49277,"2,2":523.3377,"2,3":560.76587,"3,0":543.88837,"3,1":534.6068,"3,2":522.6544,"3,3":571.3978,"4,0":592.19385,"4,1":595.3774,"4,2":529.2562,"4,3":551.40424}

src/common/runtime/scripts/priority_workload_cpu_usage_time.json

@@ -0,0 +1,22 @@
+{
+    "0,0": 262.3619,
+    "0,1": 256.2622,
+    "0,2": 516.4156,
+    "0,3": 522.557,
+    "1,0": 342.16656,
+    "1,1": 389.12814,
+    "1,2": 523.02625,
+    "1,3": 551.30365,
+    "2,0": 437.40067,
+    "2,1": 449.49277,
+    "2,2": 523.3377,
+    "2,3": 560.76587,
+    "3,0": 543.88837,
+    "3,1": 534.6068,
+    "3,2": 522.6544,
+    "3,3": 571.3978,
+    "4,0": 592.19385,
+    "4,1": 595.3774,
+    "4,2": 529.2562,
+    "4,3": 551.40424
+}

src/common/runtime/src/future_throttle_count_mode.rs

@@ -0,0 +1,128 @@
+// use core::panicking::panic;
+use std::any::TypeId;
+use std::collections::{BTreeMap, HashMap, VecDeque};
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use std::sync::Arc;
+use std::task::{Context, Poll, Waker};
+use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
+
+use futures::FutureExt;
+use parking_lot::lock_api::RwLock;
+use parking_lot::Mutex;
+use prometheus::core::{Atomic, AtomicF64};
+use rand::{thread_rng, Rng};
+use tokio::time::Sleep;
+
+use crate::ratelimit::Ratelimiter;
+use crate::runtime_throttle_count_mode::RuntimeThrottleShareWithFuture;
+
+struct OneFutureRecord {
+    total_micro: u64,
+    cpu_micro: u64,
+}
+
+// static ref EACH_FUTURE_RECORD: parking_lot::RwLock<Option<std::sync::mpsc::Sender<(Priority,Waker)>>> = RwLock::new(None);
+
+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
+    pub pend_cnt: u32, // track the pending count for test
+    /// count of inserted pendings
+    // pub inserted_pend_cnt: u32,
+    // sche_time: u32,
+    // poll_time: 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,
+            // inserted_pend_cnt: 0,
+            state: State::Common,
+            // poll_time: 0,
+            // sche_time: 0,
+        }
+    }
+}
+
+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();
+        // let sche_begin = Instant::now();
+        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,
+            },
+        };
+
+        // let inter = 5;
+
+        // if (*this.pend_cnt + 1) % inter == 0 {
+        if let Some(ratelimiter) = &this.handle.ratelimiter {
+            *this.pend_cnt += 1;
+            // println!("try wait for {}", avg_poll_time);
+            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_begin = Instant::now();
+        let poll_res = this.future.poll(cx);
+        let poll_time = poll_begin.elapsed().as_micros() as u64;
+
+        match poll_res {
+            Poll::Ready(r) => Poll::Ready(r),
+            Poll::Pending => {
+                *this.pend_cnt += 1;
+                Poll::Pending
+            }
+        }
+    }
+}