diff --git a/codex-rs/otel/src/provider.rs b/codex-rs/otel/src/provider.rs index 26f732da71..4c33f9a29b 100644 --- a/codex-rs/otel/src/provider.rs +++ b/codex-rs/otel/src/provider.rs @@ -39,6 +39,8 @@ use opentelemetry_sdk::trace::span_processor_with_async_runtime::BatchSpanProces use opentelemetry_semantic_conventions as semconv; use std::collections::BTreeMap; use std::error::Error; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use tracing::debug; use tracing_subscriber::Layer; @@ -58,12 +60,17 @@ pub struct OtelProvider { pub tracer_provider: Option, pub tracer: Option, pub metrics: Option, + shutdown_started: AtomicBool, } impl OtelProvider { + /// Flushes and shuts down configured exporters at most once. pub fn shutdown(&self) { + if self.shutdown_started.swap(/*val*/ true, Ordering::AcqRel) { + return; + } + if let Some(tracer_provider) = &self.tracer_provider { - let _ = tracer_provider.force_flush(); let _ = tracer_provider.shutdown(); } if let Some(metrics) = &self.metrics { @@ -150,6 +157,7 @@ impl OtelProvider { tracer_provider, tracer, metrics, + shutdown_started: AtomicBool::default(), })) } @@ -196,16 +204,7 @@ impl OtelProvider { impl Drop for OtelProvider { fn drop(&mut self) { - if let Some(tracer_provider) = &self.tracer_provider { - let _ = tracer_provider.force_flush(); - let _ = tracer_provider.shutdown(); - } - if let Some(metrics) = &self.metrics { - let _ = metrics.shutdown(); - } - if let Some(logger) = &self.logger { - let _ = logger.shutdown(); - } + self.shutdown(); } } @@ -456,6 +455,10 @@ fn build_tracer_provider( .build()) } +#[cfg(test)] +#[path = "provider_shutdown_tests.rs"] +mod shutdown_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/codex-rs/otel/src/provider_shutdown_tests.rs b/codex-rs/otel/src/provider_shutdown_tests.rs new file mode 100644 index 0000000000..939bffe7ac --- /dev/null +++ b/codex-rs/otel/src/provider_shutdown_tests.rs @@ -0,0 +1,85 @@ +use super::Context; +use super::OTelSdkResult; +use super::OtelProvider; +use super::SdkTracerProvider; +use super::Span; +use super::SpanData; +use super::SpanProcessor; +use pretty_assertions::assert_eq; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +#[derive(Debug, Default)] +struct ShutdownState { + force_flushes: AtomicUsize, + shutdowns: AtomicUsize, +} + +#[derive(Debug)] +struct ControlledSpanProcessor { + state: Arc, +} + +impl SpanProcessor for ControlledSpanProcessor { + fn on_start(&self, _span: &mut Span, _context: &Context) {} + + fn on_end(&self, _span: SpanData) {} + + fn force_flush(&self) -> OTelSdkResult { + self.state + .force_flushes + .fetch_add(/*val*/ 1, Ordering::Relaxed); + Ok(()) + } + + fn shutdown_with_timeout(&self, _timeout: Duration) -> OTelSdkResult { + self.state.shutdowns.fetch_add(/*val*/ 1, Ordering::Relaxed); + Ok(()) + } +} + +fn test_provider() -> (OtelProvider, Arc) { + let state = Arc::new(ShutdownState::default()); + let processor = ControlledSpanProcessor { + state: Arc::clone(&state), + }; + let tracer_provider = SdkTracerProvider::builder() + .with_span_processor(processor) + .build(); + + ( + OtelProvider { + logger: None, + tracer_provider: Some(tracer_provider), + tracer: None, + metrics: None, + shutdown_started: AtomicBool::default(), + }, + state, + ) +} + +#[test] +fn explicit_shutdown_and_drop_shut_down_exporters_once_without_force_flush() { + let (provider, state) = test_provider(); + + provider.shutdown(); + provider.shutdown(); + drop(provider); + + assert_eq!(state.shutdowns.load(Ordering::Relaxed), 1); + assert_eq!(state.force_flushes.load(Ordering::Relaxed), 0); +} + +#[test] +fn drop_shuts_down_exporters_without_force_flush() { + let (provider, state) = test_provider(); + + drop(provider); + + assert_eq!(state.shutdowns.load(Ordering::Relaxed), 1); + assert_eq!(state.force_flushes.load(Ordering::Relaxed), 0); +}