From d5de747260dce3109fa20b5be1c68cbb16a246b8 Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Mon, 29 Jun 2026 22:50:00 -0700 Subject: [PATCH] telemetry(exec-server): trace remote client lifecycle --- codex-rs/exec-server/src/client.rs | 128 +++++++++++++++++++-- codex-rs/exec-server/src/environment.rs | 15 ++- codex-rs/exec-server/src/remote.rs | 53 ++++++++- codex-rs/exec-server/src/remote_process.rs | 13 ++- 4 files changed, 189 insertions(+), 20 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 305b64c518..55b70acf1f 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -7,6 +7,7 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::Instant; use arc_swap::ArcSwap; use codex_exec_server_protocol::JSONRPCNotification; @@ -22,7 +23,9 @@ use tokio_util::task::AbortOnDropHandle; use tokio::time::timeout; use tracing::Instrument; +use tracing::Span; use tracing::debug; +use tracing::field; use crate::ProcessId; use crate::client_api::ExecServerClientConnectOptions; @@ -267,11 +270,17 @@ impl LazyRemoteExecServerClient { return None; } let client = self.clone(); - Some(AbortOnDropHandle::new(tokio::spawn(async move { - if let Err(error) = client.wait_until_ready().await { - debug!(%error, "exec-server environment startup failed"); + Some(AbortOnDropHandle::new(tokio::spawn( + async move { + if let Err(error) = client.wait_until_ready().await { + debug!(%error, "exec-server environment startup failed"); + } } - }))) + .instrument(tracing::info_span!( + "exec_server.environment.start_connecting", + trigger = "registration", + )), + ))) } pub(crate) fn startup_finished(&self) -> bool { @@ -282,26 +291,42 @@ impl LazyRemoteExecServerClient { self.initial_client().await.map(drop) } + #[tracing::instrument( + name = "exec_server.remote_client.get", + level = "info", + skip_all, + fields(path = field::Empty) + )] pub(crate) async fn get(&self) -> Result { if let Some(client) = self.connected_client() { + Span::current().record("path", "cached_connected"); return Ok(client); } let Some(cached_client) = self.cached_client() else { + Span::current().record("path", "initial_connect"); let client = self.initial_client().await?; if !client.is_disconnected() || !self.can_reconnect() { return Ok(client); } + Span::current().record("path", "initial_then_reconnect"); return self.reconnect().await; }; if !self.can_reconnect() { + Span::current().record("path", "cached_disconnected_nonrecoverable"); return Ok(cached_client); } + Span::current().record("path", "reconnect"); self.reconnect().await } + #[tracing::instrument( + name = "exec_server.remote_client.initial_client", + level = "info", + skip_all + )] async fn initial_client(&self) -> Result { // The first caller starts the work; every other caller waits for that same result. let result = self @@ -323,6 +348,7 @@ impl LazyRemoteExecServerClient { } } + #[tracing::instrument(name = "exec_server.remote_client.reconnect", level = "info", skip_all)] async fn reconnect(&self) -> Result { // Callers handling the same outage share one reconnect attempt. let attempt = { @@ -384,6 +410,11 @@ impl LazyRemoteExecServerClient { } } +#[tracing::instrument( + name = "exec_server.remote_client.connect_once", + level = "info", + skip_all +)] async fn connect_once(transport_params: ExecServerTransportParams) -> ConnectionResult { ExecServerClient::connect_for_transport(transport_params) .await @@ -647,31 +678,75 @@ impl ExecServerClient { self.call(FS_COPY_METHOD, ¶ms).await } + #[tracing::instrument( + name = "exec_server.client.start_process", + level = "info", + skip_all, + fields( + process_id = %params.process_id, + tty = params.tty, + connection_ready_ms = field::Empty, + session_registered_ms = field::Empty, + task_spawned_ms = field::Empty, + task_first_polled_ms = field::Empty, + response_received_ms = field::Empty, + ) + )] pub(crate) async fn start_process( &self, params: ExecParams, ) -> Result { + let start_process_started_at = Instant::now(); loop { - let rpc_client = self.inner.rpc_client().await?; + let rpc_client = self + .inner + .rpc_client() + .instrument(tracing::info_span!("exec_server.client.rpc_client_ready")) + .await?; + Span::current().record( + "connection_ready_ms", + start_process_started_at.elapsed().as_secs_f64() * 1_000.0, + ); if !self.inner.begin_process_start(&rpc_client) { continue; } let process_id = params.process_id.clone(); let state = Arc::new(SessionState::new(/*recoverable*/ false)); - if let Err(error) = self.inner.insert_session(&process_id, Arc::clone(&state)) { + let insert_result = tracing::info_span!("exec_server.client.register_process_session") + .in_scope(|| self.inner.insert_session(&process_id, Arc::clone(&state))); + if let Err(error) = insert_result { self.inner.finish_process_start(); return Err(error); } + Span::current().record( + "session_registered_ms", + start_process_started_at.elapsed().as_secs_f64() * 1_000.0, + ); let active_start = ActiveProcessStart { inner: Arc::clone(&self.inner), }; let client = self.clone(); let (result_tx, result_rx) = tokio::sync::oneshot::channel(); + let start_process_span = Span::current(); + let process_start_task_started_at = start_process_started_at; let process_start_task = async move { + start_process_span.record( + "task_first_polled_ms", + process_start_task_started_at.elapsed().as_secs_f64() * 1_000.0, + ); + tracing::info_span!( + "exec_server.client.process_start_task_first_poll", + process_id = %process_id, + ) + .in_scope(|| {}); let _active_start = active_start; match client .call_rpc::<_, ExecResponse>(&rpc_client, EXEC_METHOD, ¶ms) + .instrument(tracing::info_span!( + "exec_server.client.process_start_rpc", + process_id = %process_id, + )) .await { Ok(_) => { @@ -701,9 +776,23 @@ impl ExecServerClient { } }; tokio::spawn(process_start_task.in_current_span()); - return result_rx.await.map_err(|_| { - ExecServerError::Protocol("process start task stopped unexpectedly".to_string()) - })?; + Span::current().record( + "task_spawned_ms", + start_process_started_at.elapsed().as_secs_f64() * 1_000.0, + ); + let result = result_rx + .instrument(tracing::info_span!( + "exec_server.client.await_process_start_result" + )) + .await + .map_err(|_| { + ExecServerError::Protocol("process start task stopped unexpectedly".to_string()) + })?; + Span::current().record( + "response_received_ms", + start_process_started_at.elapsed().as_secs_f64() * 1_000.0, + ); + return result; } } @@ -970,6 +1059,19 @@ impl Session { self.state.subscribe_events() } + #[tracing::instrument( + name = "exec_server.remote.read", + level = "info", + skip_all, + fields( + process_id = %self.process_id, + after_seq, + wait_ms, + chunks = tracing::field::Empty, + exited = tracing::field::Empty, + closed = tracing::field::Empty, + ) + )] pub(crate) async fn read( &self, after_seq: Option, @@ -991,7 +1093,13 @@ impl Session { }) .await { - Ok(response) => return Ok(response), + Ok(response) => { + let span = Span::current(); + span.record("chunks", response.chunks.len()); + span.record("exited", response.exited); + span.record("closed", response.closed); + return Ok(response); + } Err(error) if is_transport_closed_error(&error) && !self.client.inner.is_failed() => { diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 66f7381103..0b0aaeddee 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -27,6 +27,7 @@ use crate::remote::NoiseRendezvousEnvironmentConfig; use crate::remote_file_system::RemoteFileSystem; use crate::remote_process::RemoteProcess; use tokio_util::task::AbortOnDropHandle; +use tracing::Instrument; pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; pub const CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR: &str = @@ -585,11 +586,17 @@ impl Environment { .unwrap_or_else(std::sync::PoisonError::into_inner); if startup_task.is_none() { let environment = Arc::clone(environment); - *startup_task = Some(AbortOnDropHandle::new(tokio::spawn(async move { - if let Err(error) = environment.wait_until_ready().await { - tracing::debug!(%error, "exec-server environment startup failed"); + *startup_task = Some(AbortOnDropHandle::new(tokio::spawn( + async move { + if let Err(error) = environment.wait_until_ready().await { + tracing::debug!(%error, "exec-server environment startup failed"); + } } - }))); + .instrument(tracing::info_span!( + "exec_server.environment.start_connecting", + trigger = "first_use", + )), + ))); } } diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 8870e0efe9..e932f284cc 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -15,6 +15,8 @@ use tokio_tungstenite::MaybeTlsStream; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::connect_async_with_config; use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tracing::Instrument; +use tracing::Span; use tracing::debug; use tracing::info; use tracing::warn; @@ -94,6 +96,8 @@ impl EnvironmentRegistryClient { otel.kind = "client", otel.name = "codex.exec_server.remote.register", result = tracing::field::Empty, + send_ms = tracing::field::Empty, + response_parse_ms = tracing::field::Empty, ) )] async fn register_environment( @@ -117,6 +121,7 @@ impl EnvironmentRegistryClient { environment_id: &str, executor_public_key: &NoiseChannelPublicKey, ) -> Result { + let send_started_at = Instant::now(); let response = self .http .post(endpoint_url( @@ -130,9 +135,22 @@ impl EnvironmentRegistryClient { executor_public_key: executor_public_key.clone(), }) .send() + .instrument(tracing::info_span!( + "codex.exec_server.remote.register.send" + )) .await?; - let response: EnvironmentRegistryRegistrationResponse = - self.parse_json_response(response).await?; + Span::current().record("send_ms", duration_ms(send_started_at.elapsed())); + let response_parse_started_at = Instant::now(); + let response: EnvironmentRegistryRegistrationResponse = self + .parse_json_response(response) + .instrument(tracing::info_span!( + "codex.exec_server.remote.register.parse_response" + )) + .await?; + Span::current().record( + "response_parse_ms", + duration_ms(response_parse_started_at.elapsed()), + ); if response.environment_id != environment_id { return Err(ExecServerError::Protocol( "environment registry returned a different environment id".to_string(), @@ -551,6 +569,9 @@ pub async fn run_remote_environment( otel.kind = "client", otel.name = "codex.exec_server.remote.rendezvous.connect", result = tracing::field::Empty, + endpoint_host = tracing::field::Empty, + request_build_ms = tracing::field::Empty, + websocket_connect_ms = tracing::field::Empty, ) )] async fn connect_rendezvous( @@ -561,20 +582,38 @@ async fn connect_rendezvous( tokio_tungstenite::tungstenite::Error, > { let started_at = Instant::now(); + if let Ok(parsed) = url.parse::() + && let Some(host) = parsed.host() + { + Span::current().record("endpoint_host", host); + } let result = async { + let request_build_started_at = Instant::now(); let mut request = url.into_client_request()?; request .headers_mut() .extend(current_trace_context_headers()); - connect_async_with_config( + Span::current().record( + "request_build_ms", + duration_ms(request_build_started_at.elapsed()), + ); + let websocket_connect_started_at = Instant::now(); + let result = connect_async_with_config( request, Some(noise_relay_websocket_config()), // Rendezvous sends small, latency-sensitive frames, so avoid Nagle's coalescing delay. /*disable_nagle*/ true, ) - .await - .map(|(websocket, _)| websocket) + .instrument(tracing::info_span!( + "codex.exec_server.remote.rendezvous.websocket_connect" + )) + .await; + Span::current().record( + "websocket_connect_ms", + duration_ms(websocket_connect_started_at.elapsed()), + ); + result.map(|(websocket, _)| websocket) } .await; let result_name = if result.is_ok() { "success" } else { "error" }; @@ -583,6 +622,10 @@ async fn connect_rendezvous( result } +fn duration_ms(duration: Duration) -> f64 { + duration.as_secs_f64() * 1_000.0 +} + fn normalize_environment_id(environment_id: String) -> Result { let environment_id = environment_id.trim().to_string(); if environment_id.is_empty() { diff --git a/codex-rs/exec-server/src/remote_process.rs b/codex-rs/exec-server/src/remote_process.rs index 72f41ae70e..ff9ff0162f 100644 --- a/codex-rs/exec-server/src/remote_process.rs +++ b/codex-rs/exec-server/src/remote_process.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use tokio::sync::watch; +use tracing::Instrument; use tracing::trace; use crate::ExecBackend; @@ -31,11 +32,21 @@ impl RemoteProcess { Self { client } } + #[tracing::instrument( + name = "exec_server.remote.start", + level = "info", + skip_all, + fields(process_id = %params.process_id, tty = params.tty) + )] async fn start( &self, params: ExecParams, ) -> Result { - let client = self.client.get().await?; + let client = self + .client + .get() + .instrument(tracing::info_span!("exec_server.remote.client_ready")) + .await?; let session = client.start_process(params).await?; Ok(StartedExecProcess {