adk_telemetry/

init.rs

//! Telemetry initialization and configuration

use std::sync::{
    Arc, Once, OnceLock,
    atomic::{AtomicBool, Ordering},
};
use tracing_subscriber::{
    EnvFilter,
    filter::filter_fn,
    layer::{Layer, SubscriberExt},
    util::SubscriberInitExt,
};

use crate::span_exporter::{AdkSpanExporter, AdkSpanLayer, is_runtime_span};

pub(crate) static INIT: Once = Once::new();
static ADK_EXPORTER: OnceLock<Arc<AdkSpanExporter>> = OnceLock::new();
static ADK_EXPORTER_INSTALLED: AtomicBool = AtomicBool::new(false);

/// Error returned by telemetry initialization functions.
#[derive(Debug, thiserror::Error)]
pub enum TelemetryError {
    /// Failed to build the tracing/OTLP pipeline.
    #[error("telemetry init failed: {0}")]
    Init(String),
}

/// Initialize basic telemetry with console logging.
///
/// # Arguments
/// * `service_name` - Name of the service for trace identification
///
/// # Example
/// ```
/// use adk_telemetry::init_telemetry;
/// init_telemetry("my-agent-service").expect("Failed to initialize telemetry");
/// ```
pub fn init_telemetry(service_name: &str) -> Result<(), TelemetryError> {
    INIT.call_once(|| {
        let filter = EnvFilter::try_from_default_env()
            .or_else(|_| EnvFilter::try_new("info"))
            .unwrap_or_else(|_| EnvFilter::new("info"));

        tracing_subscriber::registry()
            .with(filter)
            .with(
                tracing_subscriber::fmt::layer()
                    .with_target(true)
                    .with_thread_ids(true)
                    .with_line_number(true),
            )
            .init();

        tracing::info!(service.name = service_name, "telemetry initialized");
    });

    Ok(())
}

/// Initialize telemetry with OpenTelemetry OTLP export.
///
/// Enables distributed tracing by exporting spans to an OTLP collector.
///
/// # Arguments
/// * `service_name` - Name of the service for trace identification
/// * `endpoint` - OTLP collector endpoint (e.g., "http://localhost:4317")
///
/// # Example
/// ```no_run
/// use adk_telemetry::init_with_otlp;
/// init_with_otlp("my-agent", "http://localhost:4317")
///     .expect("Failed to initialize telemetry");
/// ```
#[cfg(feature = "otlp")]
pub fn init_with_otlp(service_name: &str, endpoint: &str) -> Result<(), TelemetryError> {
    use opentelemetry::trace::TracerProvider;
    use opentelemetry_otlp::WithExportConfig;
    use tracing_opentelemetry::OpenTelemetryLayer;

    let endpoint = endpoint.to_string();
    let service_name = service_name.to_string();

    let init_error: std::sync::Mutex<Option<String>> = std::sync::Mutex::new(None);

    INIT.call_once(|| {
        let resource = opentelemetry_sdk::Resource::builder_empty()
            .with_attributes([opentelemetry::KeyValue::new("service.name", service_name.clone())])
            .build();

        // Build OTLP span exporter
        let span_exporter = match opentelemetry_otlp::SpanExporter::builder()
            .with_tonic()
            .with_endpoint(&endpoint)
            .build()
        {
            Ok(e) => e,
            Err(e) => {
                *init_error.lock().unwrap_or_else(|p| p.into_inner()) =
                    Some(format!("failed to build OTLP span exporter: {e}"));
                return;
            }
        };

        // Build tracer provider with batch exporter
        let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
            .with_batch_exporter(span_exporter)
            .with_resource(resource.clone())
            .build();

        let tracer = tracer_provider.tracer("adk-telemetry");
        opentelemetry::global::set_tracer_provider(tracer_provider);

        // Initialize metrics
        let metric_exporter = match opentelemetry_otlp::MetricExporter::builder()
            .with_tonic()
            .with_endpoint(&endpoint)
            .build()
        {
            Ok(e) => e,
            Err(e) => {
                *init_error.lock().unwrap_or_else(|p| p.into_inner()) =
                    Some(format!("failed to build OTLP metric exporter: {e}"));
                return;
            }
        };

        let meter_provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder()
            .with_periodic_exporter(metric_exporter)
            .with_resource(resource)
            .build();

        opentelemetry::global::set_meter_provider(meter_provider);

        let telemetry_layer = OpenTelemetryLayer::new(tracer);

        let filter = EnvFilter::try_from_default_env()
            .or_else(|_| EnvFilter::try_new("info"))
            .unwrap_or_else(|_| EnvFilter::new("info"));

        tracing_subscriber::registry()
            .with(
                tracing_subscriber::fmt::layer()
                    .with_target(true)
                    .with_thread_ids(true)
                    .with_line_number(true)
                    .with_filter(filter),
            )
            .with(telemetry_layer)
            .init();

        tracing::info!(
            service.name = service_name,
            otlp.endpoint = %endpoint,
            "telemetry initialized with OpenTelemetry"
        );
    });

    if let Some(err) = init_error.lock().unwrap_or_else(|p| p.into_inner()).take() {
        return Err(TelemetryError::Init(err));
    }

    Ok(())
}

/// The tonic OTLP pipeline with a configuration hook on each exporter builder.
///
/// The hook is the seam that lets the `gcp` module inject TLS settings and a
/// per-request auth interceptor without duplicating the exporter/provider
/// plumbing shared with [`build_otlp_layer`].
#[cfg(feature = "otlp")]
pub(crate) mod otlp_pipeline {
    use super::TelemetryError;
    use opentelemetry::trace::TracerProvider;
    use opentelemetry_otlp::{WithExportConfig, WithTonicConfig};

    /// Configures a tonic exporter builder before it is built.
    pub(crate) trait ExporterHook {
        /// Returns the builder with hook-specific configuration applied.
        fn configure<B: WithTonicConfig>(&self, builder: B) -> B;
    }

    /// Hook that leaves the builder unchanged — the plain-collector path.
    pub(crate) struct NoopHook;

    impl ExporterHook for NoopHook {
        fn configure<B: WithTonicConfig>(&self, builder: B) -> B {
            builder
        }
    }

    /// Builds the OTLP span pipeline (exporter → batch tracer provider →
    /// global registration) and returns a tracer for layer construction.
    pub(crate) fn build_tracer<H: ExporterHook>(
        resource: opentelemetry_sdk::Resource,
        endpoint: &str,
        hook: &H,
    ) -> Result<opentelemetry_sdk::trace::SdkTracer, TelemetryError> {
        let span_exporter = hook
            .configure(
                opentelemetry_otlp::SpanExporter::builder().with_tonic().with_endpoint(endpoint),
            )
            .build()
            .map_err(|e| {
                TelemetryError::Init(format!("failed to build OTLP span exporter: {e}"))
            })?;

        let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
            .with_batch_exporter(span_exporter)
            .with_resource(resource)
            .build();

        let tracer = tracer_provider.tracer("adk-telemetry");
        opentelemetry::global::set_tracer_provider(tracer_provider);
        Ok(tracer)
    }
}

/// Build an OTLP tracing layer without initializing a global subscriber.
///
/// Returns a boxed [`tracing_subscriber::Layer`] that can be composed with any
/// subscriber via `.with()`. Also configures the global OpenTelemetry tracer
/// and meter providers.
///
/// The layer is returned as `Box<dyn Layer<S>>` rather than `impl Layer` so it
/// can be stored, composed across crate boundaries, and used in `Layered<...>`
/// chains without running into opaque-type limitations.
///
/// Unlike [`init_with_otlp`], this function does **not** call `.init()` on a
/// subscriber and does **not** use the `INIT` [`Once`] guard. The caller is
/// responsible for composing the returned layer into their own subscriber stack.
///
/// # Arguments
/// * `service_name` - Name of the service for trace identification
/// * `endpoint` - OTLP collector endpoint (e.g., `"http://localhost:4317"`)
///
/// # Errors
/// Returns [`TelemetryError::Init`] if the OTLP span or metric exporter fails to build.
///
/// # Example
/// ```no_run
/// use adk_telemetry::build_otlp_layer;
/// use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
///
/// let otlp_layer = build_otlp_layer("my-agent", "http://localhost:4317")
///     .expect("Failed to build OTLP layer");
///
/// tracing_subscriber::registry()
///     .with(otlp_layer)
///     .with(tracing_subscriber::fmt::layer())
///     .init();
/// ```
#[cfg(feature = "otlp")]
pub fn build_otlp_layer<S>(
    service_name: &str,
    endpoint: &str,
) -> Result<Box<dyn tracing_subscriber::Layer<S> + Send + Sync>, TelemetryError>
where
    S: tracing::Subscriber
        + for<'span> tracing_subscriber::registry::LookupSpan<'span>
        + Send
        + Sync,
{
    use opentelemetry_otlp::WithExportConfig;
    use tracing_opentelemetry::OpenTelemetryLayer;

    let resource = opentelemetry_sdk::Resource::builder_empty()
        .with_attributes([opentelemetry::KeyValue::new("service.name", service_name.to_string())])
        .build();

    let tracer = otlp_pipeline::build_tracer(resource.clone(), endpoint, &otlp_pipeline::NoopHook)?;

    // Build OTLP metric exporter
    let metric_exporter = opentelemetry_otlp::MetricExporter::builder()
        .with_tonic()
        .with_endpoint(endpoint)
        .build()
        .map_err(|e| TelemetryError::Init(format!("failed to build OTLP metric exporter: {e}")))?;

    let meter_provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder()
        .with_periodic_exporter(metric_exporter)
        .with_resource(resource)
        .build();

    opentelemetry::global::set_meter_provider(meter_provider);

    Ok(Box::new(OpenTelemetryLayer::new(tracer)))
}

/// Shutdown telemetry and flush any pending spans.
///
/// Should be called before application exit to ensure all telemetry data is sent.
/// In OTel 0.28+, the tracer provider is shut down when the last reference is dropped.
/// This function is kept for backward compatibility and explicitly drops the global provider.
pub fn shutdown_telemetry() {
    #[cfg(feature = "otlp")]
    {
        // In OTel 0.28, shutdown_tracer_provider() was removed.
        // The SdkTracerProvider shuts down automatically when the last reference is dropped.
        // We trigger this by replacing the global provider with a no-op, which drops the old one.
        opentelemetry::global::set_tracer_provider(
            opentelemetry::trace::noop::NoopTracerProvider::new(),
        );
    }
}

/// Initialize telemetry with ADK span exporter.
///
/// Creates a shared span exporter that can be used by both telemetry and the debug API.
/// Returns the exporter so it can be passed to the debug controller.
pub fn init_with_adk_exporter(service_name: &str) -> Result<Arc<AdkSpanExporter>, TelemetryError> {
    initialize_adk_exporter_with(&INIT, &ADK_EXPORTER, &ADK_EXPORTER_INSTALLED, |exporter| {
        let filter = EnvFilter::try_from_default_env()
            .or_else(|_| EnvFilter::try_new("info"))
            .unwrap_or_else(|_| EnvFilter::new("info"));

        let adk_layer = AdkSpanLayer::new(exporter).with_filter(filter_fn(|metadata| {
            metadata.is_span() && is_runtime_span(metadata.name())
        }));

        tracing_subscriber::registry()
            .with(
                tracing_subscriber::fmt::layer()
                    .with_target(true)
                    .with_thread_ids(true)
                    .with_line_number(true)
                    .with_filter(filter),
            )
            .with(adk_layer)
            .init();

        tracing::info!(service.name = service_name, "telemetry initialized with ADK span exporter");
    })
}

fn initialize_adk_exporter_with<F>(
    init: &Once,
    exporter_cell: &OnceLock<Arc<AdkSpanExporter>>,
    installed: &AtomicBool,
    install: F,
) -> Result<Arc<AdkSpanExporter>, TelemetryError>
where
    F: FnOnce(Arc<AdkSpanExporter>),
{
    let exporter = exporter_cell.get_or_init(|| Arc::new(AdkSpanExporter::new())).clone();
    init.call_once(|| {
        install(exporter.clone());
        installed.store(true, Ordering::Release);
    });

    if installed.load(Ordering::Acquire) {
        Ok(exporter)
    } else {
        Err(TelemetryError::Init(
            "global telemetry was already initialized without the ADK in-process exporter; \
             initialize the ADK exporter before other global telemetry modes"
                .to_string(),
        ))
    }
}

/// Initialize telemetry with direct SQLite span export — zero-infrastructure
/// tracing with no collector or backend to deploy.
///
/// Spans are persisted to the database file at `db_path` (created if needed)
/// by a background writer thread; the traced code path never blocks on I/O.
/// Read them back with [`SqliteTraceReader`](crate::sqlite::SqliteTraceReader)
/// or any SQLite client.
///
/// Returns the exporter so callers can [`flush`](crate::sqlite::SqliteSpanExporter::flush)
/// before exiting (the subscriber keeps it alive for the process lifetime, so
/// drop-based flushing never fires for globally installed subscribers).
///
/// # Example
/// ```no_run
/// use adk_telemetry::init_with_sqlite;
///
/// let exporter = init_with_sqlite("my-agent", "traces.db")
///     .expect("Failed to initialize telemetry");
/// // ... run the agent ...
/// exporter.flush().ok();
/// ```
#[cfg(feature = "sqlite")]
pub fn init_with_sqlite(
    service_name: &str,
    db_path: impl AsRef<std::path::Path>,
) -> Result<Arc<crate::sqlite::SqliteSpanExporter>, TelemetryError> {
    // Create the exporter (and surface db errors) before the irreversible
    // global subscriber installation.
    let exporter = Arc::new(crate::sqlite::SqliteSpanExporter::new(db_path)?);
    let exporter_clone = exporter.clone();

    INIT.call_once(|| {
        let filter = EnvFilter::try_from_default_env()
            .or_else(|_| EnvFilter::try_new("info"))
            .unwrap_or_else(|_| EnvFilter::new("info"));

        let adk_layer = AdkSpanLayer::new(exporter_clone).with_filter(filter_fn(|metadata| {
            metadata.is_span() && is_runtime_span(metadata.name())
        }));

        tracing_subscriber::registry()
            .with(
                tracing_subscriber::fmt::layer()
                    .with_target(true)
                    .with_thread_ids(true)
                    .with_line_number(true)
                    .with_filter(filter),
            )
            .with(adk_layer)
            .init();

        tracing::info!(
            service.name = service_name,
            "telemetry initialized with SQLite span exporter"
        );
    });

    Ok(exporter)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicUsize;

    #[test]
    fn repeated_adk_initialization_reuses_the_registered_exporter() {
        let init = Once::new();
        let exporter = OnceLock::new();
        let installed = AtomicBool::new(false);
        let installations = AtomicUsize::new(0);

        let first = initialize_adk_exporter_with(&init, &exporter, &installed, |_| {
            installations.fetch_add(1, Ordering::Relaxed);
        })
        .unwrap();
        let second = initialize_adk_exporter_with(&init, &exporter, &installed, |_| {
            installations.fetch_add(1, Ordering::Relaxed);
        })
        .unwrap();

        assert!(Arc::ptr_eq(&first, &second));
        assert_eq!(installations.load(Ordering::Relaxed), 1);
    }

    #[test]
    fn adk_initialization_rejects_an_incompatible_existing_global_mode() {
        let init = Once::new();
        init.call_once(|| {});
        let exporter = OnceLock::new();
        let installed = AtomicBool::new(false);

        let error = initialize_adk_exporter_with(&init, &exporter, &installed, |_| {})
            .expect_err("an earlier telemetry mode must not return a disconnected exporter");

        assert!(error.to_string().contains("already initialized without the ADK"));
    }
}