mockforge_tracing/
tracer.rs1use 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#[derive(Debug, Clone)]
13pub struct TracingConfig {
14 pub service_name: String,
16 pub exporter_type: ExporterType,
18 pub jaeger_endpoint: Option<String>,
20 pub otlp_endpoint: Option<String>,
22 pub sampling_rate: f64,
24 pub environment: String,
26 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 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 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 pub fn with_sampling_rate(mut self, rate: f64) -> Self {
69 self.sampling_rate = rate;
70 self
71 }
72
73 pub fn with_environment(mut self, env: String) -> Self {
75 self.environment = env;
76 self
77 }
78
79 pub fn with_service_version(mut self, version: String) -> Self {
81 self.service_version = Some(version);
82 self
83 }
84}
85
86pub 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
96fn 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 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 let tracer = opentelemetry::global::tracer("mockforge");
110 Ok(tracer)
111}
112
113fn 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 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 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 let exporter = exporter_builder.build_span_exporter()?;
139
140 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 opentelemetry::global::set_tracer_provider(tracer_provider.clone());
154
155 let tracer = opentelemetry::global::tracer("mockforge");
157
158 Ok(tracer)
159}
160
161pub 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}