From ba16dfb1e897d239d95aa6249a61d83e6383b974 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Tue, 24 Oct 2023 16:58:49 +0100 Subject: [PATCH] add tokio metrics collector --- .cargo/config.toml | 1 + Cargo.lock | 13 +- Cargo.toml | 2 +- libs/metrics/Cargo.toml | 1 + libs/metrics/src/lib.rs | 3 + libs/metrics/src/tokio_metrics.rs | 307 ++++++++++++++++++++++++++++++ 6 files changed, 320 insertions(+), 7 deletions(-) create mode 100644 libs/metrics/src/tokio_metrics.rs diff --git a/.cargo/config.toml b/.cargo/config.toml index cc767a7f68..bd99b53e9b 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -16,6 +16,7 @@ opt-level = 1 # This is only present for local builds, as it will be overridden # by the RUSTDOCFLAGS env var in CI. rustdocflags = ["-Arustdoc::private_intra_doc_links"] +rustflags = ["--cfg=tokio_unstable"] [alias] build_testing = ["build", "--features", "testing"] diff --git a/Cargo.lock b/Cargo.lock index f7598a79cf..075b9ad7f3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2676,9 +2676,9 @@ checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" [[package]] name = "libc" -version = "0.2.144" +version = "0.2.149" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b00cc1c228a6782d0f076e7b232802e0c5689d41bb5df366f2a6b6621cfdfe1" +checksum = "a08173bc88b7955d1b3145aa561539096c421ac8debde8cbc3612ec635fee29b" [[package]] name = "libloading" @@ -2790,6 +2790,7 @@ dependencies = [ "libc", "once_cell", "prometheus", + "tokio", "workspace_hack", ] @@ -5357,18 +5358,18 @@ dependencies = [ [[package]] name = "tokio" -version = "1.28.1" +version = "1.33.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0aa32867d44e6f2ce3385e89dceb990188b8bb0fb25b0cf576647a6f98ac5105" +checksum = "4f38200e3ef7995e5ef13baec2f432a6da0aa9ac495b2c0e8f3b7eec2c92d653" dependencies = [ - "autocfg", + "backtrace", "bytes", "libc", "mio", "num_cpus", "pin-project-lite", "signal-hook-registry", - "socket2 0.4.9", + "socket2 0.5.3", "tokio-macros", "windows-sys 0.48.0", ] diff --git a/Cargo.toml b/Cargo.toml index a0be7bb9ac..2d598e3f9b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -136,7 +136,7 @@ tar = "0.4" test-context = "0.1" thiserror = "1.0" tls-listener = { version = "0.7", features = ["rustls", "hyper-h1"] } -tokio = { version = "1.17", features = ["macros"] } +tokio = { version = "1.29", features = ["macros"] } tokio-io-timeout = "1.2.0" tokio-postgres-rustls = "0.10.0" tokio-rustls = "0.24" diff --git a/libs/metrics/Cargo.toml b/libs/metrics/Cargo.toml index d4323ae766..f7a800dfc2 100644 --- a/libs/metrics/Cargo.toml +++ b/libs/metrics/Cargo.toml @@ -9,5 +9,6 @@ prometheus.workspace = true libc.workspace = true once_cell.workspace = true chrono.workspace = true +tokio.workspace = true workspace_hack.workspace = true diff --git a/libs/metrics/src/lib.rs b/libs/metrics/src/lib.rs index b4026d93e1..d29bfcc884 100644 --- a/libs/metrics/src/lib.rs +++ b/libs/metrics/src/lib.rs @@ -21,6 +21,9 @@ pub use prometheus::{register_int_gauge_vec, IntGaugeVec}; pub use prometheus::{Encoder, TextEncoder}; use prometheus::{Registry, Result}; +#[cfg(tokio_unstable)] +pub mod tokio_metrics; + pub mod launch_timestamp; mod wrappers; pub use wrappers::{CountedReader, CountedWriter}; diff --git a/libs/metrics/src/tokio_metrics.rs b/libs/metrics/src/tokio_metrics.rs new file mode 100644 index 0000000000..4ab4dea6f7 --- /dev/null +++ b/libs/metrics/src/tokio_metrics.rs @@ -0,0 +1,307 @@ +use std::collections::HashMap; + +use once_cell::sync::Lazy; +use prometheus::{ + core::Desc, + proto::{self, LabelPair}, +}; + +#[derive(Default)] +pub struct TokioCollector { + handles: Vec, +} + +impl TokioCollector { + pub fn add_runtime(&mut self, handle: tokio::runtime::Handle) { + self.handles.push(handle); + } + + fn guage( + &self, + desc: &Desc, + f: impl Fn(&tokio::runtime::RuntimeMetrics) -> f64, + ) -> proto::MetricFamily { + self.metric(desc, proto::MetricType::GAUGE, |rt| { + let mut g = proto::Gauge::default(); + g.set_value(f(rt)); + let mut m = proto::Metric::default(); + m.set_gauge(g); + m + }) + } + + fn metric( + &self, + desc: &Desc, + typ: proto::MetricType, + f: impl Fn(&tokio::runtime::RuntimeMetrics) -> proto::Metric, + ) -> proto::MetricFamily { + let mut m = proto::MetricFamily::default(); + m.set_name(desc.fq_name.clone()); + m.set_help(desc.help.clone()); + m.set_field_type(typ); + let mut metrics = vec![]; + for rt in &self.handles { + let rt_metrics = rt.metrics(); + let mut m = f(&rt_metrics); + + let mut label_id = LabelPair::default(); + label_id.set_name("id".to_owned()); + label_id.set_value(rt.id().to_string()); + + m.set_label(vec![label_id]); + metrics.push(m); + } + + m.set_metric(metrics); + m + } + + // fn guage_per_worker( + // &self, + // desc: &Desc, + // f: impl Fn(&tokio::runtime::RuntimeMetrics, usize) -> f64, + // ) -> proto::MetricFamily { + // self.metric_per_worker(desc, proto::MetricType::GAUGE, |m, i| { + // let mut g = proto::Gauge::default(); + // g.set_value(f(m, i)); + // let mut m = proto::Metric::default(); + // m.set_gauge(g); + // m + // }) + // } + + fn counter_per_worker( + &self, + desc: &Desc, + f: impl Fn(&tokio::runtime::RuntimeMetrics, usize) -> f64, + ) -> proto::MetricFamily { + self.metric_per_worker(desc, proto::MetricType::COUNTER, |m, i| { + let mut g = proto::Counter::default(); + g.set_value(f(m, i)); + let mut m = proto::Metric::default(); + m.set_counter(g); + m + }) + } + + fn metric_per_worker( + &self, + desc: &Desc, + typ: proto::MetricType, + f: impl Fn(&tokio::runtime::RuntimeMetrics, usize) -> proto::Metric, + ) -> proto::MetricFamily { + let mut m = proto::MetricFamily::default(); + m.set_name(desc.fq_name.clone()); + m.set_help(desc.help.clone()); + m.set_field_type(typ); + let mut metrics = vec![]; + for rt in &self.handles { + let rt_metrics = rt.metrics(); + let workers = rt_metrics.num_workers(); + + for i in 0..workers { + let mut m = f(&rt_metrics, i); + + let mut label_id = LabelPair::default(); + label_id.set_name("id".to_owned()); + label_id.set_value(rt.id().to_string()); + + let mut label_worker = LabelPair::default(); + label_worker.set_name("worker".to_owned()); + label_worker.set_value(i.to_string()); + + m.set_label(vec![label_id, label_worker]); + metrics.push(m); + } + } + + m.set_metric(metrics); + m + } +} + +static ACTIVE_TASK_COUNT: Lazy = Lazy::new(|| { + Desc::new( + "tokio_active_task_count_total".to_owned(), + "the number of active tasks in the runtime".to_owned(), + vec!["id".to_owned()], + HashMap::default(), + ) + .expect("should be a valid description") +}); +static WORKER_STEAL_COUNT: Lazy = Lazy::new(|| { + Desc::new( + "tokio_worker_steal_count_total".to_owned(), + "the number of tasks the given worker thread stole from another worker thread".to_owned(), + vec!["id".to_owned(), "worker".to_owned()], + HashMap::default(), + ) + .expect("should be a valid description") +}); + +impl prometheus::core::Collector for TokioCollector { + fn desc(&self) -> Vec<&Desc> { + vec![&ACTIVE_TASK_COUNT, &WORKER_STEAL_COUNT] + } + + fn collect(&self) -> Vec { + vec![ + self.guage(&ACTIVE_TASK_COUNT, |m| m.active_tasks_count() as f64), + self.counter_per_worker(&WORKER_STEAL_COUNT, |m, i| m.worker_steal_count(i) as f64), + ] + } +} + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use prometheus::{Registry, TextEncoder}; + use tokio::{runtime, sync::Mutex, task::JoinSet}; + + use crate::tokio_metrics::TokioCollector; + + #[tokio::test] + async fn gather() { + let registry = Registry::new(); + let mut collector = TokioCollector::default(); + collector.add_runtime(tokio::runtime::Handle::current()); + registry.register(Box::new(collector)).unwrap(); + + let lock = Arc::new(Mutex::new(0)); + let guard = lock.lock().await; + let mut joinset = JoinSet::new(); + for _ in 0..10 { + let lock = lock.clone(); + joinset.spawn(async move { + let mut _guard = lock.lock().await; + }); + } + + let text = TextEncoder.encode_to_string(®istry.gather()).unwrap(); + assert_eq!( + text, + r#"# HELP tokio_active_task_count_total the number of active tasks in the runtime +# TYPE tokio_active_task_count_total gauge +tokio_active_task_count_total{id="1"} 10 +# HELP tokio_worker_steal_count_total the number of tasks the given worker thread stole from another worker thread +# TYPE tokio_worker_steal_count_total counter +tokio_worker_steal_count_total{id="1",worker="0"} 0 +"# + ); + + drop(guard); + while let Some(_x) = joinset.join_next().await {} + + let text = TextEncoder.encode_to_string(®istry.gather()).unwrap(); + assert_eq!( + text, + r#"# HELP tokio_active_task_count_total the number of active tasks in the runtime +# TYPE tokio_active_task_count_total gauge +tokio_active_task_count_total{id="1"} 0 +# HELP tokio_worker_steal_count_total the number of tasks the given worker thread stole from another worker thread +# TYPE tokio_worker_steal_count_total counter +tokio_worker_steal_count_total{id="1",worker="0"} 0 +"# + ); + } + + #[test] + fn gather_multiple_runtimes() { + let registry = Registry::new(); + let mut collector = TokioCollector::default(); + + let runtime1 = runtime::Builder::new_current_thread().build().unwrap(); + let runtime2 = runtime::Builder::new_current_thread().build().unwrap(); + + collector.add_runtime(runtime1.handle().clone()); + collector.add_runtime(runtime2.handle().clone()); + + registry.register(Box::new(collector)).unwrap(); + + std::thread::scope(|s| { + let lock1 = Arc::new(Mutex::new(0)); + let lock2 = Arc::new(Mutex::new(0)); + + let guard1 = lock1.clone().try_lock_owned().unwrap(); + let guard2 = lock2.clone().try_lock_owned().unwrap(); + + s.spawn(move || { + runtime1.block_on(async { + let mut joinset = JoinSet::new(); + for _ in 0..5 { + let lock = lock1.clone(); + joinset.spawn(async move { + let mut _guard = lock.lock().await; + }); + } + while let Some(_x) = joinset.join_next().await {} + }) + }); + + s.spawn(move || { + runtime2.block_on(async { + let mut joinset = JoinSet::new(); + for _ in 0..5 { + let lock = lock2.clone(); + joinset.spawn(async move { + let mut _guard = lock.lock().await; + }); + } + while let Some(_x) = joinset.join_next().await {} + }) + }); + + std::thread::sleep(Duration::from_millis(10)); + + let text = TextEncoder.encode_to_string(®istry.gather()).unwrap(); + assert_eq!( + text, + r#"# HELP tokio_active_task_count_total the number of active tasks in the runtime +# TYPE tokio_active_task_count_total gauge +tokio_active_task_count_total{id="1"} 5 +tokio_active_task_count_total{id="2"} 5 +# HELP tokio_worker_steal_count_total the number of tasks the given worker thread stole from another worker thread +# TYPE tokio_worker_steal_count_total counter +tokio_worker_steal_count_total{id="1",worker="0"} 0 +tokio_worker_steal_count_total{id="2",worker="0"} 0 +"# + ); + + drop(guard1); + std::thread::sleep(Duration::from_millis(10)); + + let text = TextEncoder.encode_to_string(®istry.gather()).unwrap(); + assert_eq!( + text, + r#"# HELP tokio_active_task_count_total the number of active tasks in the runtime +# TYPE tokio_active_task_count_total gauge +tokio_active_task_count_total{id="1"} 0 +tokio_active_task_count_total{id="2"} 5 +# HELP tokio_worker_steal_count_total the number of tasks the given worker thread stole from another worker thread +# TYPE tokio_worker_steal_count_total counter +tokio_worker_steal_count_total{id="1",worker="0"} 0 +tokio_worker_steal_count_total{id="2",worker="0"} 0 +"# + ); + + drop(guard2); + std::thread::sleep(Duration::from_millis(10)); + + let text = TextEncoder.encode_to_string(®istry.gather()).unwrap(); + assert_eq!( + text, + r#"# HELP tokio_active_task_count_total the number of active tasks in the runtime +# TYPE tokio_active_task_count_total gauge +tokio_active_task_count_total{id="1"} 0 +tokio_active_task_count_total{id="2"} 0 +# HELP tokio_worker_steal_count_total the number of tasks the given worker thread stole from another worker thread +# TYPE tokio_worker_steal_count_total counter +tokio_worker_steal_count_total{id="1",worker="0"} 0 +tokio_worker_steal_count_total{id="2",worker="0"} 0 +"# + ); + }); + } +}