|
| 1 | +use log::info; |
| 2 | +use std::time::Duration; |
| 3 | +use tokio::runtime::Handle; |
| 4 | +use tokio_metrics::RuntimeMonitor; |
| 5 | + |
| 6 | +pub struct TokioMetricsCollector { |
| 7 | + _reporter_handle: tokio::task::JoinHandle<()>, |
| 8 | +} |
| 9 | + |
| 10 | +impl TokioMetricsCollector { |
| 11 | + /// Start collecting metrics from the given runtime handle |
| 12 | + /// |
| 13 | + /// # Arguments |
| 14 | + /// * `handle` - Tokio runtime handle to monitor |
| 15 | + /// * `runtime_name` - Name prefix for metrics (e.g., "io_runtime" or "cpu_runtime") |
| 16 | + /// * `interval` - How often to poll and report metrics |
| 17 | + pub fn start(handle: Handle, runtime_name: String, interval: Duration) -> Self { |
| 18 | + info!( |
| 19 | + "Starting tokio metrics collection for {} with interval {:?}", |
| 20 | + runtime_name, interval |
| 21 | + ); |
| 22 | + |
| 23 | + let runtime_monitor = RuntimeMonitor::new(&handle); |
| 24 | + |
| 25 | + let reporter_handle = tokio::spawn(async move { |
| 26 | + let mut intervals = runtime_monitor.intervals(); |
| 27 | + let mut interval_timer = tokio::time::interval(interval); |
| 28 | + interval_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); |
| 29 | + |
| 30 | + loop { |
| 31 | + interval_timer.tick().await; |
| 32 | + |
| 33 | + if let Some(metrics) = intervals.next() { |
| 34 | + // Report metrics using metrics crate |
| 35 | + use metrics::{counter, gauge}; |
| 36 | + |
| 37 | + let prefix = &runtime_name; |
| 38 | + |
| 39 | + // Always available metrics |
| 40 | + gauge!(format!("{}_workers_count", prefix)).set(metrics.workers_count as f64); |
| 41 | + counter!(format!("{}_total_park_count", prefix)) |
| 42 | + .absolute(metrics.total_park_count); |
| 43 | + gauge!(format!("{}_max_park_count", prefix)).set(metrics.max_park_count as f64); |
| 44 | + gauge!(format!("{}_min_park_count", prefix)).set(metrics.min_park_count as f64); |
| 45 | + gauge!(format!("{}_total_busy_duration_us", prefix)) |
| 46 | + .set(metrics.total_busy_duration.as_micros() as f64); |
| 47 | + gauge!(format!("{}_max_busy_duration_us", prefix)) |
| 48 | + .set(metrics.max_busy_duration.as_micros() as f64); |
| 49 | + gauge!(format!("{}_min_busy_duration_us", prefix)) |
| 50 | + .set(metrics.min_busy_duration.as_micros() as f64); |
| 51 | + gauge!(format!("{}_global_queue_depth", prefix)) |
| 52 | + .set(metrics.global_queue_depth as f64); |
| 53 | + gauge!(format!("{}_elapsed_seconds", prefix)) |
| 54 | + .set(metrics.elapsed.as_secs_f64()); |
| 55 | + } |
| 56 | + } |
| 57 | + }); |
| 58 | + |
| 59 | + Self { |
| 60 | + _reporter_handle: reporter_handle, |
| 61 | + } |
| 62 | + } |
| 63 | + |
| 64 | + /// Start monitoring the current runtime |
| 65 | + pub fn start_current(runtime_name: String, interval: Duration) -> Self { |
| 66 | + let handle = Handle::current(); |
| 67 | + Self::start(handle, runtime_name, interval) |
| 68 | + } |
| 69 | +} |
| 70 | + |
| 71 | +#[cfg(test)] |
| 72 | +mod tests { |
| 73 | + use super::*; |
| 74 | + |
| 75 | + #[tokio::test] |
| 76 | + async fn test_metrics_collector_starts() { |
| 77 | + let _collector = TokioMetricsCollector::start_current( |
| 78 | + "test_runtime".to_string(), |
| 79 | + Duration::from_secs(1), |
| 80 | + ); |
| 81 | + |
| 82 | + // Sleep to allow at least one metrics report |
| 83 | + tokio::time::sleep(Duration::from_millis(1100)).await; |
| 84 | + |
| 85 | + // Collector should still be alive |
| 86 | + // In a real test, we'd verify metrics are being recorded |
| 87 | + } |
| 88 | +} |
0 commit comments