tokio metrics

Signed-off-by: Jason Volk <jason@zemos.net>

src/core/Cargo.toml

@@ -82,6 +82,7 @@ tikv-jemalloc-ctl.workspace = true
 tikv-jemalloc-sys.optional = true
 tikv-jemalloc-sys.workspace = true
 tokio.workspace = true
+tokio-metrics.workspace = true
 tracing-core.workspace = true
 tracing-subscriber.workspace = true
 tracing.workspace = true

src/core/metrics/mod.rs

@@ -0,0 +1,72 @@
+use std::sync::atomic::AtomicU32;
+
+use tokio::runtime;
+use tokio_metrics::TaskMonitor;
+#[cfg(tokio_unstable)]
+use tokio_metrics::{RuntimeIntervals, RuntimeMonitor};
+
+pub struct Metrics {
+	_runtime: Option<runtime::Handle>,
+
+	runtime_metrics: Option<runtime::RuntimeMetrics>,
+
+	task_monitor: Option<TaskMonitor>,
+
+	#[cfg(tokio_unstable)]
+	_runtime_monitor: Option<RuntimeMonitor>,
+
+	#[cfg(tokio_unstable)]
+	runtime_intervals: std::sync::Mutex<Option<RuntimeIntervals>>,
+
+	// TODO: move stats
+	pub requests_spawn_active: AtomicU32,
+	pub requests_spawn_finished: AtomicU32,
+	pub requests_handle_active: AtomicU32,
+	pub requests_handle_finished: AtomicU32,
+	pub requests_panic: AtomicU32,
+}
+
+impl Metrics {
+	#[must_use]
+	pub fn new(runtime: Option<runtime::Handle>) -> Self {
+		#[cfg(tokio_unstable)]
+		let runtime_monitor = runtime.as_ref().map(RuntimeMonitor::new);
+
+		#[cfg(tokio_unstable)]
+		let runtime_intervals = runtime_monitor.as_ref().map(RuntimeMonitor::intervals);
+
+		Self {
+			_runtime: runtime.clone(),
+
+			runtime_metrics: runtime.as_ref().map(runtime::Handle::metrics),
+
+			task_monitor: runtime.map(|_| TaskMonitor::new()),
+
+			#[cfg(tokio_unstable)]
+			_runtime_monitor: runtime_monitor,
+
+			#[cfg(tokio_unstable)]
+			runtime_intervals: std::sync::Mutex::new(runtime_intervals),
+
+			requests_spawn_active: AtomicU32::new(0),
+			requests_spawn_finished: AtomicU32::new(0),
+			requests_handle_active: AtomicU32::new(0),
+			requests_handle_finished: AtomicU32::new(0),
+			requests_panic: AtomicU32::new(0),
+		}
+	}
+
+	#[cfg(tokio_unstable)]
+	pub fn runtime_interval(&self) -> Option<tokio_metrics::RuntimeMetrics> {
+		self.runtime_intervals
+			.lock()
+			.expect("locked")
+			.as_mut()
+			.map(Iterator::next)
+			.expect("next interval")
+	}
+
+	pub fn task_root(&self) -> Option<&TaskMonitor> { self.task_monitor.as_ref() }
+
+	pub fn runtime_metrics(&self) -> Option<&runtime::RuntimeMetrics> { self.runtime_metrics.as_ref() }
+}

src/core/mod.rs

@@ -3,6 +3,7 @@ pub mod config;
 pub mod debug;
 pub mod error;
 pub mod log;
+pub mod metrics;
 pub mod mods;
 pub mod pdu;
 pub mod server;

src/core/server.rs

@@ -1,11 +1,11 @@
 use std::{
-	sync::atomic::{AtomicBool, AtomicU32, Ordering},
+	sync::atomic::{AtomicBool, Ordering},
 	time::SystemTime,
 };
 
 use tokio::{runtime, sync::broadcast};
 
-use crate::{config::Config, log, Err, Result};
+use crate::{config::Config, log::Log, metrics::Metrics, Err, Result};
 
 /// Server runtime state; public portion
 pub struct Server {
@@ -33,33 +33,25 @@ pub struct Server {
 	pub signal: broadcast::Sender<&'static str>,
 
 	/// Logging subsystem state
-	pub log: log::Log,
+	pub log: Log,
 
-	/// TODO: move stats
-	pub requests_spawn_active: AtomicU32,
-	pub requests_spawn_finished: AtomicU32,
-	pub requests_handle_active: AtomicU32,
-	pub requests_handle_finished: AtomicU32,
-	pub requests_panic: AtomicU32,
+	/// Metrics subsystem state
+	pub metrics: Metrics,
 }
 
 impl Server {
 	#[must_use]
-	pub fn new(config: Config, runtime: Option<runtime::Handle>, log: log::Log) -> Self {
+	pub fn new(config: Config, runtime: Option<runtime::Handle>, log: Log) -> Self {
 		Self {
 			config,
 			started: SystemTime::now(),
 			stopping: AtomicBool::new(false),
 			reloading: AtomicBool::new(false),
 			restarting: AtomicBool::new(false),
-			runtime,
+			runtime: runtime.clone(),
 			signal: broadcast::channel::<&'static str>(1).0,
 			log,
-			requests_spawn_active: AtomicU32::new(0),
-			requests_spawn_finished: AtomicU32::new(0),
-			requests_handle_active: AtomicU32::new(0),
-			requests_handle_finished: AtomicU32::new(0),
-			requests_panic: AtomicU32::new(0),
+			metrics: Metrics::new(runtime),
 		}
 	}