klieo-otel 3.3.0

//! Lifetime guard returned by [`crate::install`].

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use opentelemetry_sdk::trace::TracerProvider as SdkTracerProvider;

/// Guard that flushes and shuts down the OTel tracer provider on
/// `Drop`.
///
/// Bind it to a long-lived binding in `main` so that the last batch of
/// spans is exported before the process exits. Dropping the guard:
///
/// 1. Calls [`force_flush`](SdkTracerProvider::force_flush) to push any
///    pending spans through the OTLP exporter.
/// 2. Calls [`shutdown`](SdkTracerProvider::shutdown) to release the
///    exporter's network resources.
///
/// Any errors in either step are logged via `tracing::warn!` rather
/// than panicking — a crashing observability layer must not crash the
/// process it observes.
#[must_use = "OtelGuard drops immediately if not bound to a named local; use `let _otel = install(...)?;` (NOT `let _ = install(...)?;`) to keep the guard alive for the process lifetime — see install's doc-comment"]
pub struct OtelGuard {
    provider: Option<SdkTracerProvider>,
    /// Set to `true` when shutdown completes — exposed for tests.
    shutdown_marker: Arc<AtomicBool>,
}

impl OtelGuard {
    /// Construct a new guard wrapping the given provider.
    pub(crate) fn new(provider: SdkTracerProvider) -> Self {
        Self {
            provider: Some(provider),
            shutdown_marker: Arc::new(AtomicBool::new(false)),
        }
    }

    /// Returns `true` once the guard's `Drop` impl has run flush +
    /// shutdown. Test-only hook.
    #[doc(hidden)]
    #[must_use]
    pub fn shutdown_marker(&self) -> Arc<AtomicBool> {
        Arc::clone(&self.shutdown_marker)
    }
}

impl Drop for OtelGuard {
    fn drop(&mut self) {
        if let Some(provider) = self.provider.take() {
            for result in provider.force_flush() {
                if let Err(err) = result {
                    tracing::warn!(error = %err, "OTel force_flush returned an error");
                }
            }
            if let Err(err) = provider.shutdown() {
                tracing::warn!(error = %err, "OTel provider shutdown returned an error");
            }
        }
        self.shutdown_marker.store(true, Ordering::SeqCst);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{build_provider, config::OtelConfig};
    use std::time::Duration;

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn drop_marks_shutdown_complete() {
        // >=2 worker threads: the batch processor needs one, the
        // synchronous `force_flush` will block another while it waits
        // on the oneshot reply.
        let cfg = OtelConfig::new("svc")
            .endpoint("http://127.0.0.1:1")
            .timeout(Duration::from_millis(50));
        let provider = build_provider(&cfg, &cfg.endpoint).expect("build provider");
        let guard = OtelGuard::new(provider);
        let marker = guard.shutdown_marker();
        assert!(!marker.load(Ordering::SeqCst));
        tokio::task::spawn_blocking(move || drop(guard))
            .await
            .expect("drop must not panic");
        assert!(marker.load(Ordering::SeqCst));
    }
}