Skip to main content

stormchaser_api/
telemetry.rs

1use opentelemetry::{global, KeyValue};
2use opentelemetry_otlp::WithExportConfig;
3use opentelemetry_sdk::{
4    metrics::{PeriodicReader, SdkMeterProvider},
5    runtime,
6    trace::TracerProvider,
7    Resource,
8};
9use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
10
11/// Initializes OpenTelemetry tracing and metrics.
12pub fn init_telemetry() -> anyhow::Result<()> {
13    let service_name =
14        std::env::var("OTEL_SERVICE_NAME").unwrap_or_else(|_| "stormchaser-api".to_string());
15    let otlp_endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT")
16        .unwrap_or_else(|_| "http://localhost:4317".to_string());
17
18    let resource = Resource::new(vec![KeyValue::new(
19        opentelemetry_semantic_conventions::resource::SERVICE_NAME,
20        service_name,
21    )]);
22
23    // Configure Tracing (Optional)
24    let tracer = if std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").is_ok() {
25        let exporter = opentelemetry_otlp::SpanExporter::builder()
26            .with_tonic()
27            .with_endpoint(&otlp_endpoint)
28            .build()?;
29
30        let tracer_provider = TracerProvider::builder()
31            .with_batch_exporter(exporter, runtime::Tokio)
32            .with_resource(resource.clone())
33            .build();
34
35        global::set_tracer_provider(tracer_provider.clone());
36
37        use opentelemetry::trace::TracerProvider as _;
38        Some(tracer_provider.tracer("stormchaser-api"))
39    } else {
40        None
41    };
42
43    // Configure Metrics (Optional)
44    if std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").is_ok() {
45        let metric_exporter = opentelemetry_otlp::MetricExporter::builder()
46            .with_tonic()
47            .with_endpoint(&otlp_endpoint)
48            .build()?;
49
50        let reader = PeriodicReader::builder(metric_exporter, runtime::Tokio).build();
51
52        let meter_provider = SdkMeterProvider::builder()
53            .with_reader(reader)
54            .with_resource(resource)
55            .build();
56
57        global::set_meter_provider(meter_provider);
58    }
59
60    // Initialize Subscriber
61    let env_filter =
62        tracing_subscriber::EnvFilter::new(std::env::var("RUST_LOG").unwrap_or_else(|_| {
63            "stormchaser_api=debug,tower_http=debug,axum::rejection=trace".into()
64        }));
65
66    let registry = tracing_subscriber::registry()
67        .with(env_filter)
68        .with(tracing_subscriber::fmt::layer());
69
70    if let Some(tracer) = tracer {
71        registry
72            .with(tracing_opentelemetry::layer().with_tracer(tracer))
73            .init();
74    } else {
75        registry.init();
76    }
77
78    Ok(())
79}
80
81/// Shuts down the OpenTelemetry tracer provider, ensuring all telemetry is exported
82pub fn shutdown_telemetry() {
83    global::shutdown_tracer_provider();
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use std::sync::Once;
90
91    static INIT: Once = Once::new();
92
93    #[test]
94    fn test_init_telemetry_no_otel_endpoint() {
95        // Ensure OTEL endpoint is NOT set for this test
96        std::env::remove_var("OTEL_EXPORTER_OTLP_ENDPOINT");
97
98        // We can only init once per process safely
99        INIT.call_once(|| {
100            init_telemetry().unwrap();
101        });
102    }
103
104    #[test]
105    fn test_shutdown_telemetry() {
106        shutdown_telemetry();
107    }
108}