1use actix_web::{
14 body::MessageBody,
15 dev::{ServiceRequest, ServiceResponse},
16 Error,
17};
18use async_trait::async_trait;
19use opentelemetry::{global, trace::TraceId, KeyValue};
20use opentelemetry_otlp::WithExportConfig;
21use opentelemetry_sdk::{
22 propagation::TraceContextPropagator, runtime::TokioCurrentThread, trace, Resource,
23};
24use tracing::Span;
25use tracing_actix_web::{DefaultRootSpanBuilder, RootSpanBuilder, TracingLogger};
26use tracing_subscriber::{
27 fmt::{self, format::FmtSpan},
28 layer::SubscriberExt,
29 EnvFilter, Registry,
30};
31
32use std::{borrow::Cow, cell::RefCell};
33
34#[derive(Clone, Default, Debug)]
39pub struct TelemetryConfig {
40 pub app_name: String,
42 pub env: String,
44 pub endpoint_url: Option<String>,
46 pub tracer_id: Option<String>,
48}
49
50#[async_trait]
52pub trait TelemetryInit {
53 async fn init(&self) -> Result<(), Box<dyn std::error::Error>>;
58}
59
60impl TelemetryConfig {
61 pub fn get_trace_id(&self) -> TraceId {
65 use opentelemetry::trace::TraceContextExt as _; use tracing_opentelemetry::OpenTelemetrySpanExt as _; tracing::Span::current()
69 .context()
70 .span()
71 .span_context()
72 .trace_id()
73 }
74}
75
76#[async_trait]
84impl TelemetryInit for TelemetryConfig {
85 async fn init(&self) -> Result<(), Box<dyn std::error::Error>> {
86 let env_filter =
87 EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
88 let sampler = trace::Sampler::AlwaysOn;
89 let resource = Resource::new(vec![KeyValue::new("service.name", self.app_name.clone())]);
90 let trace_config = trace::config()
91 .with_sampler(sampler)
92 .with_resource(resource);
93 global::set_text_map_propagator(TraceContextPropagator::new());
94
95 match &self.endpoint_url {
97 Some(endpoint_url) => {
98 let exporter = opentelemetry_otlp::new_exporter()
99 .tonic()
100 .with_endpoint(endpoint_url);
101 let tracer = opentelemetry_otlp::new_pipeline()
102 .tracing()
103 .with_exporter(exporter)
104 .with_trace_config(trace_config)
105 .install_batch(TokioCurrentThread)?;
106 let telemetry = tracing_opentelemetry::layer().with_tracer(tracer);
107 if self.env == "development" {
108 let logger = fmt::layer().compact();
109 let subscriber = Registry::default()
110 .with(telemetry)
111 .with(logger)
112 .with(env_filter);
113 tracing::subscriber::set_global_default(subscriber)
114 .expect("setting default subscriber failed");
115 } else {
116 let subscriber = Registry::default()
117 .with(telemetry)
118 .with(env_filter)
119 .with(fmt::layer().json().with_span_events(FmtSpan::NONE));
120 tracing::subscriber::set_global_default(subscriber)
121 .expect("setting default subscriber failed");
122 };
123 }
124 None => {
125 if self.env == "development" {
126 let logger = fmt::layer().compact();
127 let subscriber = Registry::default().with(logger).with(env_filter);
128 tracing::subscriber::set_global_default(subscriber)
129 .expect("setting default subscriber failed");
130 } else {
131 let subscriber = Registry::default()
132 .with(fmt::layer().json().with_span_events(FmtSpan::NONE))
133 .with(env_filter);
134 tracing::subscriber::set_global_default(subscriber)
135 .expect("setting default subscriber failed");
136 }
137 }
138 }
139 if let Some(tracer_id) = &self.tracer_id {
140 let name: Cow<'static, str> = tracer_id.to_string().into();
141 global::tracer(name);
142 }
143
144 tracing_log::LogTracer::init()?;
148 Ok(())
149 }
150}
151
152thread_local! {
153 static EXCLUDED_ROUTES: RefCell<Vec<String>> = RefCell::new(Vec::new());
157}
158
159pub struct CustomFilterRootSpanBuilder;
164
165impl CustomFilterRootSpanBuilder {
166 pub fn set_excluded_routes(routes: Vec<String>) {
172 EXCLUDED_ROUTES.with(|excluded| {
173 *excluded.borrow_mut() = routes;
174 });
175 }
176}
177
178impl RootSpanBuilder for CustomFilterRootSpanBuilder {
179 fn on_request_start(request: &ServiceRequest) -> Span {
180 let should_exclude = EXCLUDED_ROUTES
181 .with(|excluded| excluded.borrow().contains(&request.path().to_string()));
182
183 if should_exclude {
184 Span::none()
185 } else {
186 tracing_actix_web::root_span!(level = tracing::Level::INFO, request)
187 }
188 }
189
190 fn on_request_end<B: MessageBody>(span: Span, outcome: &Result<ServiceResponse<B>, Error>) {
191 DefaultRootSpanBuilder::on_request_end(span, outcome);
192 }
193}
194
195pub struct CustomLoggerBuilder {
199 excluded_routes: Vec<String>,
200}
201
202impl CustomLoggerBuilder {
203 pub fn new() -> Self {
205 Self {
206 excluded_routes: Vec::new(),
207 }
208 }
209
210 pub fn exclude(mut self, route: &str) -> Self {
216 self.excluded_routes.push(route.to_string());
217 self
218 }
219
220 pub fn build(self) -> TracingLogger<CustomFilterRootSpanBuilder> {
224 CustomFilterRootSpanBuilder::set_excluded_routes(self.excluded_routes);
226
227 TracingLogger::<CustomFilterRootSpanBuilder>::new()
229 }
230}
231
232impl Default for CustomLoggerBuilder {
233 fn default() -> Self {
234 Self::new()
235 }
236}
237
238pub fn get_tracing_logger() -> CustomLoggerBuilder {
242 CustomLoggerBuilder::new()
243}
244
245#[cfg(test)]
246mod tests {
247 use super::*;
248 use actix_web::test::TestRequest;
249
250 #[test]
251 fn test_telemetry_config_defaults() {
252 let config = TelemetryConfig::default();
253 assert_eq!(config.app_name, "");
254 assert_eq!(config.env, "");
255 assert!(config.endpoint_url.is_none());
256 assert!(config.tracer_id.is_none());
257 }
258
259 #[tokio::test]
260 async fn test_init_with_defaults() {
261 let config = TelemetryConfig::default();
262 let result = config.init().await;
263 assert!(result.is_ok());
264 }
265
266 #[test]
267 fn test_excluded_route() {
268 CustomFilterRootSpanBuilder::set_excluded_routes(vec!["/health/liveness".to_string()]);
269 let req = TestRequest::get().uri("/health/liveness").to_srv_request();
270 let span = CustomFilterRootSpanBuilder::on_request_start(&req);
271 assert!(span.is_none());
272 }
273
274 #[tokio::test]
275 async fn test_non_excluded_route() {
276 fn mock_on_request_start(request: &ServiceRequest) -> bool {
277 let should_exclude = EXCLUDED_ROUTES
278 .with(|excluded| excluded.borrow().contains(&request.path().to_string()));
279 !should_exclude
280 }
281 CustomFilterRootSpanBuilder::set_excluded_routes(vec!["/health/liveness".to_string()]);
282 let req = TestRequest::get().uri("/some/other/route").to_srv_request();
283 let should_log = mock_on_request_start(&req);
284 assert!(should_log);
285 }
286}