tembo_telemetry/
lib.rs

1//! `tembo-telemetry` is a crate that provides logging and telemetry exporters for Tembo.io applications.
2//! It integrates with the OpenTelemetry ecosystem to provide detailed traces.
3//!
4//! # Features
5//! - Configurable telemetry setup via `TelemetryConfig`.
6//! - Integration with the OpenTelemetry and tracing ecosystems.
7//! - Out-of-the-box support for OTLP exporters.
8//! - Environment-specific logger configurations.
9//!
10//! # Usage
11//! Refer to the `TelemetryConfig` and `TelemetryInit` traits for setting up and initializing telemetry.
12
13use 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/// Configuration for telemetry setup.
35///
36/// This struct provides fields to set up OpenTelemetry exporters, specify the application name,
37/// environment, endpoint URL, and an optional tracer ID.
38#[derive(Clone, Default, Debug)]
39pub struct TelemetryConfig {
40    /// Name of the application.
41    pub app_name: String,
42    /// Specifies the environment (e.g., "development" or "production").
43    pub env: String,
44    /// Optional URL for the OTLP exporter.
45    pub endpoint_url: Option<String>,
46    /// Optional tracer ID.
47    pub tracer_id: Option<String>,
48}
49
50/// Trait to initialize telemetry based on the provided configuration.
51#[async_trait]
52pub trait TelemetryInit {
53    /// Initializes telemetry based on the configuration.
54    ///
55    /// This method sets up the global tracer provider, OTLP exporter (if specified),
56    /// and logger based on the environment.
57    async fn init(&self) -> Result<(), Box<dyn std::error::Error>>;
58}
59
60impl TelemetryConfig {
61    /// Retrieves the current trace ID.
62    ///
63    /// This method fetches the trace ID from the current span context.
64    pub fn get_trace_id(&self) -> TraceId {
65        use opentelemetry::trace::TraceContextExt as _; // opentelemetry::Context -> opentelemetry::trace::Span
66        use tracing_opentelemetry::OpenTelemetrySpanExt as _; // tracing::Span to opentelemetry::Context
67
68        tracing::Span::current()
69            .context()
70            .span()
71            .span_context()
72            .trace_id()
73    }
74}
75
76/// Initializes telemetry based on the provided configuration.
77///
78/// This method will:
79/// - Set the global text map propagator to `TraceContextPropagator`.
80/// - Check for an OTLP endpoint and set up the OTLP exporter if present.
81/// - Configure a logger based on the environment (`development` or other).
82/// - Optionally, set a global tracer if `tracer_id` is provided.
83#[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        // Check if OPENTELEMERTY_OTLP_ENDPOINT is set, if not enable standard logger
96        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        // Setup bridge between tracing crate and the log crate.  If someone
145        // uses this crate, then if they use the log crate, they will get
146        // the logs printed into the tracing session.
147        tracing_log::LogTracer::init()?;
148        Ok(())
149    }
150}
151
152thread_local! {
153    /// Thread-local storage for excluded routes.
154    ///
155    /// Contains a list of routes (endpoints) that should not be logged.
156    static EXCLUDED_ROUTES: RefCell<Vec<String>> = RefCell::new(Vec::new());
157}
158
159/// Custom root span builder that allows for filtering out specific routes.
160///
161/// This builder will check if a request's path is in the list of excluded routes,
162/// and if so, it won't log that request.
163pub struct CustomFilterRootSpanBuilder;
164
165impl CustomFilterRootSpanBuilder {
166    /// Sets the routes to be excluded from logging.
167    ///
168    /// # Arguments
169    ///
170    /// * `routes` - A list of route paths to exclude.
171    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
195/// Builder for creating a custom logging middleware.
196///
197/// This builder provides methods to specify which routes to exclude from logging.
198pub struct CustomLoggerBuilder {
199    excluded_routes: Vec<String>,
200}
201
202impl CustomLoggerBuilder {
203    /// Creates a new instance of `CustomLoggerBuilder` with no excluded routes.
204    pub fn new() -> Self {
205        Self {
206            excluded_routes: Vec::new(),
207        }
208    }
209
210    /// Specifies a route to be excluded from logging.
211    ///
212    /// # Arguments
213    ///
214    /// * `route` - The path of the route to exclude.
215    pub fn exclude(mut self, route: &str) -> Self {
216        self.excluded_routes.push(route.to_string());
217        self
218    }
219
220    /// Builds and returns a custom logging middleware.
221    ///
222    /// This middleware will use `CustomFilterRootSpanBuilder` to filter out the specified routes.
223    pub fn build(self) -> TracingLogger<CustomFilterRootSpanBuilder> {
224        // Set the excluded routes for our custom builder
225        CustomFilterRootSpanBuilder::set_excluded_routes(self.excluded_routes);
226
227        // Return a TracingLogger with our custom builder
228        TracingLogger::<CustomFilterRootSpanBuilder>::new()
229    }
230}
231
232impl Default for CustomLoggerBuilder {
233    fn default() -> Self {
234        Self::new()
235    }
236}
237
238/// Convenience function to obtain a `CustomLoggerBuilder`.
239///
240/// This can be used to start the builder chain for constructing the custom logger.
241pub 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}