use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use tracing::{debug, info, instrument, warn};
use std::time::{Duration, Instant};
use std::time::SystemTime;
pub type TraceId = String;
pub type SpanId = String;
#[derive(Debug, Clone, PartialEq)]
pub enum TraceEventType {
RequestStart,
RequestEnd,
TaskStart,
TaskEnd,
NetworkCall,
DatabaseQuery,
CacheAccess,
JitCompilation,
GcEvent,
Custom(&'static str),
}
#[derive(Debug, Clone)]
pub struct TraceEvent {
pub trace_id: TraceId,
pub span_id: SpanId,
pub parent_span_id: Option<SpanId>,
pub event_type: TraceEventType,
pub timestamp: u64, pub duration: Option<Duration>,
pub node_id: String,
pub service_name: String,
pub operation_name: String,
pub tags: HashMap<String, String>,
pub baggage: HashMap<String, String>,
}
#[derive(Debug, Clone)]
pub struct Span {
pub trace_id: TraceId,
pub span_id: SpanId,
pub parent_span_id: Option<SpanId>,
pub operation_name: String,
pub service_name: String,
pub start_time: Instant,
pub end_time: Option<Instant>,
pub tags: HashMap<String, String>,
pub events: Vec<TraceEvent>,
pub baggage: HashMap<String, String>,
}
impl Span {
pub fn new(
trace_id: TraceId,
span_id: SpanId,
parent_span_id: Option<SpanId>,
operation_name: String,
service_name: String,
) -> Self {
Self {
trace_id,
span_id,
parent_span_id,
operation_name,
service_name,
start_time: Instant::now(),
end_time: None,
tags: HashMap::new(),
events: Vec::new(),
baggage: HashMap::new(),
}
}
pub fn finish(&mut self) {
self.end_time = Some(Instant::now());
}
pub fn duration(&self) -> Option<Duration> {
self.end_time.map(|end| end.duration_since(self.start_time))
}
pub fn with_tag(mut self, key: String, value: String) -> Self {
self.tags.insert(key, value);
self
}
pub fn with_baggage(mut self, key: String, value: String) -> Self {
self.baggage.insert(key, value);
self
}
}
#[derive(Debug, Clone)]
pub struct Trace {
pub trace_id: TraceId,
pub root_span: Span,
pub spans: HashMap<SpanId, Span>,
pub start_time: Instant,
pub end_time: Option<Instant>,
pub duration: Option<Duration>,
}
impl Trace {
pub fn new(trace_id: TraceId, root_span: Span) -> Self {
let mut spans = HashMap::new();
spans.insert(root_span.span_id.clone(), root_span.clone());
Self {
trace_id,
root_span,
spans,
start_time: Instant::now(),
end_time: None,
duration: None,
}
}
pub fn add_span(&mut self, span: Span) {
self.spans.insert(span.span_id.clone(), span);
}
pub fn finish(&mut self) {
self.end_time = Some(Instant::now());
self.duration = self.end_time.map(|end| end.duration_since(self.start_time));
for span in self.spans.values_mut() {
if span.end_time.is_none() {
span.finish();
}
}
}
pub fn duration(&self) -> Option<Duration> {
self.duration
}
pub fn span_count(&self) -> usize {
self.spans.len()
}
pub fn get_span(&self, span_id: &SpanId) -> Option<&Span> {
self.spans.get(span_id)
}
}
#[derive(Debug, Clone)]
pub struct PerformanceStats {
pub total_traces: u64,
pub total_spans: u64,
pub average_trace_duration: Duration,
pub p50_trace_duration: Duration,
pub p90_trace_duration: Duration,
pub p99_trace_duration: Duration,
pub slowest_operations: Vec<(String, Duration)>,
pub operation_counts: HashMap<String, u64>,
}
#[derive(Debug, Clone)]
pub struct TracingConfig {
pub max_traces: usize,
pub max_spans_per_trace: usize,
pub trace_retention: Duration,
pub enable_sampling: bool,
pub sampling_rate: f64, pub enable_performance_analysis: bool,
}
#[derive(Clone, Debug)]
pub struct DistributedTracer {
config: TracingConfig,
node_manager: Arc<NodeManager>,
task_executor: Arc<TaskExecutor>,
active_traces: Arc<RwLock<HashMap<TraceId, Trace>>>,
completed_traces: Arc<RwLock<HashMap<TraceId, Trace>>>,
performance_stats: Arc<RwLock<PerformanceStats>>,
}
impl DistributedTracer {
pub fn new(
config: TracingConfig,
node_manager: Arc<NodeManager>,
task_executor: Arc<TaskExecutor>,
) -> Self {
let initial_stats: _ = PerformanceStats {
total_traces: 0,
total_spans: 0,
average_trace_duration: Duration::from_millis(0),
p50_trace_duration: Duration::from_millis(0),
p90_trace_duration: Duration::from_millis(0),
p99_trace_duration: Duration::from_millis(0),
slowest_operations: Vec::new(),
operation_counts: HashMap::new(),
};
Self {
config,
node_manager,
task_executor,
active_traces: Arc::new(Mutex::new(HashMap::new())),
completed_traces: Arc::new(Mutex::new(HashMap::new())),
performance_stats: Arc::new(Mutex::new(initial_stats)),
}
}
pub async fn start(&self) -> Result<(), String> {
info!("Starting distributed tracer...");
let tracer: _ = self.clone();
tokio::spawn(async move {
let mut interval_timer = interval(Duration::from_secs(5));
loop {
interval_timer.tick().await;
if let Err(e) = tracer.cleanup_old_traces().await {
warn!("Failed to cleanup old traces: {}", e);
}
if let Err(e) = tracer.update_performance_stats().await {
warn!("Failed to update performance stats: {}", e);
}
}
});
info!("Distributed tracer started");
Ok(())
}
pub async fn start_trace(&self, operation_name: String, service_name: String) -> TraceId {
let trace_id: _ = self.generate_trace_id();
let root_span_id: _ = self.generate_span_id();
let root_span: _ = Span::new(
trace_id.clone(),
root_span_id,
None,
operation_name,
service_name,
);
let trace: _ = Trace::new(trace_id.clone(), root_span);
let mut traces = self.active_traces.write().await;
traces.insert(trace_id.clone(), trace);
trace_id
}
pub async fn finish_trace(&self, trace_id: &TraceId) -> Option<Trace> {
let mut traces = self.active_traces.write().await;
let trace: _ = traces.remove(trace_id)?;
let mut completed = self.completed_traces.write().await;
let mut trace_clone = trace.clone();
trace_clone.finish();
completed.insert(trace_id.clone(), trace_clone.clone());
self.update_trace_stats(&trace_clone).await;
Some(trace_clone)
}
pub async fn start_span(
&self,
trace_id: &TraceId,
operation_name: String,
service_name: String,
) -> Option<SpanId> {
let span_id: _ = self.generate_span_id();
let parent_span_id: _ = self.get_current_span_id(trace_id).await;
let span: _ = Span::new(
trace_id.clone(),
span_id.clone(),
parent_span_id,
operation_name,
service_name,
);
let mut traces = self.active_traces.write().await;
if let Some(trace) = traces.get_mut(trace_id) {
trace.add_span(span);
Some(span_id)
} else {
None
}
}
pub async fn finish_span(&self, trace_id: &TraceId, span_id: &SpanId) -> Option<()> {
let mut traces = self.active_traces.write().await;
if let Some(trace) = traces.get_mut(trace_id) {
if let Some(span) = trace.spans.get_mut(span_id) {
span.finish();
Some(())
} else {
None
}
} else {
None
}
}
#[instrument(skip(self))]
pub async fn add_event(
&self,
trace_id: &TraceId,
span_id: &SpanId,
event_type: TraceEventType,
node_id: String,
tags: HashMap<String, String>,
) {
let event: _ = TraceEvent {
trace_id: trace_id.clone(),
span_id: span_id.clone(),
parent_span_id: self.get_parent_span_id(trace_id, span_id).await,
event_type,
timestamp: std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(),
duration: None,
node_id,
service_name: "unknown".to_string(),
operation_name: "unknown".to_string(),
tags,
baggage: HashMap::new(),
};
let mut traces = self.active_traces.write().await;
if let Some(trace) = traces.get_mut(trace_id) {
if let Some(span) = trace.spans.get_mut(span_id) {
span.events.push(event);
}
}
}
pub async fn get_trace(&self, trace_id: &TraceId) -> Option<Trace> {
let traces: _ = self.active_traces.read().await;
traces.get(trace_id).cloned()
}
pub async fn get_all_traces(&self) -> Vec<Trace> {
let traces: _ = self.active_traces.read().await;
traces.values().cloned().collect()
}
pub async fn get_completed_traces(&self) -> Vec<Trace> {
let traces: _ = self.completed_traces.read().await;
traces.values().cloned().collect()
}
pub async fn get_performance_stats(&self) -> PerformanceStats {
let stats: _ = self.performance_stats.read().await;
stats.clone()
}
async fn cleanup_old_traces(&self) -> Result<(), String> {
let cutoff: _ = Instant::now() - self.config.trace_retention;
let mut completed = self.completed_traces.write().await;
completed.retain(|_, trace| {
trace.start_time > cutoff
});
Ok(())
}
async fn update_performance_stats(&self) -> Result<(), String> {
let completed: _ = self.completed_traces.read().await;
let traces: Vec<&Trace> = completed.values().collect();
if traces.is_empty() {
return Ok(());
}
let total_traces: _ = traces.len() as u64;
let total_spans: u64 = traces.iter().map(|t| t.span_count() as u64).sum();
let mut durations: Vec<Duration> = traces.iter()
.filter_map(|t| t.duration())
.collect();
durations.sort();
let average_duration: _ = if !durations.is_empty() {
durations.iter().sum::<Duration>() / durations.len() as u32
} else {
Duration::from_millis(0)
};
let p50_duration: _ = if durations.len() > 1 {
durations[durations.len() * 50 / 100]
} else {
Duration::from_millis(0)
};
let p90_duration: _ = if durations.len() > 1 {
durations[durations.len() * 90 / 100]
} else {
Duration::from_millis(0)
};
let p99_duration: _ = if durations.len() > 1 {
durations[durations.len() * 99 / 100]
} else {
Duration::from_millis(0)
};
let mut slowest_operations = Vec::new();
let mut operation_counts = HashMap::new();
for trace in traces {
for span in trace.spans.values() {
if let Some(duration) = span.duration() {
slowest_operations.push((span.operation_name.clone(), duration));
*operation_counts.entry(span.operation_name.clone()).or_insert(0) += 1;
}
}
}
slowest_operations.sort_by(|a, b| b.1.cmp(&a.1));
slowest_operations.truncate(10);
let stats: _ = PerformanceStats {
total_traces,
total_spans,
average_trace_duration: average_duration,
p50_trace_duration: p50_duration,
p90_trace_duration: p90_duration,
p99_trace_duration: p99_duration,
slowest_operations,
operation_counts,
};
let mut current_stats = self.performance_stats.write().await;
*current_stats = stats;
Ok(())
}
async fn update_trace_stats(&self, trace: &Trace) {
debug!("Updated stats for trace: {}", trace.trace_id);
}
fn generate_trace_id(&self) -> TraceId {
format!("trace-{:x}", rand::random::<u128>())
}
fn generate_span_id(&self) -> SpanId {
format!("span-{:x}", rand::random::<u64>())
}
async fn get_current_span_id(&self, _trace_id: &TraceId) -> Option<SpanId> {
None
}
async fn get_parent_span_id(&self, trace_id: &TraceId, span_id: &SpanId) -> Option<SpanId> {
let traces: _ = self.active_traces.read().await;
if let Some(trace) = traces.get(trace_id) {
if let Some(span) = trace.get_span(span_id) {
return span.parent_span_id.clone();
}
}
None
}
}
#[derive(Clone, Debug)]
pub struct TraceContext {
pub trace_id: TraceId,
pub span_id: SpanId,
pub baggage: HashMap<String, String>,
}
impl TraceContext {
pub fn new(trace_id: TraceId, span_id: SpanId) -> Self {
Self {
trace_id,
span_id,
baggage: HashMap::new(),
}
}
pub fn with_baggage(mut self, key: String, value: String) -> Self {
self.baggage.insert(key, value);
self
}
}