mirror of
https://github.com/openai/codex.git
synced 2026-08-23 13:09:46 +00:00
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:
@@ -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"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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(())
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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(())
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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(¬ify_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");
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user