Skip to main content

aster_bench/
error_capture.rs

1use crate::bench_session::BenchAgentError;
2use chrono::Utc;
3use once_cell::sync::Lazy;
4use std::sync::Arc;
5use std::sync::RwLock;
6use tokio::sync::Mutex;
7use tracing::{Event, Subscriber};
8use tracing_subscriber::layer::Context;
9use tracing_subscriber::Layer;
10
11// Type alias to reduce complexity
12type ErrorRegistry = RwLock<Option<Arc<Mutex<Vec<BenchAgentError>>>>>;
13
14// Global registry for error vectors
15static ERROR_REGISTRY: Lazy<ErrorRegistry> = Lazy::new(|| RwLock::new(None));
16
17pub struct ErrorCaptureLayer;
18
19impl Default for ErrorCaptureLayer {
20    fn default() -> Self {
21        Self
22    }
23}
24
25impl ErrorCaptureLayer {
26    pub fn new() -> Self {
27        Self
28    }
29
30    pub fn register_error_vector(errors: Arc<Mutex<Vec<BenchAgentError>>>) {
31        if let Ok(mut registry) = ERROR_REGISTRY.write() {
32            *registry = Some(errors);
33        }
34    }
35}
36
37impl<S> Layer<S> for ErrorCaptureLayer
38where
39    S: Subscriber,
40{
41    fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
42        // Only capture error and warning level events
43        if *event.metadata().level() <= tracing::Level::WARN {
44            let mut visitor = JsonVisitor::new();
45            event.record(&mut visitor);
46
47            if let Some(message) = visitor.recorded_fields.get("message") {
48                let error = BenchAgentError {
49                    message: message.to_string(),
50                    level: event.metadata().level().to_string(),
51                    timestamp: Utc::now(),
52                };
53
54                // Get the current error vector from the registry
55                if let Ok(registry) = ERROR_REGISTRY.read() {
56                    if let Some(errors) = registry.clone() {
57                        tokio::spawn(async move {
58                            let mut errors = errors.lock().await;
59                            errors.push(error);
60                        });
61                    }
62                }
63            }
64        }
65    }
66}
67
68struct JsonVisitor {
69    recorded_fields: serde_json::Map<String, serde_json::Value>,
70}
71
72impl JsonVisitor {
73    fn new() -> Self {
74        Self {
75            recorded_fields: serde_json::Map::new(),
76        }
77    }
78}
79
80impl tracing::field::Visit for JsonVisitor {
81    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
82        self.recorded_fields.insert(
83            field.name().to_string(),
84            serde_json::Value::String(value.to_string()),
85        );
86    }
87
88    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
89        self.recorded_fields.insert(
90            field.name().to_string(),
91            serde_json::Value::String(format!("{:?}", value)),
92        );
93    }
94}