feat(runtime): add weighted workload scheduler

Cargo.lock

@@ -1667,6 +1667,14 @@ dependencies = [
  "tokio-stream",
 ]
 
+[[package]]
+name = "catio"
+version = "0.1.0"
+source = "git+https://github.com/waynexia/catio.git?rev=ddcd3e9e5a5fdd1e58866ed4e6dbbc2062bae44a#ddcd3e9e5a5fdd1e58866ed4e6dbbc2062bae44a"
+dependencies = [
+ "tokio",
+]
+
 [[package]]
 name = "cbc"
 version = "0.1.2"
@@ -2762,6 +2770,7 @@ name = "common-runtime"
 version = "1.2.0"
 dependencies = [
  "async-trait",
+ "catio",
  "clap",
  "common-error",
  "common-macro",

Cargo.toml

@@ -271,6 +271,7 @@ zstd = "0.13"
 api = { path = "src/api" }
 auth = { path = "src/auth" }
 cache = { path = "src/cache" }
+catio = { git = "https://github.com/waynexia/catio.git", rev = "ddcd3e9e5a5fdd1e58866ed4e6dbbc2062bae44a" }
 catalog = { path = "src/catalog" }
 cli = { path = "src/cli" }
 client = { path = "src/client" }

config/config.md

@@ -25,6 +25,11 @@
 | `runtime` | -- | -- | The runtime options. |
 | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. |
 | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute the runtime for global write operations. |
+| `runtime.experimental_workload_scheduler` | -- | -- | Experimental weighted, work-conserving query/write task scheduler. |
+| `runtime.experimental_workload_scheduler.enable` | Bool | `false` | Enable the scheduler. Disabled by default. |
+| `runtime.experimental_workload_scheduler.max_concurrent_polls` | Integer | `0` | Maximum task polls admitted to Tokio at once. Zero uses 4 * global_rt_size. |
+| `runtime.experimental_workload_scheduler.query_weight` | Integer | `2` | Relative query share while both query and write workloads are backlogged. |
+| `runtime.experimental_workload_scheduler.write_weight` | Integer | `8` | Relative write share while both query and write workloads are backlogged. |
 | `http` | -- | -- | The HTTP server options. |
 | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. |
 | `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout. |
@@ -491,6 +496,11 @@
 | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute the runtime for global write operations. |
 | `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.<br/>Defaults to max(num_cpus - 1, 1). |
 | `runtime.ingest_rt_size` | Integer | `8` | The number of threads to execute datanode ingestion operations. |
+| `runtime.experimental_workload_scheduler` | -- | -- | Experimental weighted, work-conserving query/write task scheduler. |
+| `runtime.experimental_workload_scheduler.enable` | Bool | `false` | Enable the scheduler. Disabled by default. |
+| `runtime.experimental_workload_scheduler.max_concurrent_polls` | Integer | `0` | Maximum task polls admitted to Tokio at once. Zero uses 4 * global_rt_size. |
+| `runtime.experimental_workload_scheduler.query_weight` | Integer | `2` | Relative query share while both query and write workloads are backlogged. |
+| `runtime.experimental_workload_scheduler.write_weight` | Integer | `8` | Relative write share while both query and write workloads are backlogged. |
 | `meta_client` | -- | -- | The metasrv client options. |
 | `meta_client.metasrv_addrs` | Array | -- | The addresses of the metasrv. |
 | `meta_client.timeout` | String | `3s` | Operation timeout. |

config/datanode.example.toml

@@ -88,6 +88,17 @@ watch = false
 ## The number of threads to execute datanode ingestion operations.
 #+ ingest_rt_size = 8
 
+## Experimental weighted, work-conserving query/write task scheduler.
+#+ [runtime.experimental_workload_scheduler]
+## Enable the scheduler. Disabled by default.
+#+ enable = false
+## Maximum task polls admitted to Tokio at once. Zero uses 4 * global_rt_size.
+#+ max_concurrent_polls = 0
+## Relative query share while both query and write workloads are backlogged.
+#+ query_weight = 2
+## Relative write share while both query and write workloads are backlogged.
+#+ write_weight = 8
+
 ## The metasrv client options.
 [meta_client]
 ## The addresses of the metasrv.

config/standalone.example.toml

@@ -56,6 +56,17 @@ max_concurrent_queries = 0
 ## The number of threads to execute the runtime for global write operations.
 #+ compact_rt_size = 4
 
+## Experimental weighted, work-conserving query/write task scheduler.
+#+ [runtime.experimental_workload_scheduler]
+## Enable the scheduler. Disabled by default.
+#+ enable = false
+## Maximum task polls admitted to Tokio at once. Zero uses 4 * global_rt_size.
+#+ max_concurrent_polls = 0
+## Relative query share while both query and write workloads are backlogged.
+#+ query_weight = 2
+## Relative write share while both query and write workloads are backlogged.
+#+ write_weight = 8
+
 ## The HTTP server options.
 [http]
 ## The address to bind the HTTP server.

src/cmd/tests/load_config_test.rs

@@ -48,6 +48,12 @@ fn test_load_datanode_runtime_options_from_runtime_section() {
         compact_rt_size = 4
         ingest_rt_size = 8
         query_rt_size = 7
+
+        [runtime.experimental_workload_scheduler]
+        enable = true
+        max_concurrent_polls = 6
+        query_weight = 1
+        write_weight = 4
     "#;
 
     let options: GreptimeOptions<DatanodeOptions> = toml::from_str(toml).unwrap();
@@ -56,6 +62,22 @@ fn test_load_datanode_runtime_options_from_runtime_section() {
     assert_eq!(4, options.runtime.compact_rt_size);
     assert_eq!(8, options.runtime.ingest_rt_size);
     assert_eq!(7, options.runtime.query_rt_size);
+    assert!(options.runtime.experimental_workload_scheduler.enable);
+    assert_eq!(
+        6,
+        options
+            .runtime
+            .experimental_workload_scheduler
+            .max_concurrent_polls
+    );
+    assert_eq!(
+        1,
+        options.runtime.experimental_workload_scheduler.query_weight
+    );
+    assert_eq!(
+        4,
+        options.runtime.experimental_workload_scheduler.write_weight
+    );
 }
 
 #[allow(deprecated)]

src/common/runtime/Cargo.toml

@@ -16,6 +16,7 @@ workspace = true
 
 [dependencies]
 async-trait.workspace = true
+catio.workspace = true
 clap.workspace = true
 common-error.workspace = true
 common-macro.workspace = true

src/common/runtime/src/global.rs

@@ -16,17 +16,49 @@
 use std::future::Future;
 use std::sync::{Mutex, Once};
 
-use common_telemetry::info;
+use catio::{Scheduler, SchedulerStats, TaskClass};
+use common_telemetry::{info, warn};
 use once_cell::sync::Lazy;
 use paste::paste;
 use serde::{Deserialize, Serialize};
 
+use crate::metrics::register_workload_scheduler_metrics;
 use crate::runtime::{BuilderBuild, RuntimeTrait};
 use crate::{Builder, JoinHandle, Runtime};
 
 const GLOBAL_WORKERS: usize = 8;
 const COMPACT_WORKERS: usize = 4;
 const HB_WORKERS: usize = 2;
+pub(crate) const QUERY_TASK_CLASS: TaskClass = TaskClass::new(1);
+pub(crate) const WRITE_TASK_CLASS: TaskClass = TaskClass::new(2);
+
+/// Experimental options for sharing Tokio capacity between query and write
+/// workloads.
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[serde(default)]
+pub struct WorkloadSchedulerOptions {
+    /// Enables policy-controlled query and write task spawning.
+    pub enable: bool,
+    /// Maximum polls admitted to Tokio at once. Zero uses four times
+    /// `global_rt_size` to keep worker queues fed without making them
+    /// effectively unbounded.
+    pub max_concurrent_polls: usize,
+    /// Relative share for query polls while writes are also backlogged.
+    pub query_weight: u32,
+    /// Relative share for write polls while queries are also backlogged.
+    pub write_weight: u32,
+}
+
+impl Default for WorkloadSchedulerOptions {
+    fn default() -> Self {
+        Self {
+            enable: false,
+            max_concurrent_polls: 0,
+            query_weight: 2,
+            write_weight: 8,
+        }
+    }
+}
 
 /// The options for the global runtimes.
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
@@ -40,6 +72,8 @@ pub struct RuntimeOptions {
     pub query_rt_size: usize,
     /// The number of threads to execute datanode ingestion operations.
     pub ingest_rt_size: usize,
+    /// Experimental weighted scheduler for query and write workloads.
+    pub experimental_workload_scheduler: WorkloadSchedulerOptions,
 }
 
 impl Default for RuntimeOptions {
@@ -50,6 +84,7 @@ impl Default for RuntimeOptions {
             compact_rt_size: usize::max(cpus / 2, 1),
             query_rt_size: usize::max(cpus.saturating_sub(1), 1),
             ingest_rt_size: cpus,
+            experimental_workload_scheduler: WorkloadSchedulerOptions::default(),
         }
     }
 }
@@ -72,6 +107,7 @@ struct GlobalRuntimes {
     hb_runtime: Runtime,
     query_runtime: Runtime,
     ingest_runtime: Runtime,
+    workload_scheduler: Option<Scheduler>,
 }
 
 macro_rules! define_spawn {
@@ -101,12 +137,42 @@ macro_rules! define_spawn {
     };
 }
 
+macro_rules! define_scheduled_spawn {
+    ($type: ident, $class: ident) => {
+        paste! {
+            fn [<spawn_ $type>]<F>(&self, future: F) -> JoinHandle<F::Output>
+            where
+                F: Future + Send + 'static,
+                F::Output: Send + 'static,
+            {
+                match &self.workload_scheduler {
+                    Some(scheduler) => self.[<$type _runtime>]
+                        .spawn(scheduler.schedule_in($class, future)),
+                    None => self.[<$type _runtime>].spawn(future),
+                }
+            }
+
+            fn [<spawn_blocking_ $type>]<F, R>(&self, future: F) -> JoinHandle<R>
+            where
+                F: FnOnce() -> R + Send + 'static,
+                R: Send + 'static,
+            {
+                self.[<$type _runtime>].spawn_blocking(future)
+            }
+
+            fn [<block_on_ $type>]<F: Future>(&self, future: F) -> F::Output {
+                self.[<$type _runtime>].block_on(future)
+            }
+        }
+    };
+}
+
 impl GlobalRuntimes {
     define_spawn!(global);
     define_spawn!(compact);
     define_spawn!(hb);
-    define_spawn!(query);
-    define_spawn!(ingest);
+    define_scheduled_spawn!(query, QUERY_TASK_CLASS);
+    define_scheduled_spawn!(ingest, WRITE_TASK_CLASS);
 
     fn new(
         global: Option<Runtime>,
@@ -114,6 +180,7 @@ impl GlobalRuntimes {
         heartbeat: Option<Runtime>,
         query: Option<Runtime>,
         ingest: Option<Runtime>,
+        workload_scheduler: Option<Scheduler>,
     ) -> Self {
         let global_runtime =
             global.unwrap_or_else(|| create_runtime("global", "global-worker", GLOBAL_WORKERS));
@@ -128,6 +195,7 @@ impl GlobalRuntimes {
                 .unwrap_or_else(|| create_runtime("heartbeat", "hb-worker", HB_WORKERS)),
             query_runtime,
             ingest_runtime,
+            workload_scheduler,
         }
     }
 }
@@ -139,6 +207,7 @@ struct ConfigRuntimes {
     hb_runtime: Option<Runtime>,
     query_runtime: Option<Runtime>,
     ingest_runtime: Option<Runtime>,
+    workload_scheduler: Option<Scheduler>,
     already_init: bool,
 }
 
@@ -149,9 +218,17 @@ static GLOBAL_RUNTIMES: Lazy<GlobalRuntimes> = Lazy::new(|| {
     let heartbeat = c.hb_runtime.take();
     let query = c.query_runtime.take();
     let ingest = c.ingest_runtime.take();
+    let workload_scheduler = c.workload_scheduler.take();
     c.already_init = true;
 
-    GlobalRuntimes::new(global, compact, heartbeat, query, ingest)
+    GlobalRuntimes::new(
+        global,
+        compact,
+        heartbeat,
+        query,
+        ingest,
+        workload_scheduler,
+    )
 });
 
 static CONFIG_RUNTIMES: Lazy<Mutex<ConfigRuntimes>> =
@@ -178,9 +255,50 @@ pub fn init_global_runtimes(options: &RuntimeOptions) {
             options.compact_rt_size,
         ));
         c.hb_runtime = Some(create_runtime("heartbeat", "hb-worker", HB_WORKERS));
+        c.workload_scheduler = create_workload_scheduler(options);
     });
 }
 
+fn create_workload_scheduler(options: &RuntimeOptions) -> Option<Scheduler> {
+    let scheduler_options = &options.experimental_workload_scheduler;
+    if !scheduler_options.enable {
+        return None;
+    }
+    if scheduler_options.query_weight == 0 || scheduler_options.write_weight == 0 {
+        warn!(
+            "The experimental workload scheduler is disabled because query_weight and \
+             write_weight must both be greater than zero"
+        );
+        return None;
+    }
+
+    let max_concurrent_polls = if scheduler_options.max_concurrent_polls == 0 {
+        options.global_rt_size.saturating_mul(4)
+    } else {
+        scheduler_options.max_concurrent_polls
+    };
+    if max_concurrent_polls == 0 {
+        warn!(
+            "The experimental workload scheduler is disabled because max_concurrent_polls \
+             resolved to zero"
+        );
+        return None;
+    }
+
+    let scheduler = Scheduler::builder()
+        .max_concurrent_polls(max_concurrent_polls)
+        .weight(QUERY_TASK_CLASS, scheduler_options.query_weight)
+        .weight(WRITE_TASK_CLASS, scheduler_options.write_weight)
+        .build();
+    register_workload_scheduler_metrics(scheduler.clone());
+    info!(
+        "Enabled the experimental workload scheduler: max_concurrent_polls={}, \
+         query_weight={}, write_weight={}",
+        max_concurrent_polls, scheduler_options.query_weight, scheduler_options.write_weight
+    );
+    Some(scheduler)
+}
+
 /// Initialize the datanode-specific global runtimes.
 ///
 /// # Panics
@@ -244,6 +362,15 @@ define_global_runtime_spawn!(hb);
 define_global_runtime_spawn!(query);
 define_global_runtime_spawn!(ingest);
 
+/// Returns scheduler counters when the experimental workload scheduler is
+/// enabled.
+pub fn workload_scheduler_stats() -> Option<SchedulerStats> {
+    GLOBAL_RUNTIMES
+        .workload_scheduler
+        .as_ref()
+        .map(Scheduler::stats)
+}
+
 #[cfg(test)]
 mod tests {
     use tokio_test::assert_ok;
@@ -259,6 +386,10 @@ mod tests {
         assert_eq!(usize::max(cpus / 2, 1), options.compact_rt_size);
         assert_eq!(usize::max(cpus.saturating_sub(1), 1), options.query_rt_size);
         assert_eq!(cpus, options.ingest_rt_size);
+        assert_eq!(
+            WorkloadSchedulerOptions::default(),
+            options.experimental_workload_scheduler
+        );
     }
 
     #[test]
@@ -269,6 +400,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         );
 
         assert_eq!("test-global", runtimes.global_runtime.name());
@@ -276,6 +408,50 @@ mod tests {
         assert_eq!("test-global", runtimes.ingest_runtime.name());
     }
 
+    #[test]
+    fn test_workload_scheduler_default_admission_window() {
+        let mut options = RuntimeOptions {
+            global_rt_size: 3,
+            ..RuntimeOptions::default()
+        };
+        options.experimental_workload_scheduler.enable = true;
+
+        let scheduler = create_workload_scheduler(&options).unwrap();
+        let stats = scheduler.stats();
+        assert_eq!(12, stats.max_concurrent_polls);
+        assert_eq!(2, stats.classes[&QUERY_TASK_CLASS].weight);
+        assert_eq!(8, stats.classes[&WRITE_TASK_CLASS].weight);
+    }
+
+    #[test]
+    fn test_workload_scheduler_wraps_query_and_write_spawns() {
+        let runtime = create_runtime("test-workload", "test-workload-worker", 2);
+        let scheduler = Scheduler::builder()
+            .max_concurrent_polls(2)
+            .weight(QUERY_TASK_CLASS, 2)
+            .weight(WRITE_TASK_CLASS, 8)
+            .build();
+        let runtimes = GlobalRuntimes::new(
+            Some(runtime.clone()),
+            Some(runtime.clone()),
+            Some(runtime.clone()),
+            Some(runtime.clone()),
+            Some(runtime.clone()),
+            Some(scheduler.clone()),
+        );
+
+        let query = runtimes.spawn_query(async { "query" });
+        let write = runtimes.spawn_ingest(async { "write" });
+        let (query, write) =
+            runtime.block_on(async { (query.await.unwrap(), write.await.unwrap()) });
+
+        assert_eq!("query", query);
+        assert_eq!("write", write);
+        let stats = scheduler.stats();
+        assert_eq!(1, stats.classes[&QUERY_TASK_CLASS].polls);
+        assert_eq!(1, stats.classes[&WRITE_TASK_CLASS].polls);
+    }
+
     #[test]
     fn test_datanode_runtime_spawn_block_on() {
         let handle = spawn_query(async { 1 + 1 });

src/common/runtime/src/lib.rs

@@ -25,7 +25,7 @@ pub use global::{
     create_runtime, global_runtime, ingest_runtime, init_datanode_runtimes, init_global_runtimes,
     query_runtime, spawn_blocking_compact, spawn_blocking_global, spawn_blocking_hb,
     spawn_blocking_ingest, spawn_blocking_query, spawn_compact, spawn_global, spawn_hb,
-    spawn_ingest, spawn_query,
+    spawn_ingest, spawn_query, workload_scheduler_stats,
 };
 
 pub use crate::repeated_task::{BoxedTaskFunction, RepeatedTask, TaskFunction};

src/common/runtime/src/metrics.rs

@@ -13,9 +13,14 @@
 // limitations under the License.
 
 //! Runtime metrics
+use catio::Scheduler;
 use lazy_static::lazy_static;
+use prometheus::core::{Collector, Desc};
+use prometheus::proto::MetricFamily;
 use prometheus::*;
 
+use crate::global::{QUERY_TASK_CLASS, WRITE_TASK_CLASS};
+
 pub const THREAD_NAME_LABEL: &str = "thread_name";
 
 lazy_static! {
@@ -32,3 +37,76 @@ lazy_static! {
     )
     .unwrap();
 }
+
+#[derive(Clone)]
+struct WorkloadSchedulerCollector {
+    scheduler: Scheduler,
+    polls: IntGaugeVec,
+    queued: IntGaugeVec,
+    active: IntGauge,
+}
+
+impl WorkloadSchedulerCollector {
+    fn new(scheduler: Scheduler) -> Self {
+        Self {
+            scheduler,
+            polls: IntGaugeVec::new(
+                Opts::new(
+                    "greptime_workload_scheduler_polls",
+                    "Cumulative task polls admitted by the workload scheduler",
+                ),
+                &["workload"],
+            )
+            .unwrap(),
+            queued: IntGaugeVec::new(
+                Opts::new(
+                    "greptime_workload_scheduler_queued_tasks",
+                    "Tasks queued in the workload scheduler",
+                ),
+                &["workload"],
+            )
+            .unwrap(),
+            active: IntGauge::new(
+                "greptime_workload_scheduler_active_polls",
+                "Task polls admitted to Tokio but not yet completed",
+            )
+            .unwrap(),
+        }
+    }
+
+    fn update(&self) {
+        let stats = self.scheduler.stats();
+        for (class, workload) in [(QUERY_TASK_CLASS, "query"), (WRITE_TASK_CLASS, "write")] {
+            let class_stats = stats.classes.get(&class).cloned().unwrap_or_default();
+            self.polls
+                .with_label_values(&[workload])
+                .set(class_stats.polls.min(i64::MAX as u64) as i64);
+            self.queued
+                .with_label_values(&[workload])
+                .set(class_stats.queued.min(i64::MAX as usize) as i64);
+        }
+        self.active
+            .set(stats.active_polls.min(i64::MAX as usize) as i64);
+    }
+}
+
+impl Collector for WorkloadSchedulerCollector {
+    fn desc(&self) -> Vec<&Desc> {
+        let mut desc = self.polls.desc();
+        desc.extend(self.queued.desc());
+        desc.extend(self.active.desc());
+        desc
+    }
+
+    fn collect(&self) -> Vec<MetricFamily> {
+        self.update();
+        let mut families = self.polls.collect();
+        families.extend(self.queued.collect());
+        families.extend(self.active.collect());
+        families
+    }
+}
+
+pub(crate) fn register_workload_scheduler_metrics(scheduler: Scheduler) {
+    let _ = register(Box::new(WorkloadSchedulerCollector::new(scheduler)));
+}