Skip to main content

stormchaser_engine/
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, logging, and metrics based on environment variables.
12pub fn init_telemetry(rust_log: &str) -> anyhow::Result<()> {
13    let service_name =
14        std::env::var("OTEL_SERVICE_NAME").unwrap_or_else(|_| "stormchaser-engine".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-engine"))
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 = tracing_subscriber::EnvFilter::new(rust_log);
62
63    let registry = tracing_subscriber::registry()
64        .with(env_filter)
65        .with(tracing_subscriber::fmt::layer());
66
67    if let Some(tracer) = tracer {
68        registry
69            .with(tracing_opentelemetry::layer().with_tracer(tracer))
70            .init();
71    } else {
72        registry.init();
73    }
74
75    Ok(())
76}
77
78/// Shutdown telemetry.
79pub fn shutdown_telemetry() {
80    global::shutdown_tracer_provider();
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use std::sync::Once;
87
88    static INIT: Once = Once::new();
89
90    #[test]
91    fn test_init_telemetry_no_otel_endpoint() {
92        // Ensure OTEL endpoint is NOT set for this test
93        std::env::remove_var("OTEL_EXPORTER_OTLP_ENDPOINT");
94
95        // We can only init once per process safely
96        INIT.call_once(|| {
97            init_telemetry("debug").unwrap();
98        });
99    }
100
101    #[test]
102    fn test_shutdown_telemetry() {
103        shutdown_telemetry();
104    }
105}