mockforge_tracing/
tracer.rs

1//! OpenTelemetry tracer initialization and configuration
2
3use crate::exporter::ExporterType;
4use opentelemetry::global;
5use opentelemetry::KeyValue;
6use opentelemetry_otlp::WithExportConfig;
7use opentelemetry_sdk::trace::Tracer;
8use opentelemetry_sdk::Resource;
9use std::error::Error;
10use std::time::Duration;
11
12/// Tracing configuration
13#[derive(Debug, Clone)]
14pub struct TracingConfig {
15    /// Service name for traces
16    pub service_name: String,
17    /// Exporter type (Jaeger or OTLP)
18    pub exporter_type: ExporterType,
19    /// Jaeger endpoint (e.g., "http://localhost:14268/api/traces")
20    pub jaeger_endpoint: Option<String>,
21    /// OTLP endpoint (e.g., "http://localhost:4317")
22    pub otlp_endpoint: Option<String>,
23    /// Sampling rate (0.0 to 1.0)
24    pub sampling_rate: f64,
25    /// Environment (e.g., "development", "production")
26    pub environment: String,
27    /// Service version
28    pub service_version: Option<String>,
29}
30
31impl Default for TracingConfig {
32    fn default() -> Self {
33        Self {
34            service_name: "mockforge".to_string(),
35            exporter_type: ExporterType::Jaeger,
36            jaeger_endpoint: Some("http://localhost:14268/api/traces".to_string()),
37            otlp_endpoint: None,
38            sampling_rate: 1.0,
39            environment: "development".to_string(),
40            service_version: None,
41        }
42    }
43}
44
45impl TracingConfig {
46    /// Create config for Jaeger exporter
47    pub fn with_jaeger(service_name: String, endpoint: String) -> Self {
48        Self {
49            service_name,
50            exporter_type: ExporterType::Jaeger,
51            jaeger_endpoint: Some(endpoint),
52            otlp_endpoint: None,
53            ..Default::default()
54        }
55    }
56
57    /// Create config for OTLP exporter
58    pub fn with_otlp(service_name: String, endpoint: String) -> Self {
59        Self {
60            service_name,
61            exporter_type: ExporterType::Otlp,
62            jaeger_endpoint: None,
63            otlp_endpoint: Some(endpoint),
64            ..Default::default()
65        }
66    }
67
68    /// Set sampling rate
69    pub fn with_sampling_rate(mut self, rate: f64) -> Self {
70        self.sampling_rate = rate;
71        self
72    }
73
74    /// Set environment
75    pub fn with_environment(mut self, env: String) -> Self {
76        self.environment = env;
77        self
78    }
79
80    /// Set service version
81    pub fn with_service_version(mut self, version: String) -> Self {
82        self.service_version = Some(version);
83        self
84    }
85}
86
87/// Initialize the OpenTelemetry tracer
88pub fn init_tracer(config: TracingConfig) -> Result<Tracer, Box<dyn Error + Send + Sync>> {
89    match config.exporter_type {
90        ExporterType::Jaeger => init_jaeger_tracer(config),
91        ExporterType::Otlp => init_otlp_tracer(config),
92    }
93}
94
95/// Initialize Jaeger tracer
96fn init_jaeger_tracer(config: TracingConfig) -> Result<Tracer, Box<dyn Error + Send + Sync>> {
97    let endpoint = config.jaeger_endpoint.ok_or("Jaeger endpoint not configured")?;
98
99    let tracer = opentelemetry_jaeger::new_agent_pipeline()
100        .with_service_name(&config.service_name)
101        .with_endpoint(&endpoint)
102        .install_batch(opentelemetry_sdk::runtime::Tokio)?;
103
104    Ok(tracer)
105}
106
107/// Initialize OTLP tracer
108fn init_otlp_tracer(config: TracingConfig) -> Result<Tracer, Box<dyn Error + Send + Sync>> {
109    let endpoint = config.otlp_endpoint.ok_or("OTLP endpoint not configured")?;
110
111    // Build resource attributes
112    let mut resource_attrs = vec![
113        KeyValue::new("service.name", config.service_name.clone()),
114        KeyValue::new("deployment.environment", config.environment.clone()),
115    ];
116
117    if let Some(version) = config.service_version {
118        resource_attrs.push(KeyValue::new("service.version", version));
119    }
120
121    let resource = Resource::new(resource_attrs);
122
123    // Create OTLP exporter with gRPC protocol
124    let exporter = opentelemetry_otlp::new_exporter()
125        .tonic()
126        .with_endpoint(&endpoint)
127        .with_timeout(Duration::from_secs(10));
128
129    // Build tracer pipeline
130    let tracer = opentelemetry_otlp::new_pipeline()
131        .tracing()
132        .with_exporter(exporter)
133        .with_trace_config(
134            opentelemetry_sdk::trace::Config::default()
135                .with_resource(resource)
136                .with_sampler(opentelemetry_sdk::trace::Sampler::TraceIdRatioBased(
137                    config.sampling_rate,
138                )),
139        )
140        .install_batch(opentelemetry_sdk::runtime::Tokio)?;
141
142    Ok(tracer)
143}
144
145/// Shutdown the tracer and flush pending spans
146pub fn shutdown_tracer() {
147    global::shutdown_tracer_provider();
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    #[test]
155    fn test_default_config() {
156        let config = TracingConfig::default();
157        assert_eq!(config.service_name, "mockforge");
158        assert_eq!(config.sampling_rate, 1.0);
159        assert_eq!(config.environment, "development");
160        assert_eq!(config.exporter_type, ExporterType::Jaeger);
161        assert!(config.jaeger_endpoint.is_some());
162        assert!(config.otlp_endpoint.is_none());
163    }
164
165    #[test]
166    fn test_jaeger_config() {
167        let config = TracingConfig::with_jaeger(
168            "test-service".to_string(),
169            "http://custom:14268/api/traces".to_string(),
170        )
171        .with_sampling_rate(0.5)
172        .with_environment("staging".to_string())
173        .with_service_version("1.0.0".to_string());
174
175        assert_eq!(config.service_name, "test-service");
176        assert_eq!(config.exporter_type, ExporterType::Jaeger);
177        assert_eq!(config.jaeger_endpoint, Some("http://custom:14268/api/traces".to_string()));
178        assert_eq!(config.sampling_rate, 0.5);
179        assert_eq!(config.environment, "staging");
180        assert_eq!(config.service_version, Some("1.0.0".to_string()));
181    }
182
183    #[test]
184    fn test_otlp_config() {
185        let config = TracingConfig::with_otlp(
186            "test-service".to_string(),
187            "http://otel-collector:4317".to_string(),
188        )
189        .with_sampling_rate(0.8)
190        .with_environment("production".to_string())
191        .with_service_version("2.0.0".to_string());
192
193        assert_eq!(config.service_name, "test-service");
194        assert_eq!(config.exporter_type, ExporterType::Otlp);
195        assert_eq!(config.otlp_endpoint, Some("http://otel-collector:4317".to_string()));
196        assert!(config.jaeger_endpoint.is_none());
197        assert_eq!(config.sampling_rate, 0.8);
198        assert_eq!(config.environment, "production");
199        assert_eq!(config.service_version, Some("2.0.0".to_string()));
200    }
201}
mockforge_tracing/tracer.rs

mockforge_tracing/
tracer.rs
