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