use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct TracingConfig {
service_name: String,
service_version: String,
console_exporter: bool,
otlp_endpoint: Option<String>,
jaeger_endpoint: Option<String>,
sampling_rate: f64,
max_attributes_per_span: u32,
#[allow(dead_code)]
max_events_per_span: u32,
batch_timeout: Duration,
#[allow(dead_code)]
max_batch_size: usize,
}
impl Default for TracingConfig {
#[inline]
fn default() -> Self {
Self {
service_name: "chie-core".to_string(),
service_version: env!("CARGO_PKG_VERSION").to_string(),
console_exporter: false,
otlp_endpoint: None,
jaeger_endpoint: None,
sampling_rate: 1.0,
max_attributes_per_span: 128,
max_events_per_span: 128,
batch_timeout: Duration::from_secs(5),
max_batch_size: 512,
}
}
}
impl TracingConfig {
#[must_use]
#[inline]
pub fn new() -> Self {
Self::default()
}
#[must_use]
#[inline]
pub fn with_service_name(mut self, name: impl Into<String>) -> Self {
self.service_name = name.into();
self
}
#[must_use]
#[inline]
pub fn with_service_version(mut self, version: impl Into<String>) -> Self {
self.service_version = version.into();
self
}
#[must_use]
#[inline]
pub fn with_console_exporter(mut self, enabled: bool) -> Self {
self.console_exporter = enabled;
self
}
#[must_use]
#[inline]
pub fn with_otlp_endpoint(mut self, endpoint: impl Into<String>) -> Self {
self.otlp_endpoint = Some(endpoint.into());
self
}
#[must_use]
#[inline]
pub fn with_jaeger_endpoint(mut self, endpoint: impl Into<String>) -> Self {
self.jaeger_endpoint = Some(endpoint.into());
self
}
#[must_use]
#[inline]
pub fn with_sampling_rate(mut self, rate: f64) -> Self {
self.sampling_rate = rate.clamp(0.0, 1.0);
self
}
#[must_use]
#[inline]
pub fn with_max_attributes_per_span(mut self, max: u32) -> Self {
self.max_attributes_per_span = max;
self
}
#[must_use]
#[inline]
pub fn with_batch_timeout(mut self, timeout: Duration) -> Self {
self.batch_timeout = timeout;
self
}
#[must_use]
#[inline]
pub fn service_name(&self) -> &str {
&self.service_name
}
#[must_use]
#[inline]
pub fn service_version(&self) -> &str {
&self.service_version
}
#[must_use]
#[inline]
pub const fn sampling_rate(&self) -> f64 {
self.sampling_rate
}
#[must_use]
#[inline]
pub const fn console_exporter_enabled(&self) -> bool {
self.console_exporter
}
#[must_use]
#[inline]
pub fn otlp_endpoint(&self) -> Option<&str> {
self.otlp_endpoint.as_deref()
}
#[must_use]
#[inline]
pub fn jaeger_endpoint(&self) -> Option<&str> {
self.jaeger_endpoint.as_deref()
}
}
pub struct TracingManager {
config: TracingConfig,
initialized: bool,
}
impl TracingManager {
pub fn new(config: TracingConfig) -> Result<Self, TracingError> {
let mut manager = Self {
config,
initialized: false,
};
manager.initialize()?;
Ok(manager)
}
fn initialize(&mut self) -> Result<(), TracingError> {
if self.initialized {
return Err(TracingError::AlreadyInitialized);
}
self.initialized = true;
Ok(())
}
pub fn shutdown(self) -> Result<(), TracingError> {
if !self.initialized {
return Err(TracingError::NotInitialized);
}
Ok(())
}
#[must_use]
#[inline]
pub const fn is_initialized(&self) -> bool {
self.initialized
}
#[must_use]
#[inline]
pub const fn config(&self) -> &TracingConfig {
&self.config
}
}
#[derive(Debug, Clone, thiserror::Error)]
pub enum TracingError {
#[error("Tracing is already initialized")]
AlreadyInitialized,
#[error("Tracing is not initialized")]
NotInitialized,
#[error("Configuration error: {0}")]
ConfigError(String),
#[error("Export error: {0}")]
ExportError(String),
}
#[derive(Debug)]
pub struct Span {
name: String,
start_time: std::time::Instant,
attributes: HashMap<String, String>,
events: Vec<SpanEvent>,
}
impl Span {
#[must_use]
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
start_time: std::time::Instant::now(),
attributes: HashMap::new(),
events: Vec::new(),
}
}
#[inline]
pub fn set_attribute(&mut self, key: impl Into<String>, value: impl Into<String>) {
self.attributes.insert(key.into(), value.into());
}
#[inline]
pub fn record_event(&mut self, name: impl Into<String>) {
self.events.push(SpanEvent {
name: name.into(),
timestamp: std::time::Instant::now(),
attributes: HashMap::new(),
});
}
#[inline]
pub fn record_event_with_attributes(
&mut self,
name: impl Into<String>,
attributes: HashMap<String, String>,
) {
self.events.push(SpanEvent {
name: name.into(),
timestamp: std::time::Instant::now(),
attributes,
});
}
#[must_use]
#[inline]
pub fn finish(self) -> Duration {
self.start_time.elapsed()
}
#[must_use]
#[inline]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
#[inline]
pub fn attributes(&self) -> &HashMap<String, String> {
&self.attributes
}
#[must_use]
#[inline]
pub fn events(&self) -> &[SpanEvent] {
&self.events
}
#[must_use]
#[inline]
pub fn elapsed(&self) -> Duration {
self.start_time.elapsed()
}
}
#[derive(Debug, Clone)]
pub struct SpanEvent {
name: String,
#[allow(dead_code)]
timestamp: std::time::Instant,
attributes: HashMap<String, String>,
}
impl SpanEvent {
#[must_use]
#[inline]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
#[inline]
pub fn attributes(&self) -> &HashMap<String, String> {
&self.attributes
}
}
pub struct SpanGuard {
span: Option<Span>,
}
impl SpanGuard {
#[must_use]
#[inline]
pub fn new(span: Span) -> Self {
Self { span: Some(span) }
}
#[must_use]
#[inline]
pub fn span_mut(&mut self) -> Option<&mut Span> {
self.span.as_mut()
}
}
impl Drop for SpanGuard {
fn drop(&mut self) {
if let Some(span) = self.span.take() {
let _duration = span.finish();
}
}
}
#[must_use]
#[inline]
pub fn span_scope(name: impl Into<String>) -> SpanGuard {
SpanGuard::new(Span::new(name))
}
#[must_use]
#[inline]
pub fn span_with_attributes(
name: impl Into<String>,
attributes: HashMap<String, String>,
) -> SpanGuard {
let mut span = Span::new(name);
for (k, v) in attributes {
span.set_attribute(k, v);
}
SpanGuard::new(span)
}
#[derive(Debug, Clone)]
pub struct TraceContext {
trace_id: String,
span_id: String,
trace_flags: u8,
}
impl TraceContext {
#[must_use]
pub fn new(trace_id: String, span_id: String, trace_flags: u8) -> Self {
Self {
trace_id,
span_id,
trace_flags,
}
}
#[must_use]
#[inline]
pub fn trace_id(&self) -> &str {
&self.trace_id
}
#[must_use]
#[inline]
pub fn span_id(&self) -> &str {
&self.span_id
}
#[must_use]
#[inline]
pub const fn trace_flags(&self) -> u8 {
self.trace_flags
}
#[must_use]
#[inline]
pub const fn is_sampled(&self) -> bool {
self.trace_flags & 0x01 != 0
}
#[must_use]
pub fn to_traceparent(&self) -> String {
format!(
"00-{}-{}-{:02x}",
self.trace_id, self.span_id, self.trace_flags
)
}
pub fn from_traceparent(traceparent: &str) -> Result<Self, TracingError> {
let parts: Vec<&str> = traceparent.split('-').collect();
if parts.len() != 4 || parts[0] != "00" {
return Err(TracingError::ConfigError(
"Invalid traceparent format".to_string(),
));
}
let trace_flags = u8::from_str_radix(parts[3], 16)
.map_err(|_| TracingError::ConfigError("Invalid trace flags".to_string()))?;
Ok(Self {
trace_id: parts[1].to_string(),
span_id: parts[2].to_string(),
trace_flags,
})
}
}
#[derive(Debug, Clone, Default)]
pub struct TracingStats {
pub total_spans: u64,
pub exported_spans: u64,
pub dropped_spans: u64,
pub total_events: u64,
pub total_attributes: u64,
}
impl TracingStats {
#[must_use]
#[inline]
pub const fn new() -> Self {
Self {
total_spans: 0,
exported_spans: 0,
dropped_spans: 0,
total_events: 0,
total_attributes: 0,
}
}
#[inline]
pub fn record_span_created(&mut self) {
self.total_spans += 1;
}
#[inline]
pub fn record_span_exported(&mut self) {
self.exported_spans += 1;
}
#[inline]
pub fn record_span_dropped(&mut self) {
self.dropped_spans += 1;
}
#[inline]
pub fn record_event(&mut self) {
self.total_events += 1;
}
#[inline]
pub fn record_attribute(&mut self) {
self.total_attributes += 1;
}
#[must_use]
#[inline]
pub fn export_rate(&self) -> f64 {
if self.total_spans == 0 {
0.0
} else {
self.exported_spans as f64 / self.total_spans as f64
}
}
#[must_use]
#[inline]
pub fn drop_rate(&self) -> f64 {
if self.total_spans == 0 {
0.0
} else {
self.dropped_spans as f64 / self.total_spans as f64
}
}
}
static TRACING_STATS: std::sync::OnceLock<Arc<std::sync::RwLock<TracingStats>>> =
std::sync::OnceLock::new();
#[must_use]
pub fn get_tracing_stats() -> TracingStats {
TRACING_STATS
.get_or_init(|| Arc::new(std::sync::RwLock::new(TracingStats::new())))
.read()
.unwrap()
.clone()
}
pub fn reset_tracing_stats() {
let stats = TRACING_STATS.get_or_init(|| Arc::new(std::sync::RwLock::new(TracingStats::new())));
let mut stats_lock = stats.write().unwrap();
*stats_lock = TracingStats::new();
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tracing_config_default() {
let config = TracingConfig::default();
assert_eq!(config.service_name(), "chie-core");
assert_eq!(config.sampling_rate(), 1.0);
assert!(!config.console_exporter_enabled());
assert!(config.otlp_endpoint().is_none());
assert!(config.jaeger_endpoint().is_none());
}
#[test]
fn test_tracing_config_builder() {
let config = TracingConfig::new()
.with_service_name("test-service")
.with_service_version("1.0.0")
.with_console_exporter(true)
.with_sampling_rate(0.5)
.with_otlp_endpoint("http://localhost:4317")
.with_jaeger_endpoint("http://localhost:14268/api/traces");
assert_eq!(config.service_name(), "test-service");
assert_eq!(config.service_version(), "1.0.0");
assert!(config.console_exporter_enabled());
assert_eq!(config.sampling_rate(), 0.5);
assert_eq!(config.otlp_endpoint(), Some("http://localhost:4317"));
assert_eq!(
config.jaeger_endpoint(),
Some("http://localhost:14268/api/traces")
);
}
#[test]
fn test_sampling_rate_clamping() {
let config1 = TracingConfig::new().with_sampling_rate(-0.5);
assert_eq!(config1.sampling_rate(), 0.0);
let config2 = TracingConfig::new().with_sampling_rate(1.5);
assert_eq!(config2.sampling_rate(), 1.0);
}
#[test]
fn test_tracing_manager_initialization() {
let config = TracingConfig::default();
let manager = TracingManager::new(config).unwrap();
assert!(manager.is_initialized());
}
#[test]
fn test_tracing_manager_shutdown() {
let config = TracingConfig::default();
let manager = TracingManager::new(config).unwrap();
assert!(manager.shutdown().is_ok());
}
#[test]
fn test_span_creation() {
let span = Span::new("test_operation");
assert_eq!(span.name(), "test_operation");
assert!(span.attributes().is_empty());
assert!(span.events().is_empty());
}
#[test]
fn test_span_attributes() {
let mut span = Span::new("test");
span.set_attribute("key1", "value1");
span.set_attribute("key2", "value2");
assert_eq!(span.attributes().len(), 2);
assert_eq!(span.attributes().get("key1"), Some(&"value1".to_string()));
assert_eq!(span.attributes().get("key2"), Some(&"value2".to_string()));
}
#[test]
fn test_span_events() {
let mut span = Span::new("test");
span.record_event("event1");
span.record_event("event2");
assert_eq!(span.events().len(), 2);
assert_eq!(span.events()[0].name(), "event1");
assert_eq!(span.events()[1].name(), "event2");
}
#[test]
fn test_span_event_with_attributes() {
let mut span = Span::new("test");
let mut attrs = HashMap::new();
attrs.insert("error".to_string(), "true".to_string());
span.record_event_with_attributes("error_occurred", attrs);
assert_eq!(span.events().len(), 1);
assert_eq!(span.events()[0].name(), "error_occurred");
assert_eq!(
span.events()[0].attributes().get("error"),
Some(&"true".to_string())
);
}
#[test]
fn test_span_duration() {
let span = Span::new("test");
std::thread::sleep(Duration::from_millis(10));
let duration = span.finish();
assert!(duration >= Duration::from_millis(10));
}
#[test]
fn test_span_guard() {
let span = Span::new("test");
let guard = SpanGuard::new(span);
drop(guard);
}
#[test]
fn test_span_scope() {
{
let _guard = span_scope("scoped_operation");
std::thread::sleep(Duration::from_millis(5));
}
}
#[test]
fn test_trace_context_creation() {
let ctx = TraceContext::new(
"0123456789abcdef0123456789abcdef".to_string(),
"0123456789abcdef".to_string(),
1,
);
assert_eq!(ctx.trace_id(), "0123456789abcdef0123456789abcdef");
assert_eq!(ctx.span_id(), "0123456789abcdef");
assert_eq!(ctx.trace_flags(), 1);
assert!(ctx.is_sampled());
}
#[test]
fn test_trace_context_traceparent() {
let ctx = TraceContext::new(
"0123456789abcdef0123456789abcdef".to_string(),
"0123456789abcdef".to_string(),
1,
);
let traceparent = ctx.to_traceparent();
assert_eq!(
traceparent,
"00-0123456789abcdef0123456789abcdef-0123456789abcdef-01"
);
let parsed = TraceContext::from_traceparent(&traceparent).unwrap();
assert_eq!(parsed.trace_id(), ctx.trace_id());
assert_eq!(parsed.span_id(), ctx.span_id());
assert_eq!(parsed.trace_flags(), ctx.trace_flags());
}
#[test]
fn test_trace_context_invalid_traceparent() {
let result = TraceContext::from_traceparent("invalid");
assert!(result.is_err());
}
#[test]
fn test_tracing_stats() {
let mut stats = TracingStats::new();
assert_eq!(stats.total_spans, 0);
assert_eq!(stats.exported_spans, 0);
stats.record_span_created();
stats.record_span_created();
stats.record_span_exported();
assert_eq!(stats.total_spans, 2);
assert_eq!(stats.exported_spans, 1);
assert_eq!(stats.export_rate(), 0.5);
}
#[test]
fn test_tracing_stats_rates() {
let mut stats = TracingStats::new();
stats.record_span_created();
stats.record_span_created();
stats.record_span_created();
stats.record_span_exported();
stats.record_span_dropped();
assert_eq!(stats.export_rate(), 1.0 / 3.0);
assert_eq!(stats.drop_rate(), 1.0 / 3.0);
}
#[test]
fn test_global_stats() {
reset_tracing_stats();
let stats = get_tracing_stats();
assert_eq!(stats.total_spans, 0);
}
}