telemetry(exec-server): trace remote client lifecycle

This commit is contained in:
Anton Panasenko
2026-07-07 23:48:05 -07:00
parent f87140dc75
commit d5de747260
4 changed files with 189 additions and 20 deletions
+118 -10
View File
@@ -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<ExecServerClient, ExecServerError> {
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<ExecServerClient, ExecServerError> {
// 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<ExecServerClient, ExecServerError> {
// 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, &params).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<Session, ExecServerError> {
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, &params)
.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<u64>,
@@ -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() =>
{
+11 -4
View File
@@ -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",
)),
)));
}
}
+48 -5
View File
@@ -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<EnvironmentRegistryRegistrationResponse, ExecServerError> {
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::<http::Uri>()
&& 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<String, ExecServerError> {
let environment_id = environment_id.trim().to_string();
if environment_id.is_empty() {
+12 -1
View File
@@ -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<StartedExecProcess, crate::ExecServerError> {
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 {