codex: address PR review feedback (#18725)

codex-rs/app-server/src/otel_reload.rs

@@ -29,6 +29,7 @@ struct OtelReloadState {
     provider: Option<OtelProvider>,
     retired_providers: Vec<OtelProvider>,
     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<dyn Error>> {
+        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<tracing_subscriber::Registry>,
+            _,
+        ) = 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<Vec<u8>, Box<dyn Error>> {
         let body = timeout(Duration::from_secs(10), async {
             loop {

codex-rs/otel/src/trace_layer.rs

@@ -33,6 +33,21 @@ struct OtelTraceLayerState<S> {
     active_spans: HashMap<Id, Arc<TraceLayer<S>>>,
 }
 
+impl<S> OtelTraceLayerState<S> {
+    fn prune_inactive_retired(&mut self) {
+        let active_layers = self
+            .active_spans
+            .values()
+            .cloned()
+            .collect::<Vec<Arc<TraceLayer<S>>>>();
+        self.retired.retain(|retired| {
+            active_layers
+                .iter()
+                .any(|active| Arc::ptr_eq(active, retired))
+        });
+    }
+}
+
 impl<S> Default for OtelTraceLayerState<S> {
     fn default() -> Self {
         Self {
@@ -46,6 +61,7 @@ impl<S> Default for OtelTraceLayerState<S> {
 #[derive(Clone)]
 pub struct OtelTraceLayerHandle {
     replace_provider: Arc<ReplaceProviderFn>,
+    restore_provider: Arc<ReplaceProviderFn>,
     shutdown: Arc<dyn Fn() + Send + Sync>,
 }
 
@@ -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::<S>::downcast_raw(layer.as_ref(), id) }
     }
 }
@@ -227,6 +257,18 @@ fn replace_trace_layer<S>(
         state.retired.push(current);
     }
     state.current = next_layer;
+    state.prune_inactive_retired();
+}
+
+fn restore_trace_layer<S>(
+    state: &RwLock<OtelTraceLayerState<S>>,
+    restored_layer: Option<Arc<TraceLayer<S>>>,
+) {
+    let mut state = state
+        .write()
+        .unwrap_or_else(std::sync::PoisonError::into_inner);
+    state.current = restored_layer;
+    state.prune_inactive_retired();
 }
 
 fn trace_layer<S>(provider: Option<&OtelProvider>) -> Option<Arc<TraceLayer<S>>>
@@ -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<tracing_subscriber::Registry>, _) =
+            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());
+    }
 }