From 2aa74d1aabc71c24412daabe4098bb01dabbb1a9 Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Sat, 11 May 2024 11:33:24 +0100
Subject: [PATCH] proxy: add measured-tokio

---
 .cargo/config.toml      |  1 +
 Cargo.lock              | 12 ++++++++++++
 Cargo.toml              |  1 +
 libs/metrics/Cargo.toml |  3 +++
 libs/metrics/src/lib.rs |  5 +++++
 proxy/src/bin/proxy.rs  |  4 ++++
 6 files changed, 26 insertions(+)

diff --git a/.cargo/config.toml b/.cargo/config.toml
index 5e452974ad..430dde2b53 100644
--- a/.cargo/config.toml
+++ b/.cargo/config.toml
@@ -2,6 +2,7 @@
 # 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 6ce7180d67..5e610f15c1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3038,6 +3038,17 @@ dependencies = [
  "procfs 0.16.0",
 ]
 
+[[package]]
+name = "measured-tokio"
+version = "0.0.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b4ed0773ddbda3a85e39d0e094934549c410ca686a5095bcd72fb62100252a0"
+dependencies = [
+ "itoa",
+ "measured",
+ "tokio",
+]
+
 [[package]]
 name = "memchr"
 version = "2.6.4"
@@ -3079,6 +3090,7 @@ dependencies = [
  "libc",
  "measured",
  "measured-process",
+ "measured-tokio",
  "once_cell",
  "procfs 0.14.2",
  "prometheus",
diff --git a/Cargo.toml b/Cargo.toml
index 17f30a1327..b49f47c33e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -109,6 +109,7 @@ leaky-bucket = "1.0.1"
 libc = "0.2"
 md5 = "0.7.0"
 measured = { version = "0.0.21", features=["lasso"] }
+measured-tokio = { version = "0.0.21" }
 measured-process = { version = "0.0.21" }
 memoffset = "0.8"
 native-tls = "0.2"
diff --git a/libs/metrics/Cargo.toml b/libs/metrics/Cargo.toml
index 0bd804051c..b585c2944a 100644
--- a/libs/metrics/Cargo.toml
+++ b/libs/metrics/Cargo.toml
@@ -18,6 +18,9 @@ workspace_hack.workspace = true
 procfs.workspace = true
 measured-process.workspace = true
 
+[target.'cfg(tokio_unstable)'.dependencies]
+measured-tokio.workspace = true
+
 [dev-dependencies]
 rand = "0.8"
 rand_distr = "0.4.3"
diff --git a/libs/metrics/src/lib.rs b/libs/metrics/src/lib.rs
index 141d8a6d01..27e2bf5dcf 100644
--- a/libs/metrics/src/lib.rs
+++ b/libs/metrics/src/lib.rs
@@ -148,6 +148,11 @@ pub struct NeonMetrics {
     #[metric(init = measured_process::ProcessCollector::for_self())]
     process: measured_process::ProcessCollector,
 
+    #[cfg(tokio_unstable)]
+    #[metric(namespace = "tokio")]
+    #[metric(init = measured_tokio::NamedRuntimesCollector::new())]
+    pub tokio: measured_tokio::NamedRuntimesCollector,
+
     #[metric(namespace = "libmetrics")]
     #[metric(init = LibMetrics::new(build_info))]
     libmetrics: LibMetrics,
diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs
index be7d961b8c..3572131c0c 100644
--- a/proxy/src/bin/proxy.rs
+++ b/proxy/src/bin/proxy.rs
@@ -268,6 +268,10 @@ async fn main() -> anyhow::Result<()> {
         build_tag: BUILD_TAG,
     });
 
+    // add the current runtime to the collector
+    #[cfg(tokio_unstable)]
+    neon_metrics.tokio.add_current("proxy");
+
     let jemalloc = match proxy::jemalloc::MetricRecorder::new() {
         Ok(t) => Some(t),
         Err(e) => {