diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 37b2d87631..3d8580a22e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2726,6 +2726,7 @@ dependencies = [ "codex-context-fragments", "codex-core-plugins", "codex-core-skills", + "codex-diagnostics", "codex-exec-server", "codex-exec-server-test-support", "codex-execpolicy", @@ -2937,6 +2938,14 @@ dependencies = [ "zip", ] +[[package]] +name = "codex-diagnostics" +version = "0.0.0" +dependencies = [ + "libc", + "pretty_assertions", +] + [[package]] name = "codex-exec" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 21e9e91f11..c3ad0d023d 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -42,6 +42,7 @@ members = [ "core-api", "core-plugins", "core-skills", + "diagnostics", "hooks", "http-client", "secrets", @@ -184,6 +185,7 @@ codex-core = { path = "core" } codex-core-api = { path = "core-api" } codex-core-plugins = { path = "core-plugins" } codex-core-skills = { path = "core-skills" } +codex-diagnostics = { path = "diagnostics" } codex-exec = { path = "exec" } codex-file-system = { path = "file-system" } codex-exec-server-protocol = { path = "exec-server-protocol" } diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 0fc5d1a865..bb9292bf0d 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -35,6 +35,7 @@ codex-context-fragments = { workspace = true } codex-config = { workspace = true } codex-core-plugins = { workspace = true } codex-core-skills = { workspace = true } +codex-diagnostics = { workspace = true } codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } codex-extension-items = { workspace = true } diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index badcdca87d..ba0bfcb741 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -9,6 +9,8 @@ use crate::session::session::Session; use crate::user_message_admission::PendingUserMessageAdmissionState; use crate::user_message_admission::UserMessageAdmission; use crate::user_message_admission::UserMessageAdmissionError; +use codex_diagnostics::Gauge; +use codex_diagnostics::GaugeGuard; use codex_exec_server::SelectedCapabilityRootsStatus; use codex_extension_api::ThreadIdleCause; use codex_features::Feature; @@ -65,6 +67,8 @@ use tokio_util::sync::CancellationToken; use codex_rollout::state_db::StateDbHandle; +static LIVE_THREADS: Gauge = Gauge::new("core.threads.live"); + #[derive(Clone, Debug)] pub struct ThreadConfigSnapshot { pub model: String, @@ -197,6 +201,7 @@ pub struct CodexThread { session_configured: SessionConfiguredEvent, rollout_path: Option, out_of_band_elicitations: Mutex, + _diagnostics_guard: GaugeGuard, } #[derive(Default)] @@ -230,6 +235,7 @@ impl CodexThread { session_configured, rollout_path, out_of_band_elicitations: Mutex::new(OutOfBandElicitations::default()), + _diagnostics_guard: LIVE_THREADS.track(), } } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index fcca8b6036..876a83d1c0 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -800,6 +800,12 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { assert_eq!(manager.list_thread_ids().await, Vec::new()); assert!(manager.get_thread(thread.thread_id).await.is_err()); + assert!( + codex_diagnostics::snapshot() + .gauges + .iter() + .any(|gauge| gauge.name == "core.threads.live" && gauge.value > 0) + ); let report = manager .shutdown_all_threads_bounded(Duration::from_secs(10)) diff --git a/codex-rs/diagnostics/BUILD.bazel b/codex-rs/diagnostics/BUILD.bazel new file mode 100644 index 0000000000..104e78d070 --- /dev/null +++ b/codex-rs/diagnostics/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "diagnostics", + crate_name = "codex_diagnostics", +) diff --git a/codex-rs/diagnostics/Cargo.toml b/codex-rs/diagnostics/Cargo.toml new file mode 100644 index 0000000000..2eb2125a10 --- /dev/null +++ b/codex-rs/diagnostics/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "codex-diagnostics" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lib] +name = "codex_diagnostics" +path = "src/lib.rs" +doctest = false + +[lints] +workspace = true + +[dependencies] +libc = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/diagnostics/src/lib.rs b/codex-rs/diagnostics/src/lib.rs new file mode 100644 index 0000000000..a966e4aee5 --- /dev/null +++ b/codex-rs/diagnostics/src/lib.rs @@ -0,0 +1,210 @@ +use std::sync::Mutex; +use std::sync::Once; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +static GAUGES: Mutex> = Mutex::new(Vec::new()); + +/// A process-wide gauge that registers itself the first time it is used. +pub struct Gauge { + name: &'static str, + value: AtomicU64, + registered: Once, +} + +impl Gauge { + /// Creates a gauge suitable for use in a `static` declaration. + pub const fn new(name: &'static str) -> Self { + Self { + name, + value: AtomicU64::new(0), + registered: Once::new(), + } + } + + /// Increments the gauge and registers it if needed. + pub fn increment(&'static self) { + self.register(); + self.value.fetch_add(1, Ordering::Relaxed); + } + + /// Decrements the gauge without allowing an underflow. + pub fn decrement(&self) { + let _ = self + .value + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| { + Some(value.saturating_sub(1)) + }); + } + + /// Increments the gauge for the lifetime of the returned guard. + pub fn track(&'static self) -> GaugeGuard { + self.increment(); + GaugeGuard { gauge: self } + } + + fn register(&'static self) { + self.registered.call_once(|| { + GAUGES + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(self); + }); + } +} + +/// Decrements its gauge when the measured object is dropped. +pub struct GaugeGuard { + gauge: &'static Gauge, +} + +impl Drop for GaugeGuard { + fn drop(&mut self) { + self.gauge.decrement(); + } +} + +/// The current value of one registered diagnostic gauge. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct GaugeSnapshot { + pub name: &'static str, + pub value: u64, +} + +/// Best-effort operating-system measurements for the current process. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ProcessSnapshot { + pub id: u32, + pub resident_memory_bytes: Option, + pub physical_footprint_bytes: Option, +} + +/// Content-free diagnostic values contributed by this process. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct DiagnosticsSnapshot { + pub process: ProcessSnapshot, + pub gauges: Vec, +} + +/// Collects built-in process measurements and every registered gauge. +pub fn snapshot() -> DiagnosticsSnapshot { + let mut gauges = GAUGES + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .map(|gauge| GaugeSnapshot { + name: gauge.name, + value: gauge.value.load(Ordering::Relaxed), + }) + .collect::>(); + gauges.sort_unstable_by_key(|gauge| gauge.name); + + DiagnosticsSnapshot { + process: process_snapshot(), + gauges, + } +} + +#[cfg(target_os = "macos")] +fn process_snapshot() -> ProcessSnapshot { + let usage = unsafe { + let mut usage = std::mem::MaybeUninit::::zeroed(); + // SAFETY: the kernel initializes this correctly sized buffer on success. + if libc::proc_pid_rusage( + libc::getpid(), + libc::RUSAGE_INFO_V0, + usage.as_mut_ptr().cast(), + ) != 0 + { + return empty_process_snapshot(); + } + usage.assume_init() + }; + + ProcessSnapshot { + id: std::process::id(), + resident_memory_bytes: Some(usage.ri_resident_size), + physical_footprint_bytes: Some(usage.ri_phys_footprint), + } +} + +#[cfg(target_os = "linux")] +fn process_snapshot() -> ProcessSnapshot { + // SAFETY: querying the system page size does not access caller-owned memory. + let page_size = u64::try_from(unsafe { libc::sysconf(libc::_SC_PAGESIZE) }) + .ok() + .filter(|page_size| *page_size > 0); + let resident_pages = std::fs::read_to_string("/proc/self/statm") + .ok() + .and_then(|statm| statm.split_whitespace().nth(1)?.parse::().ok()); + + ProcessSnapshot { + id: std::process::id(), + resident_memory_bytes: resident_pages + .zip(page_size) + .map(|(pages, page_size)| pages.saturating_mul(page_size)), + physical_footprint_bytes: None, + } +} + +#[cfg(target_os = "windows")] +fn process_snapshot() -> ProcessSnapshot { + #[repr(C)] + struct ProcessMemoryCounters { + size: u32, + page_fault_count: u32, + peak_working_set_size: usize, + working_set_size: usize, + quota_peak_paged_pool_usage: usize, + quota_paged_pool_usage: usize, + quota_peak_non_paged_pool_usage: usize, + quota_non_paged_pool_usage: usize, + pagefile_usage: usize, + peak_pagefile_usage: usize, + } + + #[link(name = "kernel32")] + unsafe extern "system" { + fn GetCurrentProcess() -> *mut std::ffi::c_void; + fn K32GetProcessMemoryInfo( + process: *mut std::ffi::c_void, + counters: *mut ProcessMemoryCounters, + size: u32, + ) -> i32; + } + + let counters = unsafe { + let mut counters = std::mem::MaybeUninit::::zeroed(); + let size = u32::try_from(std::mem::size_of::()).unwrap_or(u32::MAX); + // SAFETY: the pseudo-handle is valid and the kernel initializes this + // correctly sized writable buffer on success. + if K32GetProcessMemoryInfo(GetCurrentProcess(), counters.as_mut_ptr(), size) == 0 { + return empty_process_snapshot(); + } + counters.assume_init() + }; + + ProcessSnapshot { + id: std::process::id(), + resident_memory_bytes: u64::try_from(counters.working_set_size).ok(), + physical_footprint_bytes: None, + } +} + +#[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "windows")))] +fn process_snapshot() -> ProcessSnapshot { + empty_process_snapshot() +} + +#[cfg(not(target_os = "linux"))] +fn empty_process_snapshot() -> ProcessSnapshot { + ProcessSnapshot { + id: std::process::id(), + resident_memory_bytes: None, + physical_footprint_bytes: None, + } +} + +#[cfg(test)] +#[path = "tests.rs"] +mod tests; diff --git a/codex-rs/diagnostics/src/tests.rs b/codex-rs/diagnostics/src/tests.rs new file mode 100644 index 0000000000..d5ac2f2a5a --- /dev/null +++ b/codex-rs/diagnostics/src/tests.rs @@ -0,0 +1,58 @@ +use super::Gauge; +use super::snapshot; +use pretty_assertions::assert_eq; + +static GUARD_GAUGE: Gauge = Gauge::new("diagnostics.tests.guard"); +static REGISTERED_GAUGE: Gauge = Gauge::new("diagnostics.tests.registered"); + +#[test] +fn gauge_guards_follow_the_measured_lifetime() { + let first = GUARD_GAUGE.track(); + let second = GUARD_GAUGE.track(); + let value = || { + snapshot() + .gauges + .into_iter() + .find(|gauge| gauge.name == "diagnostics.tests.guard") + .expect("used gauge should be registered") + .value + }; + + assert_eq!(value(), 2); + drop(first); + assert_eq!(value(), 1); + drop(second); + assert_eq!(value(), 0); +} + +#[test] +fn snapshot_includes_process_memory_and_registers_gauges_once() { + REGISTERED_GAUGE.increment(); + REGISTERED_GAUGE.increment(); + let diagnostics = snapshot(); + let registered = diagnostics + .gauges + .iter() + .filter(|gauge| gauge.name == "diagnostics.tests.registered") + .collect::>(); + + assert_eq!(diagnostics.process.id, std::process::id()); + assert_eq!(registered.len(), 1); + assert_eq!(registered[0].value, 2); + #[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))] + assert!( + diagnostics + .process + .resident_memory_bytes + .is_some_and(|bytes| bytes > 0) + ); + #[cfg(target_os = "macos")] + assert!( + diagnostics + .process + .physical_footprint_bytes + .is_some_and(|bytes| bytes > 0) + ); + #[cfg(not(target_os = "macos"))] + assert_eq!(diagnostics.process.physical_footprint_bytes, None); +}