1use 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
14pub static METRICS: Lazy<Arc<SimpleMetrics>> = Lazy::new(|| {
16 Arc::new(SimpleMetrics::new())
17});
18
19pub struct SimpleMetrics {
21 enabled: bool,
22
23 requests_total: AtomicU64,
25 requests_in_flight: AtomicUsize,
26 errors_total: AtomicU64,
27
28 tokens_generated_total: AtomicU64,
30 cache_hits_total: AtomicU64,
31 cache_misses_total: AtomicU64,
32
33 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 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 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 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 #[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 #[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 #[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 #[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 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 pub async fn get_health_status(&self) -> HealthStatus {
154 self.health_monitor.read().await.get_status()
155 }
156
157 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
216pub 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 pub fn complete(self) {
234 }
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#[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
258pub 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#[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#[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#[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 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 assert!(stats.cache_hit_rate > 60.0 && stats.cache_hit_rate < 70.0);
392 }
393 }
394}