Skip to main content

codex_exec_server/
telemetry.rs

1use std::sync::Arc;
2use std::sync::Mutex;
3use std::time::Duration;
4use std::time::Instant;
5
6use codex_otel::MetricsClient;
7use tracing::warn;
8
9const CONNECTIONS_ACTIVE_METRIC: &str = "exec_server_connections_active";
10const CONNECTIONS_ACTIVE_DESCRIPTION: &str = "Number of active exec-server connections.";
11const CONNECTIONS_TOTAL_METRIC: &str = "exec_server_connections_total";
12const CONNECTIONS_TOTAL_DESCRIPTION: &str = "Total number of accepted exec-server connections.";
13const REQUESTS_TOTAL_METRIC: &str = "exec_server_requests_total";
14const REQUESTS_TOTAL_DESCRIPTION: &str = "Total number of exec-server requests.";
15const REQUEST_DURATION_METRIC: &str = "exec_server_request_duration_seconds";
16const REQUEST_DURATION_DESCRIPTION: &str = "Duration of exec-server requests in seconds.";
17const PROCESSES_ACTIVE_METRIC: &str = "exec_server_processes_active";
18const PROCESSES_ACTIVE_DESCRIPTION: &str = "Number of active exec-server processes.";
19const PROCESSES_FINISHED_TOTAL_METRIC: &str = "exec_server_processes_finished_total";
20const PROCESSES_FINISHED_TOTAL_DESCRIPTION: &str =
21    "Total number of finished exec-server processes.";
22const PROCESS_DURATION_METRIC: &str = "exec_server_process_duration_seconds";
23const PROCESS_DURATION_DESCRIPTION: &str = "Duration of exec-server processes in seconds.";
24const REMOTE_REGISTRATION_METRICS: OperationMetrics = OperationMetrics {
25    total_name: "exec_server_remote_registration_total",
26    total_description: "Total number of remote exec-server registration attempts.",
27    duration_name: "exec_server_remote_registration_duration_seconds",
28    duration_description: "Duration of remote exec-server registration attempts in seconds.",
29};
30const REMOTE_RENDEZVOUS_METRICS: OperationMetrics = OperationMetrics {
31    total_name: "exec_server_remote_rendezvous_connect_total",
32    total_description: "Total number of remote exec-server rendezvous connection attempts.",
33    duration_name: "exec_server_remote_rendezvous_connect_duration_seconds",
34    duration_description: "Duration of remote exec-server rendezvous connection attempts in seconds.",
35};
36const REMOTE_RECONNECTS_TOTAL_METRIC: &str = "exec_server_remote_reconnects_total";
37const REMOTE_RECONNECTS_TOTAL_DESCRIPTION: &str = "Total number of remote exec-server reconnects.";
38
39#[derive(Clone, Copy)]
40struct OperationMetrics {
41    total_name: &'static str,
42    total_description: &'static str,
43    duration_name: &'static str,
44    duration_description: &'static str,
45}
46
47#[derive(Clone, Copy)]
48pub(crate) enum ConnectionTransport {
49    Relay,
50    Stdio,
51    WebSocket,
52}
53
54impl ConnectionTransport {
55    fn metric_tag(self) -> &'static str {
56        match self {
57            Self::Relay => "relay",
58            Self::Stdio => "stdio",
59            Self::WebSocket => "websocket",
60        }
61    }
62}
63
64#[derive(Clone, Default)]
65pub struct ExecServerTelemetry {
66    inner: Option<Arc<ExecServerTelemetryInner>>,
67}
68
69struct ExecServerTelemetryInner {
70    metrics: MetricsClient,
71    active: Arc<Mutex<ActiveCounts>>,
72}
73
74#[derive(Default)]
75struct ActiveCounts {
76    relay_connections: i64,
77    stdio_connections: i64,
78    websocket_connections: i64,
79    processes: i64,
80}
81
82impl ActiveCounts {
83    fn connections(&self, transport: ConnectionTransport) -> i64 {
84        match transport {
85            ConnectionTransport::Relay => self.relay_connections,
86            ConnectionTransport::Stdio => self.stdio_connections,
87            ConnectionTransport::WebSocket => self.websocket_connections,
88        }
89    }
90}
91
92pub(crate) struct ConnectionMetricGuard {
93    telemetry: ExecServerTelemetry,
94    transport: ConnectionTransport,
95}
96
97pub(crate) struct ProcessMetricGuard {
98    telemetry: ExecServerTelemetry,
99    started_at: Instant,
100    result: &'static str,
101}
102
103impl ExecServerTelemetry {
104    pub fn new(metrics: MetricsClient) -> Self {
105        let active = Arc::new(Mutex::new(ActiveCounts::default()));
106        register_active_gauges(&metrics, &active);
107        Self {
108            inner: Some(Arc::new(ExecServerTelemetryInner { metrics, active })),
109        }
110    }
111
112    pub(crate) fn connection_started(
113        &self,
114        transport: ConnectionTransport,
115    ) -> ConnectionMetricGuard {
116        self.with_inner(|inner| {
117            inner.adjust_connection_count(transport, /*delta*/ 1);
118            inner.counter(
119                CONNECTIONS_TOTAL_METRIC,
120                CONNECTIONS_TOTAL_DESCRIPTION,
121                &[("transport", transport.metric_tag())],
122            );
123        });
124        ConnectionMetricGuard {
125            telemetry: self.clone(),
126            transport,
127        }
128    }
129
130    pub(crate) fn request_completed(
131        &self,
132        method: &'static str,
133        result: &'static str,
134        duration: Duration,
135    ) {
136        self.with_inner(|inner| {
137            let tags = [("method", method), ("result", result)];
138            inner.counter(REQUESTS_TOTAL_METRIC, REQUESTS_TOTAL_DESCRIPTION, &tags);
139            inner.duration(
140                REQUEST_DURATION_METRIC,
141                REQUEST_DURATION_DESCRIPTION,
142                duration,
143                &tags,
144            );
145        });
146    }
147
148    pub(crate) fn remote_registration_completed(&self, result: &'static str, duration: Duration) {
149        self.record_operation(REMOTE_REGISTRATION_METRICS, result, duration);
150    }
151
152    pub(crate) fn remote_rendezvous_completed(&self, result: &'static str, duration: Duration) {
153        self.record_operation(REMOTE_RENDEZVOUS_METRICS, result, duration);
154    }
155
156    pub(crate) fn remote_reconnect(&self, reason: &'static str) {
157        self.with_inner(|inner| {
158            inner.counter(
159                REMOTE_RECONNECTS_TOTAL_METRIC,
160                REMOTE_RECONNECTS_TOTAL_DESCRIPTION,
161                &[("reason", reason)],
162            );
163        });
164    }
165
166    pub(crate) fn process_started(&self) -> ProcessMetricGuard {
167        self.with_inner(|inner| {
168            inner.adjust_process_count(/*delta*/ 1);
169        });
170        ProcessMetricGuard {
171            telemetry: self.clone(),
172            started_at: Instant::now(),
173            result: "unknown",
174        }
175    }
176
177    fn process_finished(&self, result: &'static str, duration: Duration) {
178        self.with_inner(|inner| {
179            inner.adjust_process_count(/*delta*/ -1);
180            inner.counter(
181                PROCESSES_FINISHED_TOTAL_METRIC,
182                PROCESSES_FINISHED_TOTAL_DESCRIPTION,
183                &[("result", result)],
184            );
185            inner.duration(
186                PROCESS_DURATION_METRIC,
187                PROCESS_DURATION_DESCRIPTION,
188                duration,
189                &[("result", result)],
190            );
191        });
192    }
193
194    fn connection_finished(&self, transport: ConnectionTransport) {
195        self.with_inner(|inner| {
196            inner.adjust_connection_count(transport, /*delta*/ -1);
197        });
198    }
199
200    fn with_inner(&self, emit: impl FnOnce(&ExecServerTelemetryInner)) {
201        if let Some(inner) = &self.inner {
202            emit(inner);
203        }
204    }
205
206    fn record_operation(
207        &self,
208        metrics: OperationMetrics,
209        result: &'static str,
210        duration: Duration,
211    ) {
212        self.with_inner(|inner| {
213            let tags = [("result", result)];
214            inner.counter(metrics.total_name, metrics.total_description, &tags);
215            inner.duration(
216                metrics.duration_name,
217                metrics.duration_description,
218                duration,
219                &tags,
220            );
221        });
222    }
223}
224
225impl Drop for ConnectionMetricGuard {
226    fn drop(&mut self) {
227        self.telemetry.connection_finished(self.transport);
228    }
229}
230
231impl ProcessMetricGuard {
232    pub(crate) fn finish(mut self, result: &'static str) {
233        self.result = result;
234    }
235}
236
237impl Drop for ProcessMetricGuard {
238    fn drop(&mut self) {
239        self.telemetry
240            .process_finished(self.result, self.started_at.elapsed());
241    }
242}
243
244impl ExecServerTelemetryInner {
245    fn active_counts(&self) -> std::sync::MutexGuard<'_, ActiveCounts> {
246        // These are independent integer counts, so a panic cannot leave a cross-field invariant
247        // half-updated. Recovering a poisoned lock preserves the last completed count update.
248        self.active
249            .lock()
250            .unwrap_or_else(std::sync::PoisonError::into_inner)
251    }
252
253    fn adjust_connection_count(&self, transport: ConnectionTransport, delta: i64) {
254        let mut active = self.active_counts();
255        let count = match transport {
256            ConnectionTransport::Relay => &mut active.relay_connections,
257            ConnectionTransport::Stdio => &mut active.stdio_connections,
258            ConnectionTransport::WebSocket => &mut active.websocket_connections,
259        };
260        *count += delta;
261    }
262
263    fn adjust_process_count(&self, delta: i64) {
264        let mut active = self.active_counts();
265        active.processes += delta;
266    }
267
268    fn counter(&self, name: &str, description: &str, tags: &[(&str, &str)]) {
269        if self
270            .metrics
271            .counter_with_description(name, description, /*inc*/ 1, tags)
272            .is_err()
273        {
274            warn!(metric = name, "failed to emit exec-server counter");
275        }
276    }
277
278    fn duration(&self, name: &str, description: &str, duration: Duration, tags: &[(&str, &str)]) {
279        if self
280            .metrics
281            .record_duration_seconds_with_description(name, description, duration, tags)
282            .is_err()
283        {
284            warn!(metric = name, "failed to emit exec-server duration");
285        }
286    }
287}
288
289fn register_active_gauges(metrics: &MetricsClient, active: &Arc<Mutex<ActiveCounts>>) {
290    for transport in [
291        ConnectionTransport::Relay,
292        ConnectionTransport::Stdio,
293        ConnectionTransport::WebSocket,
294    ] {
295        register_active_gauge(
296            metrics,
297            active,
298            CONNECTIONS_ACTIVE_METRIC,
299            CONNECTIONS_ACTIVE_DESCRIPTION,
300            &[("transport", transport.metric_tag())],
301            move |active| active.connections(transport),
302        );
303    }
304
305    register_active_gauge(
306        metrics,
307        active,
308        PROCESSES_ACTIVE_METRIC,
309        PROCESSES_ACTIVE_DESCRIPTION,
310        &[],
311        |active| active.processes,
312    );
313}
314
315fn register_active_gauge(
316    metrics: &MetricsClient,
317    active: &Arc<Mutex<ActiveCounts>>,
318    name: &str,
319    description: &str,
320    tags: &[(&str, &str)],
321    read: impl Fn(&ActiveCounts) -> i64 + Send + Sync + 'static,
322) {
323    let active = Arc::clone(active);
324    if metrics
325        .register_observable_gauge_with_description(
326            name,
327            description,
328            move || {
329                let active = active
330                    .lock()
331                    .unwrap_or_else(std::sync::PoisonError::into_inner);
332                read(&active)
333            },
334            tags,
335        )
336        .is_err()
337    {
338        warn!(metric = name, "failed to register exec-server gauge");
339    }
340}