cognee_observability/
guard.rs

1//! RAII guard that flushes and shuts down the OTEL pipeline on drop.
2//!
3//! The guard always exists, even when telemetry is disabled at compile
4//! time (via the absent `telemetry` feature) or at runtime (no endpoint
5//! configured). The disabled variant is a no-op so callers do not need
6//! cfg-gating around the call site.
7
8use std::time::Duration;
9
10#[cfg(feature = "telemetry")]
11use opentelemetry_sdk::trace::SdkTracerProvider;
12
13const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
14
15/// RAII handle that flushes and shuts down the global tracer provider on
16/// drop.
17///
18/// Holding the guard for the lifetime of `main()` (CLI) or for as long as
19/// `AppState` is alive (HTTP server) ensures the final batch of spans is
20/// exported before the process exits.
21#[must_use = "TelemetryGuard must be held for the lifetime of the process to flush spans on shutdown"]
22pub struct TelemetryGuard {
23    #[cfg(feature = "telemetry")]
24    provider: Option<SdkTracerProvider>,
25    timeout: Duration,
26}
27
28impl TelemetryGuard {
29    /// Construct a noop guard. Drop is free.
30    pub fn noop() -> Self {
31        Self {
32            #[cfg(feature = "telemetry")]
33            provider: None,
34            timeout: DEFAULT_SHUTDOWN_TIMEOUT,
35        }
36    }
37
38    #[cfg(feature = "telemetry")]
39    pub(crate) fn from_provider(provider: SdkTracerProvider) -> Self {
40        Self {
41            provider: Some(provider),
42            timeout: DEFAULT_SHUTDOWN_TIMEOUT,
43        }
44    }
45
46    /// Override the flush+shutdown budget (mostly useful in tests).
47    pub fn with_timeout(mut self, timeout: Duration) -> Self {
48        self.timeout = timeout;
49        self
50    }
51
52    /// Test-only inspector: returns `true` when an SDK provider is held.
53    #[cfg(all(feature = "telemetry", any(test, debug_assertions)))]
54    pub fn has_provider(&self) -> bool {
55        self.provider.is_some()
56    }
57
58    /// Test-only inspector: always `false` without `telemetry`.
59    #[cfg(all(not(feature = "telemetry"), any(test, debug_assertions)))]
60    pub fn has_provider(&self) -> bool {
61        false
62    }
63}
64
65impl Drop for TelemetryGuard {
66    fn drop(&mut self) {
67        #[cfg(feature = "telemetry")]
68        {
69            if let Some(provider) = self.provider.take() {
70                if let Err(err) = provider.force_flush() {
71                    tracing::warn!(
72                        target: "cognee.observability",
73                        ?err,
74                        "OTEL force_flush failed during TelemetryGuard drop"
75                    );
76                }
77                if let Err(err) = provider.shutdown_with_timeout(self.timeout) {
78                    tracing::warn!(
79                        target: "cognee.observability",
80                        ?err,
81                        "OTEL shutdown_with_timeout failed during TelemetryGuard drop"
82                    );
83                }
84            }
85        }
86        // Without `telemetry`, dropping is free; the timeout field is
87        // retained for signature stability.
88        let _ = self.timeout;
89    }
90}
cognee_observability/guard.rs

cognee_observability/
guard.rs
