From 2707dfc219b9bc9d38d7d37d1f691855f7b44c1e Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Wed, 5 Aug 2026 15:14:12 +0000 Subject: [PATCH] Bound interactive telemetry shutdown (#37109) ## Why A stalled telemetry exporter can block shutdown indefinitely when the interactive TUI exits. ## What changed - Add `OtelProvider::shutdown_with_timeout` to run exporter shutdown on a detached thread and report completion, worker failure, or timeout. - Give interactive TUI telemetry shutdown a 500 ms time budget and warn if cleanup does not finish successfully. ## Testing - Cover successful shutdown on current-thread and multi-thread Tokio runtimes, worker creation failure, and timeout behavior. - Verify that OTLP HTTP shutdown returns on time when a collector stalls. GitOrigin-RevId: 5b5f31a841fbc5788715cb3065d522c6dafa994f --- codex-rs/otel/src/provider.rs | 52 +++++ codex-rs/otel/src/provider_shutdown_tests.rs | 177 +++++++++++++++++- .../otel/tests/suite/otlp_http_loopback.rs | 85 ++++++++- codex-rs/tui/src/lib.rs | 16 +- 4 files changed, 319 insertions(+), 11 deletions(-) diff --git a/codex-rs/otel/src/provider.rs b/codex-rs/otel/src/provider.rs index 448bf07ac0..5c76d70c80 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::io; +use std::mem::ManuallyDrop; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; @@ -63,6 +65,11 @@ pub struct OtelProvider { shutdown_started: AtomicBool, } +struct ShutdownWorker { + provider: ManuallyDrop, + completed_tx: tokio::sync::oneshot::Sender<()>, +} + impl OtelProvider { /// Flushes and shuts down configured exporters at most once. pub fn shutdown(&self) { @@ -81,6 +88,51 @@ impl OtelProvider { } } + /// Shuts down exporters on a detached thread within an external time budget. + pub async fn shutdown_with_timeout(self, timeout: Duration) -> io::Result<()> { + self.shutdown_with_timeout_and_spawner(timeout, |worker| { + std::thread::Builder::new() + .name("codex-otel-shutdown".to_string()) + .spawn(move || { + let provider = ManuallyDrop::into_inner(worker.provider); + provider.shutdown(); + drop(provider); + let _ = worker.completed_tx.send(()); + }) + }) + .await + } + + async fn shutdown_with_timeout_and_spawner( + self, + timeout: Duration, + spawn: F, + ) -> io::Result<()> + where + F: FnOnce(ShutdownWorker) -> io::Result>, + { + let (completed_tx, completed_rx) = tokio::sync::oneshot::channel(); + // A failed spawn drops its closure on the caller. Keep the provider + // from synchronously running its potentially blocking destructor. + let worker = ShutdownWorker { + provider: ManuallyDrop::new(self), + completed_tx, + }; + let _shutdown_worker = spawn(worker)?; + + match tokio::time::timeout(timeout, completed_rx).await { + Ok(Ok(())) => Ok(()), + Ok(Err(_)) => Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "telemetry shutdown worker stopped before completing", + )), + Err(_) => Err(io::Error::new( + io::ErrorKind::TimedOut, + "telemetry shutdown exceeded its time budget", + )), + } + } + pub fn from(settings: &OtelSettings) -> Result, Box> { let log_enabled = !matches!(settings.exporter, OtelExporter::None); let trace_enabled = !matches!(settings.trace_exporter, OtelExporter::None); diff --git a/codex-rs/otel/src/provider_shutdown_tests.rs b/codex-rs/otel/src/provider_shutdown_tests.rs index 939bffe7ac..e68c8ac879 100644 --- a/codex-rs/otel/src/provider_shutdown_tests.rs +++ b/codex-rs/otel/src/provider_shutdown_tests.rs @@ -6,20 +6,35 @@ use super::Span; use super::SpanData; use super::SpanProcessor; use pretty_assertions::assert_eq; +use std::io::ErrorKind; use std::sync::Arc; +use std::sync::Condvar; +use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::sync::mpsc; use std::time::Duration; +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ShutdownBehavior { + Complete, + WaitForRelease, +} + #[derive(Debug, Default)] struct ShutdownState { force_flushes: AtomicUsize, shutdowns: AtomicUsize, + released: Mutex, + release_notification: Condvar, + started: Mutex>>, + completed: Mutex>>, } #[derive(Debug)] struct ControlledSpanProcessor { + behavior: ShutdownBehavior, state: Arc, } @@ -37,21 +52,54 @@ impl SpanProcessor for ControlledSpanProcessor { fn shutdown_with_timeout(&self, _timeout: Duration) -> OTelSdkResult { self.state.shutdowns.fetch_add(/*val*/ 1, Ordering::Relaxed); + if let Some(started) = self.state.started.lock().expect("started lock").take() { + let _ = started.send(()); + } + + if self.behavior == ShutdownBehavior::WaitForRelease { + drop( + self.state + .release_notification + .wait_while( + self.state.released.lock().expect("release lock"), + |released| !*released, + ) + .expect("release notification"), + ); + } + + if let Some(completed) = self.state.completed.lock().expect("completed lock").take() { + let _ = completed.send(()); + } Ok(()) } } -fn test_provider() -> (OtelProvider, Arc) { - let state = Arc::new(ShutdownState::default()); +struct TestProvider { + provider: OtelProvider, + state: Arc, + started: mpsc::Receiver<()>, + completed: mpsc::Receiver<()>, +} + +fn test_provider(behavior: ShutdownBehavior) -> TestProvider { + let (started_tx, started) = mpsc::channel(); + let (completed_tx, completed) = mpsc::channel(); + let state = Arc::new(ShutdownState { + started: Mutex::new(Some(started_tx)), + completed: Mutex::new(Some(completed_tx)), + ..ShutdownState::default() + }); let processor = ControlledSpanProcessor { + behavior, state: Arc::clone(&state), }; let tracer_provider = SdkTracerProvider::builder() .with_span_processor(processor) .build(); - ( - OtelProvider { + TestProvider { + provider: OtelProvider { logger: None, tracer_provider: Some(tracer_provider), tracer: None, @@ -59,27 +107,142 @@ fn test_provider() -> (OtelProvider, Arc) { shutdown_started: AtomicBool::default(), }, state, - ) + started, + completed, + } +} + +#[tokio::test(flavor = "current_thread")] +async fn bounded_shutdown_does_not_flush_when_worker_creation_fails() { + let TestProvider { + provider, + state, + started, + completed, + } = test_provider(ShutdownBehavior::Complete); + + let result = provider + .shutdown_with_timeout_and_spawner(Duration::from_secs(/*secs*/ 1), |_worker| { + Err(std::io::Error::new( + ErrorKind::WouldBlock, + "shutdown worker could not be created", + )) + }) + .await; + + assert_eq!( + result.as_ref().map_err(std::io::Error::kind), + Err(ErrorKind::WouldBlock) + ); + assert_eq!(state.shutdowns.load(Ordering::Relaxed), 0); + assert_eq!(state.force_flushes.load(Ordering::Relaxed), 0); + assert_eq!(started.try_recv(), Err(mpsc::TryRecvError::Empty)); + assert_eq!(completed.try_recv(), Err(mpsc::TryRecvError::Empty)); +} + +#[tokio::test(flavor = "current_thread")] +async fn bounded_shutdown_times_out_without_blocking_the_runtime() { + let TestProvider { + provider, + state, + started, + completed, + } = test_provider(ShutdownBehavior::WaitForRelease); + + let result = provider + .shutdown_with_timeout(Duration::from_millis(/*millis*/ 50)) + .await; + + assert_eq!( + result.as_ref().map_err(std::io::Error::kind), + Err(ErrorKind::TimedOut) + ); + started + .recv_timeout(Duration::from_secs(/*secs*/ 1)) + .expect("shutdown worker started"); + assert_eq!(state.shutdowns.load(Ordering::Relaxed), 1); + assert_eq!(state.force_flushes.load(Ordering::Relaxed), 0); + + *state.released.lock().expect("release lock") = true; + state.release_notification.notify_one(); + completed + .recv_timeout(Duration::from_secs(/*secs*/ 1)) + .expect("shutdown worker completed after release"); +} + +async fn assert_bounded_shutdown_completes() { + let TestProvider { + provider, + state, + started, + completed, + } = test_provider(ShutdownBehavior::Complete); + + provider + .shutdown_with_timeout(Duration::from_secs(/*secs*/ 1)) + .await + .expect("healthy processor shuts down"); + + started + .recv_timeout(Duration::from_secs(/*secs*/ 1)) + .expect("shutdown worker started"); + completed + .recv_timeout(Duration::from_secs(/*secs*/ 1)) + .expect("shutdown worker completed"); + assert_eq!(state.shutdowns.load(Ordering::Relaxed), 1); + assert_eq!(state.force_flushes.load(Ordering::Relaxed), 0); +} + +#[tokio::test(flavor = "current_thread")] +async fn bounded_shutdown_completes_on_current_thread_runtime() { + assert_bounded_shutdown_completes().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn bounded_shutdown_completes_on_multi_thread_runtime() { + assert_bounded_shutdown_completes().await; } #[test] fn explicit_shutdown_and_drop_shut_down_exporters_once_without_force_flush() { - let (provider, state) = test_provider(); + let TestProvider { + provider, + state, + started, + completed, + } = test_provider(ShutdownBehavior::Complete); provider.shutdown(); provider.shutdown(); drop(provider); + started + .recv_timeout(Duration::from_secs(/*secs*/ 1)) + .expect("shutdown started"); + completed + .recv_timeout(Duration::from_secs(/*secs*/ 1)) + .expect("shutdown completed"); 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(); + let TestProvider { + provider, + state, + started, + completed, + } = test_provider(ShutdownBehavior::Complete); drop(provider); + started + .recv_timeout(Duration::from_secs(/*secs*/ 1)) + .expect("shutdown started"); + completed + .recv_timeout(Duration::from_secs(/*secs*/ 1)) + .expect("shutdown completed"); assert_eq!(state.shutdowns.load(Ordering::Relaxed), 1); assert_eq!(state.force_flushes.load(Ordering::Relaxed), 0); } diff --git a/codex-rs/otel/tests/suite/otlp_http_loopback.rs b/codex-rs/otel/tests/suite/otlp_http_loopback.rs index 01bd6c25db..d9316a269d 100644 --- a/codex-rs/otel/tests/suite/otlp_http_loopback.rs +++ b/codex-rs/otel/tests/suite/otlp_http_loopback.rs @@ -529,7 +529,7 @@ fn otlp_http_exporter_sends_traces_to_collector() } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn otlp_http_exporter_sends_traces_to_collector_in_tokio_runtime() +async fn otlp_http_exporter_sends_traces_to_collector_with_bounded_shutdown_in_tokio_runtime() -> std::result::Result<(), Box> { let _trace_context_config_guard = TRACE_CONTEXT_CONFIG_LOCK .lock() @@ -598,7 +598,10 @@ async fn otlp_http_exporter_sends_traces_to_collector_in_tokio_runtime() let _guard = span.enter(); tracing::info!("trace loopback event from tokio runtime"); }); - otel.shutdown(); + tokio::task::block_in_place(|| { + tokio::runtime::Handle::current() + .block_on(otel.shutdown_with_timeout(Duration::from_secs(/*secs*/ 2))) + })?; server.join().expect("server join"); let captured = rx.recv_timeout(Duration::from_secs(1)).expect("captured"); @@ -631,6 +634,84 @@ async fn otlp_http_exporter_sends_traces_to_collector_in_tokio_runtime() Ok(()) } +#[test] +fn otlp_http_exporter_times_out_when_collector_stalls_during_bounded_shutdown() { + let _trace_context_config_guard = TRACE_CONTEXT_CONFIG_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let addr = listener.local_addr().expect("local_addr"); + let (request_started_tx, request_started_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept trace request"); + let (path, _, _) = read_http_request(&mut stream).expect("read trace request"); + request_started_tx.send(path).expect("request started"); + release_rx + .recv_timeout(Duration::from_secs(/*secs*/ 2)) + .expect("collector released"); + write_http_response(&mut stream, "202 Accepted").expect("write trace response"); + }); + + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("build tokio runtime"); + let (result, elapsed) = runtime.block_on(async move { + let otel = OtelProvider::from(&OtelSettings { + environment: "test".to_string(), + service_name: "codex-cli".to_string(), + service_version: env!("CARGO_PKG_VERSION").to_string(), + codex_home: PathBuf::from("."), + exporter: OtelExporter::None, + trace_exporter: OtelExporter::OtlpHttp { + endpoint: format!("http://{addr}/v1/traces"), + headers: HashMap::new(), + protocol: OtelHttpProtocol::Json, + tls: None, + }, + metrics_exporter: OtelExporter::None, + runtime_metrics: false, + span_attributes: BTreeMap::new(), + tracestate: BTreeMap::new(), + }) + .expect("build otel provider") + .expect("otel provider"); + let tracing_layer = otel.tracing_layer().expect("tracing layer"); + let subscriber = tracing_subscriber::registry().with(tracing_layer); + + tracing::subscriber::with_default(subscriber, || { + let span = tracing::info_span!("trace-loopback-stalled"); + let _guard = span.enter(); + tracing::info!("trace loopback event for stalled collector"); + }); + + let started = Instant::now(); + let result = otel + .shutdown_with_timeout(Duration::from_millis(/*millis*/ 50)) + .await; + (result, started.elapsed()) + }); + + let path = request_started_rx + .recv_timeout(Duration::from_secs(/*secs*/ 1)) + .expect("trace request reached collector"); + release_tx.send(()).expect("release collector"); + server.join().expect("server join"); + runtime.shutdown_timeout(Duration::from_secs(/*secs*/ 1)); + + assert_eq!(path, "/v1/traces"); + assert_eq!( + result.as_ref().map_err(std::io::Error::kind), + Err(std::io::ErrorKind::TimedOut) + ); + assert!( + elapsed < Duration::from_secs(/*secs*/ 1), + "bounded shutdown blocked for {elapsed:?}" + ); +} + #[test] fn otlp_http_exporter_sends_traces_to_collector_in_current_thread_tokio_runtime() -> std::result::Result<(), Box> { diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 53136314a4..0b255736ae 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -75,6 +75,7 @@ use std::fs::OpenOptions; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; use std::time::Instant; pub use token_usage::TokenUsage; use tracing::error; @@ -226,6 +227,7 @@ pub use public_widgets::composer_input::ComposerInput; // (tests access modules directly within the crate) const TUI_LOG_FILE_NAME: &str = "codex-tui.log"; +const INTERACTIVE_OTEL_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(/*millis*/ 500); #[cfg(unix)] const AUTO_CONNECT_DAEMON_CONNECT_TIMEOUT: std::time::Duration = @@ -1296,7 +1298,7 @@ pub async fn run_main( .with(otel_tracing_layer) .try_init(); - run_ratatui_app( + let app_result = run_ratatui_app( cli, arg0_paths, loader_overrides, @@ -1314,7 +1316,17 @@ pub async fn run_main( environment_manager, ) .await - .map_err(|err| std::io::Error::other(err.to_string())) + .map_err(|err| std::io::Error::other(err.to_string())); + + if let Some(otel) = otel + && let Err(err) = otel + .shutdown_with_timeout(INTERACTIVE_OTEL_SHUTDOWN_TIMEOUT) + .await + { + warn!(error = %err, "failed to finish interactive telemetry shutdown"); + } + + app_result } #[allow(clippy::too_many_arguments)]