Trace exec-server requests from receipt through completion (#39098)

## What changed

- Start inbound exec-server request spans when messages enter the connection queue and carry them through dispatch and response handling.
- Record request outcomes for client-handled network policy callbacks, including errors and disconnections.
- Add the `exec_server_request_queue_duration_seconds` histogram, labeled by bounded route name, while excluding synchronous route setup time.

## Testing

- Cover span lifetime and trace-parent propagation across server and client queues.
- Verify queue-duration telemetry and outcome recording for completed, rejected, and cancelled requests.

GitOrigin-RevId: ed67fe5305048bdf283a26ec874337d549e3324f
This commit is contained in:
Ben Romano
2026-08-17 22:42:38 +00:00
committed by copyberry
parent 4617d4d21d
commit fd34ad7297
11 changed files with 558 additions and 52 deletions

View File

@@ -14,12 +14,16 @@ use codex_network_proxy::NetworkPolicyDecider;
use codex_network_proxy::NetworkPolicyRequest;
use codex_network_proxy::NetworkProxyAuditMetadata;
use codex_utils_path_uri::PathUri;
use opentelemetry::trace::TracerProvider as _;
use opentelemetry_sdk::trace::InMemorySpanExporter;
use opentelemetry_sdk::trace::SdkTracerProvider;
use pretty_assertions::assert_eq;
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::time::timeout;
use tracing::instrument::WithSubscriber;
use tracing_subscriber::filter::filter_fn;
use tracing_subscriber::prelude::*;
use super::super::LazyRemoteExecServerClient;
@@ -308,6 +312,18 @@ async fn abandoned_process_start_unregisters_and_cleans_up() {
#[tokio::test]
async fn policy_requests_use_process_decider_and_cancel_on_unregister() {
let span_exporter = InMemorySpanExporter::default();
let tracer_provider = SdkTracerProvider::builder()
.with_simple_exporter(span_exporter.clone())
.build();
let subscriber = tracing_subscriber::registry().with(
tracing_opentelemetry::layer()
.with_tracer(tracer_provider.tracer("exec-server-test"))
.with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)),
);
let _subscriber = tracing::subscriber::set_default(subscriber);
tracing::callsite::rebuild_interest_cache();
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
@@ -428,6 +444,13 @@ async fn policy_requests_use_process_decider_and_cancel_on_unregister() {
let started_tx = started_tx.clone();
let dropped_tx = dropped_tx.clone();
async move {
assert_eq!(
tracing::Span::current()
.metadata()
.map(tracing::Metadata::name),
Some("codex.exec_server.request"),
"network policy decisions must run inside the inbound request span"
);
match request.host.as_str() {
"allowed.example" => NetworkDecision::Allow,
"denied.example" => NetworkDecision::deny("blocked"),
@@ -478,4 +501,40 @@ async fn policy_requests_use_process_decider_and_cancel_on_unregister() {
.await
.expect("policy routing should finish")
.expect("server task should finish");
tracer_provider.force_flush().expect("flush traces");
let spans = span_exporter.get_finished_spans().expect("span export");
let policy_spans = spans
.iter()
.filter(|span| span.name.as_ref() == NETWORK_POLICY_REQUEST_METHOD)
.collect::<Vec<_>>();
assert!(
!policy_spans.is_empty(),
"network policy requests should export server spans"
);
let outcomes = policy_spans
.iter()
.map(|span| {
span.attributes
.iter()
.find(|attribute| attribute.key.as_str() == "result")
.map(|attribute| attribute.value.as_str().into_owned())
})
.collect::<Vec<_>>();
assert!(
outcomes.iter().all(Option::is_some),
"completed, rejected, and cancelled policy requests must all record an outcome"
);
assert!(
outcomes
.iter()
.any(|outcome| outcome.as_deref() == Some("success")),
"completed and capacity-rejected requests should record successful responses"
);
assert!(
outcomes
.iter()
.any(|outcome| outcome.as_deref() == Some("disconnected")),
"cancelled requests should record disconnection"
);
}

View File

@@ -16,6 +16,7 @@ use tokio::time::sleep;
use tokio::time::timeout;
use tokio::time::timeout_at;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use tracing::debug;
use super::ConnectionStatus;
@@ -63,6 +64,23 @@ const REGISTRY_RECOVERY_INITIAL_RETRY_INTERVAL: Duration = Duration::from_millis
const REGISTRY_RECOVERY_MAX_RETRY_INTERVAL: Duration = Duration::from_secs(5);
const NETWORK_POLICY_DENIAL_REASON: &str = "not_allowed";
struct ClientRequestOutcome {
span: tracing::Span,
result: &'static str,
}
impl ClientRequestOutcome {
fn complete(&mut self, result: &'static str) {
self.result = result;
}
}
impl Drop for ClientRequestOutcome {
fn drop(&mut self) {
self.span.record("result", self.result);
}
}
impl SessionState {
fn last_published_seq(&self) -> u64 {
self.ordered_events
@@ -563,7 +581,14 @@ impl ExecServerClient {
return;
};
match event {
RpcClientEvent::Request(request) => {
RpcClientEvent::Request {
request,
request_span,
} => {
let mut request_outcome = ClientRequestOutcome {
span: request_span,
result: "disconnected",
};
if request.method != NETWORK_POLICY_REQUEST_METHOD {
let error = method_not_found(format!(
"exec-server client does not implement `{}` yet",
@@ -576,8 +601,12 @@ impl ExecServerClient {
);
return;
}
request_outcome.complete("error");
continue;
}
request_outcome
.span
.record("otel.name", NETWORK_POLICY_REQUEST_METHOD);
let request_guard = match rpc_client
.admit_inbound_request(&request.id, &rpc_inbound_request_slots)
@@ -612,6 +641,7 @@ impl ExecServerClient {
);
return;
}
request_outcome.complete("success");
continue;
}
};
@@ -630,6 +660,7 @@ impl ExecServerClient {
);
return;
}
request_outcome.complete("error");
continue;
}
};
@@ -677,7 +708,8 @@ impl ExecServerClient {
let inner = Arc::downgrade(&inner);
let rpc_client = Arc::downgrade(&rpc_client);
let connection_cancelled = connection_cancelled.clone();
tokio::spawn(async move {
let task_span = request_outcome.span.clone();
let task = async move {
let _request_guard = request_guard;
let decision = match (controller, policy_request, process_cancelled) {
(Some(controller), Some(request), Some(process_cancelled)) => {
@@ -747,8 +779,11 @@ impl ExecServerClient {
?error,
"failed to send network policy decision to exec-server"
);
} else {
request_outcome.complete("success");
}
});
};
tokio::spawn(task.instrument(task_span));
}
RpcClientEvent::Notification(notification) => {
if let Err(error) = handle_server_notification(&inner, notification).await {

View File

@@ -4,10 +4,12 @@ use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;
use axum::extract::ws::Message as AxumWebSocketMessage;
use axum::extract::ws::WebSocket as AxumWebSocket;
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCRequest;
use futures::Sink;
use futures::SinkExt;
use futures::Stream;
@@ -43,8 +45,48 @@ pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30
#[derive(Debug)]
pub(crate) enum JsonRpcConnectionEvent {
Message(JSONRPCMessage),
MalformedMessage { reason: String },
Disconnected { reason: Option<String> },
QueuedRequest {
request: JSONRPCRequest,
request_span: tracing::Span,
queued_at: Instant,
},
MalformedMessage {
reason: String,
},
Disconnected {
reason: Option<String>,
},
}
impl JsonRpcConnectionEvent {
pub(crate) fn message(message: JSONRPCMessage) -> Self {
let JSONRPCMessage::Request(request) = message else {
return Self::Message(message);
};
let queued_at = Instant::now();
let request_span = tracing::info_span!(
"codex.exec_server.request",
otel.kind = "server",
otel.name = "unknown",
method = request.method.as_str(),
result = tracing::field::Empty,
);
if let Some(trace) = &request.trace
&& !codex_otel::set_parent_from_w3c_trace_context(&request_span, trace)
{
warn!(
method = request.method.as_str(),
"ignoring invalid inbound exec-server trace carrier"
);
}
Self::QueuedRequest {
request,
request_span,
queued_at,
}
}
}
#[derive(Clone)]
@@ -327,7 +369,7 @@ impl JsonRpcConnection {
match serde_json::from_str::<JSONRPCMessage>(&line) {
Ok(message) => {
if incoming_tx_for_reader
.send(JsonRpcConnectionEvent::Message(message))
.send(JsonRpcConnectionEvent::message(message))
.await
.is_err()
{
@@ -462,7 +504,7 @@ impl JsonRpcConnection {
Some(Ok(message)) => match message.parse_jsonrpc_frame() {
Ok(JsonRpcWebSocketFrame::Message(message)) => {
if incoming_tx
.send(JsonRpcConnectionEvent::Message(message))
.send(JsonRpcConnectionEvent::message(message))
.await
.is_err()
{
@@ -698,7 +740,9 @@ mod tests {
.await?
.expect("stdio connection should report the message");
match event {
JsonRpcConnectionEvent::Message(actual) => assert_eq!(actual, message),
JsonRpcConnectionEvent::QueuedRequest { request, .. } => {
assert_eq!(JSONRPCMessage::Request(request), message)
}
event => anyhow::bail!("expected JSON-RPC message, got {event:?}"),
}
@@ -806,10 +850,12 @@ mod tests {
server_websocket
.send(Message::Binary(serde_json::to_vec(&message)?.into()))
.await?;
assert!(matches!(
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
Some(JsonRpcConnectionEvent::Message(actual)) if actual == message
));
let Some(JsonRpcConnectionEvent::QueuedRequest { request, .. }) =
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?
else {
anyhow::bail!("expected a queued JSON-RPC request");
};
assert_eq!(JSONRPCMessage::Request(request), message);
drop(connection);
Ok(())

View File

@@ -77,7 +77,7 @@ impl NoiseVirtualStream {
};
for message in self.inbound_decoder.push(&plaintext)? {
self.incoming_tx
.try_send(JsonRpcConnectionEvent::Message(message))
.try_send(JsonRpcConnectionEvent::message(message))
.map_err(|_| {
ExecServerError::Protocol(
"Noise virtual stream inbound queue is full or closed".to_string(),

View File

@@ -598,7 +598,7 @@ async fn receive_data(
for message in decoder.push(&plaintext)? {
send_incoming_event(
incoming_tx,
JsonRpcConnectionEvent::Message(message),
JsonRpcConnectionEvent::message(message),
delivery_deadline,
)
.await?;

View File

@@ -380,7 +380,7 @@ where
&mut websocket,
&mut keepalive,
&incoming_tx,
JsonRpcConnectionEvent::Message(message),
JsonRpcConnectionEvent::message(message),
)
.await
{
@@ -979,10 +979,12 @@ mod tests {
.into(),
))
.await?;
assert!(matches!(
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
Some(JsonRpcConnectionEvent::Message(actual)) if actual == message
));
let Some(JsonRpcConnectionEvent::QueuedRequest { request, .. }) =
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?
else {
anyhow::bail!("expected a queued JSON-RPC request");
};
assert_eq!(JSONRPCMessage::Request(request), message);
drop(connection);
Ok(())

View File

@@ -68,9 +68,14 @@ enum RpcCallTimeout {
#[derive(Debug)]
pub(crate) enum RpcClientEvent {
Request(JSONRPCRequest),
Request {
request: JSONRPCRequest,
request_span: tracing::Span,
},
Notification(JSONRPCNotification),
Disconnected { reason: Option<String> },
Disconnected {
reason: Option<String>,
},
}
pub(crate) enum RpcInboundRequestAdmissionError {
@@ -341,6 +346,22 @@ impl RpcClient {
break None;
}
}
JsonRpcConnectionEvent::QueuedRequest {
request,
request_span,
..
} => {
if event_tx
.send(RpcClientEvent::Request {
request,
request_span,
})
.await
.is_err()
{
break None;
}
}
JsonRpcConnectionEvent::MalformedMessage { reason } => {
let _ = reason;
break None;
@@ -775,7 +796,10 @@ async fn handle_server_message(
}
JSONRPCMessage::Request(request) => {
event_tx
.send(RpcClientEvent::Request(request))
.send(RpcClientEvent::Request {
request,
request_span: tracing::Span::none(),
})
.await
.map_err(|_| "RPC client event receiver closed".to_string())?;
}
@@ -804,6 +828,7 @@ mod tests {
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCNotification;
use codex_exec_server_protocol::JSONRPCRequest;
use codex_exec_server_protocol::JSONRPCResponse;
use codex_exec_server_protocol::RequestId;
use opentelemetry::trace::TracerProvider as _;
@@ -824,6 +849,7 @@ mod tests {
use super::RESERVED_OUTBOUND_CONTROL_MESSAGES;
use super::RpcCallError;
use super::RpcClient;
use super::RpcClientEvent;
use super::RpcNotificationSender;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
@@ -879,6 +905,79 @@ mod tests {
}
}
#[tokio::test]
async fn inbound_request_span_stays_open_until_event_consumption() {
let span_exporter = InMemorySpanExporter::default();
let tracer_provider = SdkTracerProvider::builder()
.with_simple_exporter(span_exporter.clone())
.build();
let subscriber = tracing_subscriber::registry().with(
tracing_opentelemetry::layer()
.with_tracer(tracer_provider.tracer("exec-server-test"))
.with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)),
);
let _subscriber = tracing::subscriber::set_default(subscriber);
tracing::callsite::rebuild_interest_cache();
let (outgoing_tx, _outgoing_rx) = tokio::sync::mpsc::channel(/*buffer*/ 1);
let (incoming_tx, incoming_rx) = tokio::sync::mpsc::channel(/*buffer*/ 1);
let (_disconnected_tx, disconnected_rx) = tokio::sync::watch::channel(/*init*/ false);
let connection = JsonRpcConnection {
outgoing_tx,
incoming_rx,
disconnected_rx,
task_handles: Vec::new(),
transport: JsonRpcTransport::Plain,
};
let (_client, mut events_rx) = RpcClient::new(connection);
incoming_tx
.send(JsonRpcConnectionEvent::message(JSONRPCMessage::Request(
JSONRPCRequest {
id: RequestId::Integer(1),
method: "test/callback".to_string(),
params: None,
trace: None,
},
)))
.await
.expect("queue inbound client request");
timeout(Duration::from_secs(1), async {
while events_rx.is_empty() {
tokio::task::yield_now().await;
}
})
.await
.expect("inbound request should enter the client event queue");
assert!(
span_exporter
.get_finished_spans()
.expect("request span export")
.is_empty(),
"the request span must remain open until the client consumes the event"
);
let Some(RpcClientEvent::Request {
request,
request_span,
}) = events_rx.recv().await
else {
panic!("expected an inbound client request");
};
assert_eq!(request.method, "test/callback");
request_span.record("otel.name", "test/callback");
drop(request_span);
tracer_provider.force_flush().expect("flush traces");
let spans = span_exporter.get_finished_spans().expect("span export");
assert!(
spans
.iter()
.any(|span| span.name.as_ref() == "test/callback"),
"the request span should cover the complete event queue wait"
);
}
#[tokio::test]
async fn rpc_client_matches_out_of_order_responses_by_request_id() {
let (client_stdin, server_reader) = tokio::io::duplex(4096);

View File

@@ -1,4 +1,5 @@
use std::sync::Arc;
use std::time::Instant;
use codex_exec_server_protocol::JSONRPCMessage;
use tokio::sync::mpsc;
@@ -167,13 +168,26 @@ async fn run_connection(
dispatcher.handle_malformed_message(reason).await
}
JsonRpcConnectionEvent::Message(message) => match message {
JSONRPCMessage::Request(request) => dispatcher.dispatch_request(request).await,
JSONRPCMessage::Request(request) => {
dispatcher
.dispatch_request(request, tracing::Span::none(), Instant::now())
.await
}
JSONRPCMessage::Notification(notification) => {
dispatcher.handle_notification(notification).await
}
JSONRPCMessage::Response(response) => dispatcher.handle_response(response),
JSONRPCMessage::Error(error) => dispatcher.handle_error(error),
},
JsonRpcConnectionEvent::QueuedRequest {
request,
request_span,
queued_at,
} => {
dispatcher
.dispatch_request(request, request_span, queued_at)
.await
}
JsonRpcConnectionEvent::Disconnected { reason } => {
if let Some(reason) = reason {
debug!("exec-server connection disconnected: {reason}");
@@ -217,6 +231,7 @@ fn complete_queued_client_responses(
JsonRpcConnectionEvent::Message(
JSONRPCMessage::Request(_) | JSONRPCMessage::Notification(_),
)
| JsonRpcConnectionEvent::QueuedRequest { .. }
| JsonRpcConnectionEvent::MalformedMessage { .. }
| JsonRpcConnectionEvent::Disconnected { .. } => continue,
};

View File

@@ -170,11 +170,18 @@ impl RequestDispatcher {
RequestTaskResult::Completed
}
pub(super) async fn dispatch_request(&mut self, request: JSONRPCRequest) -> RequestTaskResult {
pub(super) async fn dispatch_request(
&mut self,
request: JSONRPCRequest,
request_span: tracing::Span,
queued_at: Instant,
) -> RequestTaskResult {
let started_at = Instant::now();
let Some((method, route)) = self.router.request_route(request.method.as_str()) else {
let method = "unknown";
let span = request_span(method, &request);
self.telemetry
.request_queue_completed(method, queued_at.elapsed());
request_span.record("otel.name", method);
if self
.outgoing_tx
.send(RpcServerOutboundMessage::Error {
@@ -187,27 +194,33 @@ impl RequestDispatcher {
.await
.is_err()
{
span.record("result", "disconnected");
request_span.record("result", "disconnected");
self.telemetry
.request_completed(method, "disconnected", started_at.elapsed());
return RequestTaskResult::ConnectionClosed;
}
span.record("result", "error");
request_span.record("result", "error");
self.telemetry
.request_completed(method, "error", started_at.elapsed());
return RequestTaskResult::Completed;
};
let task_span = request_span(method, &request);
request_span.record("otel.name", method);
let route_setup_started_at = Instant::now();
let route = route(Arc::clone(&self.handler), request);
let route_setup_duration = route_setup_started_at.elapsed();
let outgoing_tx = self.outgoing_tx.clone();
let mut disconnected_rx = self.disconnected_rx.clone();
let telemetry = self.telemetry.clone();
let task = async move {
telemetry.request_queue_completed(
method,
queued_at.elapsed().saturating_sub(route_setup_duration),
);
let message = tokio::select! {
message = route.instrument(task_span.clone()) => message,
message = route.instrument(request_span.clone()) => message,
_ = disconnected_rx.changed() => {
task_span.record("result", "disconnected");
request_span.record("result", "disconnected");
telemetry.request_completed(method, "disconnected", started_at.elapsed());
return RequestTaskResult::ConnectionClosed;
}
@@ -221,11 +234,11 @@ impl RequestDispatcher {
None => true,
};
if !response_sent {
task_span.record("result", "disconnected");
request_span.record("result", "disconnected");
telemetry.request_completed(method, "disconnected", started_at.elapsed());
return RequestTaskResult::ConnectionClosed;
}
task_span.record("result", result);
request_span.record("result", result);
telemetry.request_completed(method, result, started_at.elapsed());
RequestTaskResult::Completed
};
@@ -326,23 +339,6 @@ pub(super) enum RequestTaskResult {
ConnectionClosed,
}
fn request_span(span_name: &str, request: &JSONRPCRequest) -> tracing::Span {
let method = request.method.as_str();
let span = tracing::info_span!(
"codex.exec_server.request",
otel.kind = "server",
otel.name = span_name,
method,
result = tracing::field::Empty,
);
if let Some(trace) = &request.trace
&& !codex_otel::set_parent_from_w3c_trace_context(&span, trace)
{
warn!(method, "ignoring invalid inbound exec-server trace carrier");
}
span
}
fn request_result(message: &Option<RpcServerOutboundMessage>) -> &'static str {
match message {
Some(RpcServerOutboundMessage::Error { .. }) => "error",

View File

@@ -1,18 +1,42 @@
use std::sync::Arc;
use std::time::Duration;
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCRequest;
use codex_exec_server_protocol::RequestId;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_otel::MetricsClient;
use codex_otel::MetricsConfig;
use opentelemetry::trace::SpanId;
use opentelemetry::trace::TraceId;
use opentelemetry::trace::TracerProvider as _;
use opentelemetry_sdk::metrics::InMemoryMetricExporter;
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
use opentelemetry_sdk::metrics::data::MetricData;
use opentelemetry_sdk::trace::InMemorySpanExporter;
use opentelemetry_sdk::trace::SdkTracerProvider;
use pretty_assertions::assert_eq;
use tokio::sync::Notify;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::time::timeout;
use tracing_subscriber::filter::filter_fn;
use tracing_subscriber::prelude::*;
use super::ConcurrentRequestLimit;
use super::RequestDispatchMode;
use super::request_span;
use super::RequestDispatcher;
use super::RequestTaskResult;
use crate::ExecServerRuntimePaths;
use crate::connection::JsonRpcConnectionEvent;
use crate::rpc::RpcNotificationSender;
use crate::rpc::RpcRouter;
use crate::rpc::RpcServerOutboundMessage;
use crate::server::ExecServerHandler;
use crate::server::session_registry::SessionRegistry;
use crate::telemetry::ExecServerTelemetry;
/// Public limits reject values that cannot safely enable semaphore-backed concurrency.
#[test]
@@ -54,7 +78,7 @@ fn request_dispatch_mode_parses_bounded_concurrency() {
assert_eq!(max_concurrent_requests.get(), Semaphore::MAX_PERMITS);
}
/// Request spans retain the wire method and inbound trace while bounding their exported names.
/// End-to-end request spans retain the wire method and inbound trace with bounded names.
#[test]
fn request_span_uses_bounded_name_wire_method_and_inbound_trace_parent() {
let span_exporter = InMemorySpanExporter::default();
@@ -83,13 +107,26 @@ fn request_span_uses_bounded_name_wire_method_and_inbound_trace_parent() {
params: None,
trace: Some(trace),
};
let request_span = request_span("unknown", &request);
let JsonRpcConnectionEvent::QueuedRequest { request_span, .. } =
JsonRpcConnectionEvent::message(JSONRPCMessage::Request(request))
else {
panic!("requests should start a server span before dispatch");
};
assert!(
span_exporter
.get_finished_spans()
.expect("request span export")
.is_empty(),
"the request span must remain open while the request is waiting"
);
request_span.record("otel.name", "unknown");
request_span.in_scope(|| {});
drop(request_span);
});
tracer_provider.force_flush().expect("flush traces");
let spans = span_exporter.get_finished_spans().expect("span export");
assert_eq!(spans.len(), 1);
let request_span = spans
.iter()
.find(|span| span.name.as_ref() == "unknown")
@@ -105,3 +142,206 @@ fn request_span_uses_bounded_name_wire_method_and_inbound_trace_parent() {
assert_eq!(request_span.span_context.trace_id(), trace_id);
assert_eq!(request_span.parent_span_id, parent_span_id);
}
/// One end-to-end request span covers queueing and execution; the metric isolates admission wait.
#[tokio::test]
async fn request_queue_waits_for_dispatcher_admission_before_recording_telemetry() {
let span_exporter = InMemorySpanExporter::default();
let tracer_provider = SdkTracerProvider::builder()
.with_simple_exporter(span_exporter.clone())
.build();
let subscriber = tracing_subscriber::registry().with(
tracing_opentelemetry::layer()
.with_tracer(tracer_provider.tracer("exec-server-test"))
.with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)),
);
let _subscriber = tracing::subscriber::set_default(subscriber);
tracing::callsite::rebuild_interest_cache();
let metrics = MetricsClient::new(
MetricsConfig::in_memory(
"test",
"codex-exec-server",
env!("CARGO_PKG_VERSION"),
InMemoryMetricExporter::default(),
)
.with_runtime_reader(),
)
.expect("metrics client");
let telemetry = ExecServerTelemetry::new(metrics.clone());
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*buffer*/ 1);
let notifications = RpcNotificationSender::new(outgoing_tx.clone());
let requests = notifications.request_sender();
let handler = Arc::new(ExecServerHandler::new(
SessionRegistry::new(telemetry.clone()),
notifications,
ExecServerRuntimePaths::new(
std::env::current_exe().expect("current executable"),
/*codex_linux_sandbox_exe*/ None,
)
.expect("runtime paths"),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
));
let mut router = RpcRouter::new();
let execution_started = Arc::new(Notify::new());
let release_execution = Arc::new(Notify::new());
let notify_execution_started = Arc::clone(&execution_started);
let wait_for_execution_release = Arc::clone(&release_execution);
let route_setup_duration = Duration::from_millis(200);
router.request(
"test/queued",
move |_handler: Arc<ExecServerHandler>, _params: ()| {
let execution_started = Arc::clone(&notify_execution_started);
let release_execution = Arc::clone(&wait_for_execution_release);
std::thread::sleep(route_setup_duration);
async move {
execution_started.notify_one();
release_execution.notified().await;
Ok::<_, codex_exec_server_protocol::JSONRPCErrorError>(())
}
},
);
let (_disconnected_tx, disconnected_rx) = watch::channel(/*init*/ false);
let mut dispatcher = RequestDispatcher::new(
Arc::new(router),
handler,
outgoing_tx,
disconnected_rx,
requests,
telemetry,
RequestDispatchMode::Concurrent {
max_concurrent_requests: ConcurrentRequestLimit::new(
/*max_concurrent_requests*/ 2,
)
.expect("valid request limit"),
},
);
dispatcher.initialized = true;
let admission = Arc::clone(
&dispatcher
.lanes
.as_ref()
.expect("concurrent request lanes")
.ordinary,
);
let occupied_permits = admission
.acquire_many_owned(/*n*/ 2)
.await
.expect("occupy the request admission lane");
let JsonRpcConnectionEvent::QueuedRequest {
request,
request_span,
queued_at,
} = JsonRpcConnectionEvent::message(JSONRPCMessage::Request(JSONRPCRequest {
id: RequestId::Integer(1),
method: "test/queued".to_string(),
params: None,
trace: None,
}))
else {
panic!("requests should start a server span before dispatch");
};
assert!(matches!(
dispatcher
.dispatch_request(request, request_span, queued_at)
.await,
RequestTaskResult::Completed
));
tokio::task::yield_now().await;
tokio::time::sleep(Duration::from_millis(25)).await;
assert!(
span_exporter
.get_finished_spans()
.expect("request span export")
.is_empty(),
"the end-to-end request span must remain open before admission"
);
let queued_snapshot = metrics.snapshot().expect("queued metrics snapshot");
assert!(
!queued_snapshot
.scope_metrics()
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
.any(|metric| metric.name() == "exec_server_request_queue_duration_seconds"),
"queue latency must not be recorded before request admission"
);
drop(occupied_permits);
timeout(Duration::from_secs(1), execution_started.notified())
.await
.expect("queued request should execute after admission");
assert!(
span_exporter
.get_finished_spans()
.expect("executing span export")
.is_empty(),
"the same request span must remain open during execution"
);
let snapshot = metrics.snapshot().expect("metrics snapshot");
let queue_metric = snapshot
.scope_metrics()
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
.find(|metric| metric.name() == "exec_server_request_queue_duration_seconds")
.expect("request queue duration metric");
let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = queue_metric.data() else {
panic!("request queue duration should be an f64 histogram");
};
let data_point = histogram
.data_points()
.next()
.expect("request queue duration data point");
let queue_duration = Duration::from_secs_f64(data_point.sum());
assert_eq!(data_point.count(), 1);
assert!(
queue_duration >= Duration::from_millis(25),
"queue latency should include the occupied admission lane"
);
assert_eq!(
data_point
.attributes()
.find(|attribute| attribute.key.as_str() == "method")
.map(|attribute| attribute.value.as_str().into_owned()),
Some("test/queued".to_string())
);
release_execution.notify_one();
let response = timeout(Duration::from_secs(1), outgoing_rx.recv())
.await
.expect("queued request should send its response")
.expect("queued request response");
assert!(matches!(
response,
RpcServerOutboundMessage::Response {
request_id: RequestId::Integer(1),
..
}
));
assert!(matches!(
dispatcher.join_next().await,
RequestTaskResult::Completed
));
tracer_provider.force_flush().expect("flush traces");
let spans = span_exporter.get_finished_spans().expect("span export");
assert_eq!(spans.len(), 1);
let request_span = spans.first().expect("end-to-end request span");
assert_eq!(request_span.name.as_ref(), "test/queued");
let request_duration = request_span
.end_time
.duration_since(request_span.start_time)
.expect("request span should have a valid interval");
assert!(
request_duration >= Duration::from_millis(25),
"the end-to-end request span should include the admission wait"
);
assert!(
queue_duration
<= request_duration.saturating_sub(route_setup_duration) + Duration::from_millis(5),
"queue latency must exclude synchronous request decoding and route setup"
);
metrics.shutdown().expect("shutdown metrics");
}

View File

@@ -14,6 +14,9 @@ const REQUESTS_TOTAL_METRIC: &str = "exec_server_requests_total";
const REQUESTS_TOTAL_DESCRIPTION: &str = "Total number of exec-server requests.";
const REQUEST_DURATION_METRIC: &str = "exec_server_request_duration_seconds";
const REQUEST_DURATION_DESCRIPTION: &str = "Duration of exec-server requests in seconds.";
const REQUEST_QUEUE_DURATION_METRIC: &str = "exec_server_request_queue_duration_seconds";
const REQUEST_QUEUE_DURATION_DESCRIPTION: &str =
"Time exec-server requests spend queued before execution in seconds.";
const PROCESSES_ACTIVE_METRIC: &str = "exec_server_processes_active";
const PROCESSES_ACTIVE_DESCRIPTION: &str = "Number of active exec-server processes.";
const PROCESSES_FINISHED_TOTAL_METRIC: &str = "exec_server_processes_finished_total";
@@ -145,6 +148,17 @@ impl ExecServerTelemetry {
});
}
pub(crate) fn request_queue_completed(&self, method: &'static str, duration: Duration) {
self.with_inner(|inner| {
inner.duration(
REQUEST_QUEUE_DURATION_METRIC,
REQUEST_QUEUE_DURATION_DESCRIPTION,
duration,
&[("method", method)],
);
});
}
pub(crate) fn remote_registration_completed(&self, result: &'static str, duration: Duration) {
self.record_operation(REMOTE_REGISTRATION_METRICS, result, duration);
}