diff --git a/codex-rs/app-server/src/otel_reload.rs b/codex-rs/app-server/src/otel_reload.rs index 86f73d8dfc..83073fe017 100644 --- a/codex-rs/app-server/src/otel_reload.rs +++ b/codex-rs/app-server/src/otel_reload.rs @@ -29,6 +29,7 @@ struct OtelReloadState { provider: Option, retired_providers: Vec, initialized_from_thread_config: bool, + pending_reload: bool, } #[derive(Clone)] @@ -61,6 +62,7 @@ impl OtelReloader { provider, retired_providers: Vec::new(), initialized_from_thread_config: false, + pending_reload: false, })), default_analytics_enabled, }, @@ -90,6 +92,9 @@ impl OtelReloader { if state.initialized_from_thread_config { return Err("app-server OTel config is already initialized from a different effective thread config; restart the app server to use a different project OTel config".into()); } + if state.pending_reload { + return Err("app-server OTel config reload is already pending for a different effective thread config".into()); + } let next_provider = codex_core::otel_init::build_provider( config, @@ -102,6 +107,7 @@ impl OtelReloader { let previous_key = state.key.clone(); state.key = Some(next_key.clone()); let old_provider = std::mem::replace(&mut state.provider, next_provider); + state.pending_reload = true; Ok(OtelReloadCommit::new( Arc::clone(&self.state), next_key, @@ -123,6 +129,7 @@ impl OtelReloader { .unwrap_or_else(std::sync::PoisonError::into_inner); state.provider.take(); state.retired_providers.clear(); + state.pending_reload = false; } } @@ -169,6 +176,7 @@ impl OtelReloadCommit { state.retired_providers.push(old_provider); } state.initialized_from_thread_config = true; + state.pending_reload = false; self.committed = true; } } @@ -193,10 +201,12 @@ impl Drop for OtelReloadCommit { std::mem::replace(&mut state.provider, self.old_provider.take()); state.key = self.previous_key.clone(); self.logger_layer.replace_provider(state.provider.as_ref()); - self.trace_layer.replace_provider(state.provider.as_ref()); + self.trace_layer.restore_provider(state.provider.as_ref()); OtelProvider::replace_global_metrics(state.provider.as_ref()); + state.pending_reload = false; } else if let Some(old_provider) = self.old_provider.take() { state.retired_providers.push(old_provider); + state.pending_reload = false; } } } @@ -383,6 +393,56 @@ mod tests { Ok(()) } + #[tokio::test(flavor = "multi_thread")] + async fn reload_from_config_rejects_different_config_while_pending() + -> Result<(), Box> { + let server = MockServer::start().await; + let codex_home = TempDir::new()?; + let initial_config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .build() + .await?; + let (_logger_layer, _trace_layer, reloader): ( + _, + OtelTraceLayer, + _, + ) = OtelReloader::new( + &initial_config, + /*provider*/ None, + /*default_analytics_enabled*/ false, + ); + + let mut first_config = initial_config.clone(); + first_config.otel.exporter = OtelExporterKind::OtlpHttp { + endpoint: format!("{}/v1/logs", server.uri()), + headers: HashMap::new(), + protocol: OtelHttpProtocol::Json, + tls: None, + }; + let first_reload = reloader.reload_from_config(&first_config)?; + + let mut second_config = initial_config; + second_config.otel.exporter = OtelExporterKind::OtlpHttp { + endpoint: format!("{}/other/logs", server.uri()), + headers: HashMap::new(), + protocol: OtelHttpProtocol::Json, + tls: None, + }; + let err = match reloader.reload_from_config(&second_config) { + Ok(_) => panic!("different effective OTel config should be rejected while pending"), + Err(err) => err, + }; + assert!( + err.to_string().contains("reload is already pending"), + "unexpected error: {err}" + ); + + drop(first_reload); + reloader.reload_from_config(&second_config)?.commit(); + reloader.shutdown(); + Ok(()) + } + async fn wait_for_otel_logs_payload(server: &MockServer) -> Result, Box> { let body = timeout(Duration::from_secs(10), async { loop { diff --git a/codex-rs/otel/src/trace_layer.rs b/codex-rs/otel/src/trace_layer.rs index a8c68266de..5d3c3c5b44 100644 --- a/codex-rs/otel/src/trace_layer.rs +++ b/codex-rs/otel/src/trace_layer.rs @@ -33,6 +33,21 @@ struct OtelTraceLayerState { active_spans: HashMap>>, } +impl OtelTraceLayerState { + fn prune_inactive_retired(&mut self) { + let active_layers = self + .active_spans + .values() + .cloned() + .collect::>>>(); + self.retired.retain(|retired| { + active_layers + .iter() + .any(|active| Arc::ptr_eq(active, retired)) + }); + } +} + impl Default for OtelTraceLayerState { fn default() -> Self { Self { @@ -46,6 +61,7 @@ impl Default for OtelTraceLayerState { #[derive(Clone)] pub struct OtelTraceLayerHandle { replace_provider: Arc, + restore_provider: Arc, shutdown: Arc, } @@ -54,6 +70,10 @@ impl OtelTraceLayerHandle { (self.replace_provider)(provider); } + pub fn restore_provider(&self, provider: Option<&OtelProvider>) { + (self.restore_provider)(provider); + } + pub fn shutdown(&self) { (self.shutdown)(); } @@ -73,11 +93,15 @@ where state: Arc::clone(&state), }; let replace_state = Arc::clone(&state); + let restore_state = Arc::clone(&state); let shutdown_state = Arc::clone(&state); let handle = OtelTraceLayerHandle { replace_provider: Arc::new(move |provider| { replace_trace_layer(&replace_state, trace_layer(provider)); }), + restore_provider: Arc::new(move |provider| { + restore_trace_layer(&restore_state, trace_layer(provider)); + }), shutdown: Arc::new(move || { let mut state = shutdown_state .write() @@ -123,10 +147,12 @@ where .state .write() .unwrap_or_else(std::sync::PoisonError::into_inner); - state + let layer = state .active_spans .remove(id) - .or_else(|| state.current.clone()) + .or_else(|| state.current.clone()); + state.prune_inactive_retired(); + layer } } @@ -211,7 +237,11 @@ where .state .read() .unwrap_or_else(std::sync::PoisonError::into_inner); - let layer = state.current.as_ref()?; + let layer = state + .current + .clone() + .or_else(|| state.active_spans.values().next().cloned()) + .or_else(|| state.retired.last().cloned())?; unsafe { Layer::::downcast_raw(layer.as_ref(), id) } } } @@ -227,6 +257,18 @@ fn replace_trace_layer( state.retired.push(current); } state.current = next_layer; + state.prune_inactive_retired(); +} + +fn restore_trace_layer( + state: &RwLock>, + restored_layer: Option>>, +) { + let mut state = state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.current = restored_layer; + state.prune_inactive_retired(); } fn trace_layer(provider: Option<&OtelProvider>) -> Option>> @@ -358,4 +400,39 @@ mod tests { assert_eq!(spans.len(), 1); assert_eq!(spans[0].events.len(), 1); } + + #[test] + fn downcast_uses_active_span_layer_after_provider_is_disabled() { + let (provider, _spans) = provider_with_tracer("first"); + let (layer, handle) = OtelTraceLayer::from_provider(Some(&provider)); + let subscriber = tracing_subscriber::registry().with(layer); + let _guard = tracing::subscriber::set_default(subscriber); + + let span = trace_span!("old_span"); + let _entered = span.enter(); + let before_disable = current_span_trace_id().expect("trace id before disable"); + + handle.replace_provider(/*provider*/ None); + let after_disable = current_span_trace_id().expect("trace id after disable"); + + assert_eq!(after_disable, before_disable); + } + + #[test] + fn restore_provider_reclaims_provisional_layer_without_active_spans() { + let (first_provider, _first_spans) = provider_with_tracer("first"); + let (second_provider, _second_spans) = provider_with_tracer("second"); + let (layer, handle): (OtelTraceLayer, _) = + OtelTraceLayer::from_provider(Some(&first_provider)); + + handle.replace_provider(Some(&second_provider)); + handle.restore_provider(Some(&first_provider)); + + let state = layer + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert!(state.current.is_some()); + assert!(state.retired.is_empty()); + } }