use std::{
collections::HashMap,
sync::{Arc, RwLock},
};
#[derive(Debug, Clone)]
pub struct Span {
pub span_id: String,
pub parent_span_id: Option<String>,
pub trace_id: String,
pub name: String,
pub attributes: HashMap<String, String>,
pub status: String,
pub error_message: String,
pub duration_ms: f64,
pub layer: String,
}
#[derive(Clone, Default)]
pub struct SpanRecorder {
spans: Arc<RwLock<Vec<Span>>>,
}
impl SpanRecorder {
pub fn new() -> Self {
Self {
spans: Arc::new(RwLock::new(Vec::new())),
}
}
pub fn record(&self, span: Span) {
let mut spans = self.spans.write().unwrap();
spans.push(span);
}
pub fn spans(&self) -> Vec<Span> {
self.spans.read().unwrap().clone()
}
pub fn clear(&self) {
let mut spans = self.spans.write().unwrap();
spans.clear();
}
}
#[derive(Clone, Default)]
pub struct MetricsRecorder {
counters: Arc<RwLock<HashMap<String, u64>>>,
gauges: Arc<RwLock<HashMap<String, i64>>>,
histograms: Arc<RwLock<HashMap<String, Vec<f64>>>>,
}
impl MetricsRecorder {
pub fn new() -> Self {
Self {
counters: Arc::new(RwLock::new(HashMap::new())),
gauges: Arc::new(RwLock::new(HashMap::new())),
histograms: Arc::new(RwLock::new(HashMap::new())),
}
}
pub fn current() -> Self {
Self::new()
}
pub fn increment_counter(&self, name: &str, value: u64) {
let mut counters = self.counters.write().unwrap();
*counters.entry(name.to_string()).or_insert(0) += value;
}
pub fn set_gauge(&self, name: &str, value: i64) {
let mut gauges = self.gauges.write().unwrap();
gauges.insert(name.to_string(), value);
}
pub fn record_histogram(&self, name: &str, value: f64) {
let mut histograms = self.histograms.write().unwrap();
histograms.entry(name.to_string()).or_default().push(value);
}
pub fn get_counter(&self, name: &str) -> u64 {
self.counters
.read()
.unwrap()
.get(name)
.copied()
.unwrap_or(0)
}
pub fn get_gauge(&self, name: &str) -> i64 {
self.gauges.read().unwrap().get(name).copied().unwrap_or(0)
}
pub fn get_histogram(&self, name: &str) -> Histogram {
let values = self
.histograms
.read()
.unwrap()
.get(name)
.cloned()
.unwrap_or_default();
Histogram::new(values)
}
pub fn get_counter_with_labels(&self, name: &str, _labels: &[(&str, &str)]) -> u64 {
self.get_counter(name)
}
}
#[derive(Debug, Clone)]
pub struct SpanContext {
pub trace_id: String,
pub parent_span_id: String,
pub sampled: bool,
}
impl SpanContext {
pub fn new(trace_id: &str, parent_span_id: &str) -> Self {
Self {
trace_id: trace_id.to_string(),
parent_span_id: parent_span_id.to_string(),
sampled: true,
}
}
}
pub struct Histogram {
values: Vec<f64>,
}
impl Histogram {
pub fn new(values: Vec<f64>) -> Self {
Self { values }
}
pub fn count(&self) -> usize {
self.values.len()
}
pub fn sum(&self) -> f64 {
self.values.iter().sum()
}
pub fn p50(&self) -> f64 {
self.percentile(0.5)
}
pub fn p95(&self) -> f64 {
self.percentile(0.95)
}
pub fn p99(&self) -> f64 {
self.percentile(0.99)
}
fn percentile(&self, p: f64) -> f64 {
if self.values.is_empty() {
return 0.0;
}
let mut sorted = self.values.clone();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
let index = ((self.values.len() as f64 - 1.0) * p) as usize;
sorted[index]
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_span_recorder() {
let recorder = SpanRecorder::new();
let span = Span {
span_id: "span-1".to_string(),
parent_span_id: None,
trace_id: "trace-1".to_string(),
name: "test".to_string(),
attributes: HashMap::new(),
status: "ok".to_string(),
error_message: String::new(),
duration_ms: 100.0,
layer: String::new(),
};
recorder.record(span.clone());
let spans = recorder.spans();
assert_eq!(spans.len(), 1);
assert_eq!(spans[0].span_id, "span-1");
}
#[test]
fn test_histogram() {
let hist = Histogram::new(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
assert_eq!(hist.count(), 5);
assert_eq!(hist.sum(), 15.0);
assert_eq!(hist.p50(), 3.0);
}
#[test]
fn test_metrics_recorder() {
let recorder = MetricsRecorder::new();
recorder.increment_counter("requests", 1);
recorder.increment_counter("requests", 2);
assert_eq!(recorder.get_counter("requests"), 3);
recorder.set_gauge("active_connections", 42);
assert_eq!(recorder.get_gauge("active_connections"), 42);
recorder.record_histogram("latency", 100.0);
recorder.record_histogram("latency", 200.0);
let hist = recorder.get_histogram("latency");
assert_eq!(hist.count(), 2);
}
}