tiny_agent/
telemetry.rs

1//! Optional OpenTelemetry helpers for exporting traces over OTLP.
2
3use thiserror::Error;
4
5#[derive(Debug, Error)]
6pub enum TelemetryError {
7    #[error("failed to initialize OTLP span exporter: {0}")]
8    Exporter(String),
9    #[error("failed to initialize tracing subscriber: {0}")]
10    Subscriber(String),
11}
12
13#[derive(Debug, Clone)]
14pub struct TempoConfig {
15    pub endpoint: String,
16    pub service_name: String,
17    pub service_namespace: String,
18    pub environment: String,
19    pub tracer_name: String,
20}
21
22impl TempoConfig {
23    pub fn new(service_name: impl Into<String>) -> Self {
24        let service_name = service_name.into();
25        Self {
26            endpoint: otlp_endpoint_from_env()
27                .unwrap_or_else(|| "http://127.0.0.1:4317".to_string()),
28            tracer_name: service_name.clone(),
29            service_name,
30            service_namespace: "tiny-agent".to_string(),
31            environment: "local".to_string(),
32        }
33    }
34
35    pub fn endpoint(mut self, endpoint: impl Into<String>) -> Self {
36        self.endpoint = endpoint.into();
37        self
38    }
39
40    pub fn service_namespace(mut self, namespace: impl Into<String>) -> Self {
41        self.service_namespace = namespace.into();
42        self
43    }
44
45    pub fn environment(mut self, environment: impl Into<String>) -> Self {
46        self.environment = environment.into();
47        self
48    }
49
50    pub fn tracer_name(mut self, tracer_name: impl Into<String>) -> Self {
51        self.tracer_name = tracer_name.into();
52        self
53    }
54}
55
56pub fn otel_enabled_from_env() -> bool {
57    matches!(
58        std::env::var("TINY_AGENT_OTEL").as_deref(),
59        Ok("1" | "true" | "TRUE" | "yes" | "YES")
60    ) || std::env::var_os("OTEL_EXPORTER_OTLP_ENDPOINT").is_some()
61        || std::env::var_os("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT").is_some()
62}
63
64pub fn otlp_endpoint_from_env() -> Option<String> {
65    std::env::var("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT")
66        .or_else(|_| std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT"))
67        .ok()
68}
69
70#[cfg(feature = "otel-tempo")]
71mod imp {
72    use super::{TelemetryError, TempoConfig, otel_enabled_from_env};
73    use opentelemetry::{KeyValue, global, trace::TracerProvider as _};
74    use opentelemetry_otlp::WithExportConfig;
75    use opentelemetry_sdk::{Resource, trace::SdkTracerProvider};
76    use tracing_subscriber::{
77        Layer,
78        filter::{LevelFilter, Targets},
79        layer::SubscriberExt,
80        registry::LookupSpan,
81        util::SubscriberInitExt,
82    };
83
84    pub struct TelemetryGuard {
85        provider: Option<SdkTracerProvider>,
86        endpoint: String,
87    }
88
89    impl TelemetryGuard {
90        pub fn endpoint(&self) -> &str {
91            &self.endpoint
92        }
93    }
94
95    impl Drop for TelemetryGuard {
96        fn drop(&mut self) {
97            if let Some(provider) = self.provider.take() {
98                let _ = provider.force_flush();
99                let _ = provider.shutdown();
100            }
101        }
102    }
103
104    pub fn init_tempo_from_env(
105        service_name: impl Into<String>,
106    ) -> Result<Option<TelemetryGuard>, TelemetryError> {
107        if !otel_enabled_from_env() {
108            return Ok(None);
109        }
110        init_tempo(TempoConfig::new(service_name)).map(Some)
111    }
112
113    pub fn init_tempo(config: TempoConfig) -> Result<TelemetryGuard, TelemetryError> {
114        let provider = build_provider(&config)?;
115        global::set_tracer_provider(provider.clone());
116
117        let otel_layer = tempo_layer(&provider, &config.tracer_name);
118        tracing_subscriber::registry()
119            .with(otel_layer)
120            .try_init()
121            .map_err(|e| TelemetryError::Subscriber(e.to_string()))?;
122
123        Ok(TelemetryGuard {
124            provider: Some(provider),
125            endpoint: config.endpoint,
126        })
127    }
128
129    pub fn build_provider(config: &TempoConfig) -> Result<SdkTracerProvider, TelemetryError> {
130        let exporter = opentelemetry_otlp::SpanExporter::builder()
131            .with_tonic()
132            .with_endpoint(&config.endpoint)
133            .build()
134            .map_err(|e| TelemetryError::Exporter(e.to_string()))?;
135
136        Ok(SdkTracerProvider::builder()
137            .with_batch_exporter(exporter)
138            .with_resource(
139                Resource::builder()
140                    .with_service_name(config.service_name.clone())
141                    .with_attributes([
142                        KeyValue::new("service.namespace", config.service_namespace.clone()),
143                        KeyValue::new("deployment.environment.name", config.environment.clone()),
144                    ])
145                    .build(),
146            )
147            .build())
148    }
149
150    pub fn tempo_layer<S>(
151        provider: &SdkTracerProvider,
152        tracer_name: &str,
153    ) -> impl Layer<S> + Send + Sync + 'static
154    where
155        S: tracing::Subscriber + Send + Sync + 'static,
156        for<'span> S: LookupSpan<'span>,
157    {
158        let tracer = provider.tracer(tracer_name.to_string());
159        let targets = Targets::new().with_target("tiny_agent", LevelFilter::DEBUG);
160        tracing_opentelemetry::layer()
161            .with_tracer(tracer)
162            .with_filter(targets)
163    }
164}
165
166#[cfg(feature = "otel-tempo")]
167pub use imp::{TelemetryGuard, build_provider, init_tempo, init_tempo_from_env, tempo_layer};
tiny_agent/telemetry.rs

tiny_agent/
telemetry.rs
