diff --git a/Cargo.lock b/Cargo.lock index 87827905fe..859940be5f 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 12a1d2d346..537b685af7 100644 --- a/Cargo.toml +++ b/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" } diff --git a/config/config.md b/config/config.md index 93575fcce4..3b0f99d697 100644 --- a/config/config.md +++ b/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.
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. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 9646936f00..7b81e98d3f 100644 --- a/config/datanode.example.toml +++ b/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. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 571353ec20..0fd0dd4a80 100644 --- a/config/standalone.example.toml +++ b/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. diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 5814cc33b4..9cb77ff398 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/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 = 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)] diff --git a/src/common/runtime/Cargo.toml b/src/common/runtime/Cargo.toml index 0b5ba5f096..927aa17b9c 100644 --- a/src/common/runtime/Cargo.toml +++ b/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 diff --git a/src/common/runtime/src/global.rs b/src/common/runtime/src/global.rs index e3f9b93280..edf428940f 100644 --- a/src/common/runtime/src/global.rs +++ b/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, } macro_rules! define_spawn { @@ -101,12 +137,42 @@ macro_rules! define_spawn { }; } +macro_rules! define_scheduled_spawn { + ($type: ident, $class: ident) => { + paste! { + fn [](&self, future: F) -> JoinHandle + 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 [](&self, future: F) -> JoinHandle + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + self.[<$type _runtime>].spawn_blocking(future) + } + + fn [](&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, @@ -114,6 +180,7 @@ impl GlobalRuntimes { heartbeat: Option, query: Option, ingest: Option, + workload_scheduler: Option, ) -> 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, query_runtime: Option, ingest_runtime: Option, + workload_scheduler: Option, already_init: bool, } @@ -149,9 +218,17 @@ static GLOBAL_RUNTIMES: Lazy = 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> = @@ -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 { + 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 { + 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 }); diff --git a/src/common/runtime/src/lib.rs b/src/common/runtime/src/lib.rs index 533fb57c6a..fc5b02cea3 100644 --- a/src/common/runtime/src/lib.rs +++ b/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}; diff --git a/src/common/runtime/src/metrics.rs b/src/common/runtime/src/metrics.rs index c332ccf6e6..3755fb00f4 100644 --- a/src/common/runtime/src/metrics.rs +++ b/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 { + 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))); +}