tiny-agent 0.3.0

//! Optional OpenTelemetry helpers for exporting traces over OTLP.

use thiserror::Error;

#[derive(Debug, Error)]
pub enum TelemetryError {
    #[error("failed to initialize OTLP span exporter: {0}")]
    Exporter(String),
    #[error("failed to initialize tracing subscriber: {0}")]
    Subscriber(String),
}

#[derive(Debug, Clone)]
pub struct TempoConfig {
    pub endpoint: String,
    pub service_name: String,
    pub service_namespace: String,
    pub environment: String,
    pub tracer_name: String,
}

impl TempoConfig {
    pub fn new(service_name: impl Into<String>) -> Self {
        let service_name = service_name.into();
        Self {
            endpoint: otlp_endpoint_from_env()
                .unwrap_or_else(|| "http://127.0.0.1:4317".to_string()),
            tracer_name: service_name.clone(),
            service_name,
            service_namespace: "tiny-agent".to_string(),
            environment: "local".to_string(),
        }
    }

    pub fn endpoint(mut self, endpoint: impl Into<String>) -> Self {
        self.endpoint = endpoint.into();
        self
    }

    pub fn service_namespace(mut self, namespace: impl Into<String>) -> Self {
        self.service_namespace = namespace.into();
        self
    }

    pub fn environment(mut self, environment: impl Into<String>) -> Self {
        self.environment = environment.into();
        self
    }

    pub fn tracer_name(mut self, tracer_name: impl Into<String>) -> Self {
        self.tracer_name = tracer_name.into();
        self
    }
}

pub fn otel_enabled_from_env() -> bool {
    matches!(
        std::env::var("TINY_AGENT_OTEL").as_deref(),
        Ok("1" | "true" | "TRUE" | "yes" | "YES")
    ) || std::env::var_os("OTEL_EXPORTER_OTLP_ENDPOINT").is_some()
        || std::env::var_os("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT").is_some()
}

pub fn otlp_endpoint_from_env() -> Option<String> {
    std::env::var("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT")
        .or_else(|_| std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT"))
        .ok()
}

#[cfg(feature = "otel-tempo")]
mod imp {
    use super::{TelemetryError, TempoConfig, otel_enabled_from_env};
    use opentelemetry::{KeyValue, global, trace::TracerProvider as _};
    use opentelemetry_otlp::WithExportConfig;
    use opentelemetry_sdk::{Resource, trace::SdkTracerProvider};
    use tracing_subscriber::{
        Layer,
        filter::{LevelFilter, Targets},
        layer::SubscriberExt,
        registry::LookupSpan,
        util::SubscriberInitExt,
    };

    pub struct TelemetryGuard {
        provider: Option<SdkTracerProvider>,
        endpoint: String,
    }

    impl TelemetryGuard {
        pub fn endpoint(&self) -> &str {
            &self.endpoint
        }
    }

    impl Drop for TelemetryGuard {
        fn drop(&mut self) {
            if let Some(provider) = self.provider.take() {
                let _ = provider.force_flush();
                let _ = provider.shutdown();
            }
        }
    }

    pub fn init_tempo_from_env(
        service_name: impl Into<String>,
    ) -> Result<Option<TelemetryGuard>, TelemetryError> {
        if !otel_enabled_from_env() {
            return Ok(None);
        }
        init_tempo(TempoConfig::new(service_name)).map(Some)
    }

    pub fn init_tempo(config: TempoConfig) -> Result<TelemetryGuard, TelemetryError> {
        let provider = build_provider(&config)?;
        global::set_tracer_provider(provider.clone());

        let otel_layer = tempo_layer(&provider, &config.tracer_name);
        tracing_subscriber::registry()
            .with(otel_layer)
            .try_init()
            .map_err(|e| TelemetryError::Subscriber(e.to_string()))?;

        Ok(TelemetryGuard {
            provider: Some(provider),
            endpoint: config.endpoint,
        })
    }

    pub fn build_provider(config: &TempoConfig) -> Result<SdkTracerProvider, TelemetryError> {
        let exporter = opentelemetry_otlp::SpanExporter::builder()
            .with_tonic()
            .with_endpoint(&config.endpoint)
            .build()
            .map_err(|e| TelemetryError::Exporter(e.to_string()))?;

        Ok(SdkTracerProvider::builder()
            .with_batch_exporter(exporter)
            .with_resource(
                Resource::builder()
                    .with_service_name(config.service_name.clone())
                    .with_attributes([
                        KeyValue::new("service.namespace", config.service_namespace.clone()),
                        KeyValue::new("deployment.environment.name", config.environment.clone()),
                    ])
                    .build(),
            )
            .build())
    }

    pub fn tempo_layer<S>(
        provider: &SdkTracerProvider,
        tracer_name: &str,
    ) -> impl Layer<S> + Send + Sync + 'static
    where
        S: tracing::Subscriber + Send + Sync + 'static,
        for<'span> S: LookupSpan<'span>,
    {
        let tracer = provider.tracer(tracer_name.to_string());
        let targets = Targets::new().with_target("tiny_agent", LevelFilter::DEBUG);
        tracing_opentelemetry::layer()
            .with_tracer(tracer)
            .with_filter(targets)
    }
}

#[cfg(feature = "otel-tempo")]
pub use imp::{TelemetryGuard, build_provider, init_tempo, init_tempo_from_env, tempo_layer};