use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::time::SystemTime;
#[derive(Debug, Clone)]
pub struct DistributedTracer {
pub service_name: String,
}
impl DistributedTracer {
pub fn new(service_name: String) -> Self {
Self { service_name }
}
pub fn start_span(&self, operation: &str) -> Span {
Span::new(operation, None)
}
pub fn start_span_with_context(
&self,
operation: &str,
context: &HashMap<String, String>,
) -> Span {
let trace_id: _ = context
.get("trace-id")
.cloned()
.unwrap_or_else(|| generate_trace_id());
let parent_span_id: _ = context.get("span-id").cloned();
Span::new_with_context(operation, trace_id, parent_span_id)
}
pub fn inject_context(&self, span: &Span, headers: &mut HashMap<String, String>) {
headers.insert("trace-id".to_string(), span.trace_id.clone());
headers.insert("span-id".to_string(), span.span_id.clone());
if let Some(ref parent_id) = span.parent_span_id {
headers.insert("parent-span-id".to_string(), parent_id.clone());
}
}
pub fn extract_context(&self, headers: &HashMap<String, String>) -> Option<TraceContext> {
let trace_id: _ = headers.get("trace-id")?.clone();
let span_id: _ = headers.get("span-id")?.clone();
let baggage: _ = HashMap::new();
Some(TraceContext {
trace_id,
span_id,
baggage,
})
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Span {
pub trace_id: String,
pub span_id: String,
pub parent_span_id: Option<String>,
pub operation_name: String,
pub start_time: SystemTime,
pub tags: HashMap<String, String>,
pub logs: Vec<String>,
}
impl Span {
fn new(operation: &str, parent_span_id: Option<String>) -> Self {
Self {
trace_id: generate_trace_id(),
span_id: generate_span_id(),
parent_span_id,
operation_name: operation.to_string(),
start_time: SystemTime::now(),
tags: HashMap::new(),
logs: Vec::new(),
}
}
fn new_with_context(
operation: &str,
trace_id: String,
parent_span_id: Option<String>,
) -> Self {
Self {
trace_id,
span_id: generate_span_id(),
parent_span_id,
operation_name: operation.to_string(),
start_time: SystemTime::now(),
tags: HashMap::new(),
logs: Vec::new(),
}
}
pub fn add_tag(&mut self, key: &str, value: &str) {
self.tags.insert(key.to_string(), value.to_string());
}
pub fn log_event(&mut self, event: &str) {
self.logs.push(event.to_string());
}
pub fn get_duration(&self) -> Duration {
SystemTime::now()
.duration_since(self.start_time)
.unwrap_or_else(|_| Duration::from_secs(0))
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TraceContext {
pub trace_id: String,
pub span_id: String,
pub baggage: HashMap<String, String>,
}
impl TraceContext {
pub fn new(trace_id: String, span_id: String) -> Self {
Self {
trace_id,
span_id,
baggage: HashMap::new(),
}
}
pub fn add_baggage(&mut self, key: &str, value: &str) {
self.baggage.insert(key.to_string(), value.to_string());
}
pub fn get_baggage(&self, key: &str) -> Option<&String> {
self.baggage.get(key)
}
}
fn generate_trace_id() -> String {
let bytes: [u8; 16] = rand::random();
bytes.iter().map(|b| format!("{:02x}", b)).collect()
}
fn generate_span_id() -> String {
let bytes: [u8; 8] = rand::random();
bytes.iter().map(|b| format!("{:02x}", b)).collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tracer_creation() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
assert_eq!(tracer.service_name, "amberjs-service");
}
#[test]
fn test_start_span() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
let span: _ = tracer.start_span("api_request");
assert!(!span.trace_id.is_empty());
assert!(!span.span_id.is_empty());
assert_eq!(span.operation_name, "api_request");
assert_eq!(span.parent_span_id, None);
assert_eq!(span.tags.len(), 0);
assert_eq!(span.logs.len(), 0);
}
#[test]
fn test_inject_context() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
let span: _ = tracer.start_span("database_query");
let mut headers = HashMap::new();
tracer.inject_context(&span, &mut headers);
assert_eq!(headers.get("trace-id"), Some(&span.trace_id));
assert_eq!(headers.get("span-id"), Some(&span.span_id));
assert_eq!(headers.get("parent-span-id"), None);
}
#[test]
fn test_extract_context() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
let mut headers = HashMap::new();
headers.insert("trace-id".to_string(), "test-trace-id".to_string());
headers.insert("span-id".to_string(), "test-span-id".to_string());
let context: _ = tracer.extract_context(&headers).unwrap();
assert_eq!(context.trace_id, "test-trace-id");
assert_eq!(context.span_id, "test-span-id");
}
#[test]
fn test_span_with_context() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
let mut headers = HashMap::new();
headers.insert("trace-id".to_string(), "parent-trace".to_string());
headers.insert("span-id".to_string(), "parent-span".to_string());
let child_span: _ = tracer.start_span_with_context("child_operation", &headers);
assert_eq!(child_span.trace_id, "parent-trace");
assert_eq!(child_span.parent_span_id, Some("parent-span".to_string());
assert_eq!(child_span.operation_name, "child_operation");
}
#[test]
fn test_span_tags() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
let mut span = tracer.start_span("api_request");
span.add_tag("user_id", "12345");
span.add_tag("auth_method", "oauth");
assert_eq!(span.tags.get("user_id"), Some(&"12345".to_string());
assert_eq!(span.tags.get("auth_method"), Some(&"oauth".to_string());
}
#[test]
fn test_span_logs() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
let mut span = tracer.start_span("api_request");
span.log_event("request_received");
span.log_event("response_sent");
assert_eq!(span.logs.len(), 2);
assert!(span.logs.contains(&"request_received".to_string());
assert!(span.logs.contains(&"response_sent".to_string());
}
#[test]
fn test_trace_context_baggage() {
let mut context = TraceContext::new("trace-123".to_string(), "span-456".to_string());
context.add_baggage("user_id", "12345");
context.add_baggage("tenant", "acme");
assert_eq!(context.get_baggage("user_id"), Some(&"12345".to_string());
assert_eq!(context.get_baggage("tenant"), Some(&"acme".to_string());
}
#[test]
fn test_span_duration() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
let span: _ = tracer.start_span("api_request");
std::thread::sleep(Duration::from_millis(10));
let duration: _ = span.get_duration();
assert!(duration.as_millis() >= 10);
}
#[test]
fn test_trace_id_generation() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
let span1: _ = tracer.start_span("op1");
let span2: _ = tracer.start_span("op2");
assert_ne!(span1.span_id, span2.span_id);
assert_eq!(span1.trace_id.len(), 32);
}
#[test]
fn test_span_id_generation() {
let tracer: _ = DistributedTracer::new("amberjs-service".to_string());
let span1: _ = tracer.start_span("op1");
let span2: _ = tracer.start_span("op2");
assert_ne!(span1.span_id, span2.span_id);
assert_eq!(span1.span_id.len(), 16);
}
}
use std::collections::{BTreeMap, HashMap};use std::time::Duration;