add tokio metrics collector

.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"]

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",
 ]

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"

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

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};

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<tokio::runtime::Handle>,
+}
+
+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<Desc> = 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<Desc> = 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<proto::MetricFamily> {
+        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(&registry.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(&registry.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(&registry.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(&registry.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(&registry.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
+"#
+            );
+        });
+    }
+}