Skip to main content

runtime/
simple_observability.rs

1//! Simple High-Performance Observability
2//!
3//! Zero-overhead metrics collection and monitoring with minimal dependencies.
4
5use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
6use std::sync::Arc;
7use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
8use tokio::sync::RwLock;
9use tracing::{info, warn, error, instrument, Span};
10use serde::{Deserialize, Serialize};
11use once_cell::sync::Lazy;
12use sysinfo::System;
13
14/// Global metrics instance - zero overhead when disabled
15pub static METRICS: Lazy<Arc<SimpleMetrics>> = Lazy::new(|| {
16    Arc::new(SimpleMetrics::new())
17});
18
19/// Simple atomic counter-based metrics system
20pub struct SimpleMetrics {
21    enabled: bool,
22
23    // Request metrics (atomic for zero-lock performance)
24    requests_total: AtomicU64,
25    requests_in_flight: AtomicUsize,
26    errors_total: AtomicU64,
27
28    // Inference metrics
29    tokens_generated_total: AtomicU64,
30    cache_hits_total: AtomicU64,
31    cache_misses_total: AtomicU64,
32
33    // System state
34    health_monitor: Arc<RwLock<HealthMonitor>>,
35}
36
37impl SimpleMetrics {
38    pub fn new() -> Self {
39        let enabled = std::env::var("UNILLM_OBSERVABILITY_ENABLED")
40            .unwrap_or_else(|_| "true".to_string())
41            .parse::<bool>()
42            .unwrap_or(true);
43
44        Self {
45            enabled,
46            requests_total: AtomicU64::new(0),
47            requests_in_flight: AtomicUsize::new(0),
48            errors_total: AtomicU64::new(0),
49            tokens_generated_total: AtomicU64::new(0),
50            cache_hits_total: AtomicU64::new(0),
51            cache_misses_total: AtomicU64::new(0),
52            health_monitor: Arc::new(RwLock::new(HealthMonitor::new())),
53        }
54    }
55
56    /// Initialize the metrics system
57    pub fn initialize() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
58        if !METRICS.enabled {
59            info!("Observability disabled - zero overhead mode");
60            return Ok(());
61        }
62
63        // Initialize tracing
64        use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
65
66        let env_filter = EnvFilter::try_from_default_env()
67            .unwrap_or_else(|_| EnvFilter::new("info"));
68
69        tracing_subscriber::registry()
70            .with(env_filter)
71            .with(tracing_subscriber::fmt::layer().json())
72            .init();
73
74        // Start health monitoring
75        let health_monitor = Arc::clone(&METRICS.health_monitor);
76        tokio::spawn(async move {
77            let mut interval = tokio::time::interval(Duration::from_secs(10));
78            loop {
79                interval.tick().await;
80                health_monitor.write().await.update().await;
81            }
82        });
83
84        info!("Simple observability system initialized");
85        Ok(())
86    }
87
88    #[inline(always)]
89    pub fn is_enabled(&self) -> bool {
90        self.enabled
91    }
92
93    /// Record request start (zero overhead when disabled)
94    #[inline(always)]
95    pub fn request_start(&self) -> RequestTimer {
96        if !self.enabled {
97            return RequestTimer::disabled();
98        }
99
100        self.requests_total.fetch_add(1, Ordering::Relaxed);
101        self.requests_in_flight.fetch_add(1, Ordering::Relaxed);
102        RequestTimer::new()
103    }
104
105    /// Record token generation (zero overhead when disabled)
106    #[inline(always)]
107    pub fn token_generated(&self) {
108        if !self.enabled {
109            return;
110        }
111        self.tokens_generated_total.fetch_add(1, Ordering::Relaxed);
112    }
113
114    /// Record cache hit/miss (zero overhead when disabled)
115    #[inline(always)]
116    pub fn cache_hit(&self, hit: bool) {
117        if !self.enabled {
118            return;
119        }
120
121        if hit {
122            self.cache_hits_total.fetch_add(1, Ordering::Relaxed);
123        } else {
124            self.cache_misses_total.fetch_add(1, Ordering::Relaxed);
125        }
126    }
127
128    /// Record error (zero overhead when disabled)
129    #[inline(always)]
130    pub fn error(&self, error_type: &str) {
131        if !self.enabled {
132            return;
133        }
134
135        self.errors_total.fetch_add(1, Ordering::Relaxed);
136        error!("Error recorded: {}", error_type);
137    }
138
139    /// Get current statistics
140    pub fn get_stats(&self) -> MetricsSnapshot {
141        MetricsSnapshot {
142            requests_total: self.requests_total.load(Ordering::Relaxed),
143            requests_in_flight: self.requests_in_flight.load(Ordering::Relaxed),
144            errors_total: self.errors_total.load(Ordering::Relaxed),
145            tokens_generated_total: self.tokens_generated_total.load(Ordering::Relaxed),
146            cache_hits_total: self.cache_hits_total.load(Ordering::Relaxed),
147            cache_misses_total: self.cache_misses_total.load(Ordering::Relaxed),
148            cache_hit_rate: self.calculate_cache_hit_rate(),
149        }
150    }
151
152    /// Get health status
153    pub async fn get_health_status(&self) -> HealthStatus {
154        self.health_monitor.read().await.get_status()
155    }
156
157    /// Get Prometheus-style metrics
158    pub fn get_prometheus_metrics(&self) -> String {
159        if !self.enabled {
160            return String::new();
161        }
162
163        let stats = self.get_stats();
164        format!(
165            r#"# HELP unillm_requests_total Total requests processed
166# TYPE unillm_requests_total counter
167unillm_requests_total {}
168
169# HELP unillm_requests_in_flight Currently processing requests
170# TYPE unillm_requests_in_flight gauge
171unillm_requests_in_flight {}
172
173# HELP unillm_errors_total Total errors encountered
174# TYPE unillm_errors_total counter
175unillm_errors_total {}
176
177# HELP unillm_tokens_generated_total Total tokens generated
178# TYPE unillm_tokens_generated_total counter
179unillm_tokens_generated_total {}
180
181# HELP unillm_cache_hits_total Cache hits
182# TYPE unillm_cache_hits_total counter
183unillm_cache_hits_total {}
184
185# HELP unillm_cache_misses_total Cache misses
186# TYPE unillm_cache_misses_total counter
187unillm_cache_misses_total {}
188
189# HELP unillm_cache_hit_rate Cache hit percentage
190# TYPE unillm_cache_hit_rate gauge
191unillm_cache_hit_rate {:.2}
192"#,
193            stats.requests_total,
194            stats.requests_in_flight,
195            stats.errors_total,
196            stats.tokens_generated_total,
197            stats.cache_hits_total,
198            stats.cache_misses_total,
199            stats.cache_hit_rate
200        )
201    }
202
203    fn calculate_cache_hit_rate(&self) -> f64 {
204        let hits = self.cache_hits_total.load(Ordering::Relaxed) as f64;
205        let misses = self.cache_misses_total.load(Ordering::Relaxed) as f64;
206        let total = hits + misses;
207
208        if total > 0.0 {
209            (hits / total) * 100.0
210        } else {
211            0.0
212        }
213    }
214}
215
216/// Zero-overhead request timer
217pub struct RequestTimer {
218    start: Option<Instant>,
219}
220
221impl RequestTimer {
222    fn new() -> Self {
223        Self {
224            start: Some(Instant::now()),
225        }
226    }
227
228    fn disabled() -> Self {
229        Self { start: None }
230    }
231
232    /// Complete the request and record metrics
233    pub fn complete(self) {
234        // Timer automatically decrements in_flight counter in Drop
235    }
236}
237
238impl Drop for RequestTimer {
239    fn drop(&mut self) {
240        if self.start.is_some() && METRICS.enabled {
241            METRICS.requests_in_flight.fetch_sub(1, Ordering::Relaxed);
242        }
243    }
244}
245
246/// Metrics snapshot
247#[derive(Debug, Serialize)]
248pub struct MetricsSnapshot {
249    pub requests_total: u64,
250    pub requests_in_flight: usize,
251    pub errors_total: u64,
252    pub tokens_generated_total: u64,
253    pub cache_hits_total: u64,
254    pub cache_misses_total: u64,
255    pub cache_hit_rate: f64,
256}
257
258/// System health monitoring
259pub struct HealthMonitor {
260    system: System,
261    last_update: Instant,
262    status: HealthStatus,
263}
264
265impl HealthMonitor {
266    fn new() -> Self {
267        Self {
268            system: System::new_all(),
269            last_update: Instant::now(),
270            status: HealthStatus::default(),
271        }
272    }
273
274    async fn update(&mut self) {
275        self.system.refresh_all();
276        self.last_update = Instant::now();
277
278        self.status = HealthStatus {
279            status: "healthy".to_string(),
280            timestamp: SystemTime::now()
281                .duration_since(UNIX_EPOCH)
282                .unwrap()
283                .as_secs(),
284            cpu_usage_percent: self.system.global_cpu_usage(),
285            memory_usage_percent: (self.system.used_memory() as f32 / self.system.total_memory() as f32) * 100.0,
286            uptime_seconds: System::uptime(),
287            gpu_available: self.check_gpu_availability(),
288            memory_usage_bytes: self.system.used_memory(),
289            total_memory_bytes: self.system.total_memory(),
290        };
291    }
292
293    fn get_status(&self) -> HealthStatus {
294        self.status.clone()
295    }
296
297    fn check_gpu_availability(&self) -> bool {
298        std::env::var("CUDA_VISIBLE_DEVICES").is_ok() ||
299        std::path::Path::new("/dev/nvidia0").exists() ||
300        std::env::var("HIP_VISIBLE_DEVICES").is_ok()
301    }
302}
303
304/// Health status response
305#[derive(Debug, Clone, Serialize, Deserialize)]
306pub struct HealthStatus {
307    pub status: String,
308    pub timestamp: u64,
309    pub cpu_usage_percent: f32,
310    pub memory_usage_percent: f32,
311    pub uptime_seconds: u64,
312    pub gpu_available: bool,
313    pub memory_usage_bytes: u64,
314    pub total_memory_bytes: u64,
315}
316
317impl Default for HealthStatus {
318    fn default() -> Self {
319        Self {
320            status: "unknown".to_string(),
321            timestamp: SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
322            cpu_usage_percent: 0.0,
323            memory_usage_percent: 0.0,
324            uptime_seconds: 0,
325            gpu_available: false,
326            memory_usage_bytes: 0,
327            total_memory_bytes: 0,
328        }
329    }
330}
331
332/// Create a traced span for request processing
333#[instrument]
334pub fn start_request_span(request_id: &str, endpoint: &str) -> Span {
335    tracing::info_span!("request", request_id = request_id, endpoint = endpoint)
336}
337
338/// Create a traced span for inference
339#[instrument]
340pub fn start_inference_span(model: &str, batch_size: usize) -> Span {
341    tracing::info_span!("inference", model = model, batch_size = batch_size)
342}
343
344#[cfg(test)]
345mod tests {
346    use super::*;
347
348    #[tokio::test]
349    async fn test_metrics_zero_overhead() {
350        // Test basic metrics functionality
351        let timer = METRICS.request_start();
352        METRICS.token_generated();
353        METRICS.cache_hit(true);
354        METRICS.cache_hit(false);
355        METRICS.error("test_error");
356        timer.complete();
357
358        let stats = METRICS.get_stats();
359        if METRICS.is_enabled() {
360            assert!(stats.requests_total > 0);
361            assert!(stats.tokens_generated_total > 0);
362            assert!(stats.cache_hits_total > 0);
363            assert!(stats.cache_misses_total > 0);
364            assert!(stats.errors_total > 0);
365        }
366    }
367
368    #[tokio::test]
369    async fn test_health_monitoring() {
370        let health = METRICS.get_health_status().await;
371        assert!(health.timestamp > 0);
372    }
373
374    #[test]
375    fn test_prometheus_metrics_format() {
376        let metrics = METRICS.get_prometheus_metrics();
377        if METRICS.is_enabled() {
378            assert!(metrics.contains("unillm_requests_total"));
379        }
380    }
381
382    #[test]
383    fn test_cache_hit_rate() {
384        METRICS.cache_hit(true);
385        METRICS.cache_hit(true);
386        METRICS.cache_hit(false);
387
388        let stats = METRICS.get_stats();
389        if METRICS.is_enabled() {
390            // Should be around 66.67% (2 hits out of 3 total)
391            assert!(stats.cache_hit_rate > 60.0 && stats.cache_hit_rate < 70.0);
392        }
393    }
394}