use agentos_bridge::queue_tracker::{register_limit, TrackedLimit};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
pub(crate) const TIMEOUT_GUARD_START_ERROR_CODE: &str = "ERR_TIMEOUT_GUARD_START";
#[cfg_attr(test, allow(dead_code))]
pub(crate) const CPU_BUDGET_GUARD_START_ERROR_CODE: &str = "ERR_CPU_BUDGET_GUARD_START";
#[cfg_attr(test, allow(dead_code))]
const CPU_BUDGET_POLL_INTERVAL: Duration = Duration::from_millis(50);
#[cfg(all(unix, not(target_os = "macos")))]
#[cfg_attr(test, allow(dead_code))]
#[derive(Clone, Copy)]
pub(crate) struct ThreadCpuClock {
clockid: libc::clockid_t,
}
#[cfg(all(unix, not(target_os = "macos")))]
#[cfg_attr(test, allow(dead_code))]
pub(crate) fn current_thread_cpu_clock() -> Option<ThreadCpuClock> {
unsafe {
let mut clockid: libc::clockid_t = 0;
let rc = libc::pthread_getcpuclockid(libc::pthread_self(), &mut clockid);
if rc == 0 {
Some(ThreadCpuClock { clockid })
} else {
None
}
}
}
#[cfg(all(unix, not(target_os = "macos")))]
impl ThreadCpuClock {
#[cfg_attr(test, allow(dead_code))]
fn elapsed_ms(self) -> Option<u64> {
unsafe {
let mut ts: libc::timespec = std::mem::zeroed();
if libc::clock_gettime(self.clockid, &mut ts) == 0 {
let ms = (ts.tv_sec as i128) * 1_000 + (ts.tv_nsec as i128) / 1_000_000;
Some(ms.max(0) as u64)
} else {
None
}
}
}
}
#[cfg(target_os = "macos")]
#[cfg_attr(test, allow(dead_code))]
#[derive(Clone, Copy)]
pub(crate) struct ThreadCpuClock {
port: libc::mach_port_t,
}
#[cfg(target_os = "macos")]
#[cfg_attr(test, allow(dead_code))]
pub(crate) fn current_thread_cpu_clock() -> Option<ThreadCpuClock> {
let port = unsafe { libc::pthread_mach_thread_np(libc::pthread_self()) };
if port == 0 {
None
} else {
Some(ThreadCpuClock { port })
}
}
#[cfg(target_os = "macos")]
impl ThreadCpuClock {
#[cfg_attr(test, allow(dead_code))]
fn elapsed_ms(self) -> Option<u64> {
unsafe {
let mut info = std::mem::MaybeUninit::<libc::thread_basic_info>::zeroed();
let mut count = (std::mem::size_of::<libc::thread_basic_info>()
/ std::mem::size_of::<libc::integer_t>())
as libc::mach_msg_type_number_t;
let rc = libc::thread_info(
self.port,
libc::THREAD_BASIC_INFO as libc::thread_flavor_t,
info.as_mut_ptr() as libc::thread_info_t,
&mut count,
);
if rc != libc::KERN_SUCCESS {
return None;
}
let info = info.assume_init();
let ms = (info.user_time.seconds as i128 + info.system_time.seconds as i128) * 1_000
+ (info.user_time.microseconds as i128 + info.system_time.microseconds as i128)
/ 1_000;
Some(ms.max(0) as u64)
}
}
}
pub(crate) struct CpuBudgetGuard {
cancel_tx: Option<crossbeam_channel::Sender<()>>,
fired: Arc<AtomicBool>,
join_handle: Option<thread::JoinHandle<()>>,
}
#[cfg(unix)]
impl CpuBudgetGuard {
#[cfg_attr(test, allow(dead_code))]
pub(crate) fn new(
budget_ms: u32,
cpu_clock: ThreadCpuClock,
isolate_handle: v8::IsolateHandle,
execution_abort: crate::session::SharedExecutionAbort,
) -> Result<Self, String> {
let (cancel_tx, cancel_rx) = crossbeam_channel::bounded::<()>(1);
let fired = Arc::new(AtomicBool::new(false));
let fired_clone = Arc::clone(&fired);
let baseline_ms = cpu_clock.elapsed_ms().unwrap_or(0);
let budget_ms = budget_ms as u64;
let cpu_gauge = register_limit(TrackedLimit::V8CpuTimeMs, budget_ms as usize);
let handle = thread::Builder::new()
.name("cpu-budget".into())
.spawn(move || {
let ticker = crossbeam_channel::tick(CPU_BUDGET_POLL_INTERVAL);
loop {
crossbeam_channel::select! {
recv(cancel_rx) -> _ => {
return;
}
recv(ticker) -> _ => {
let used = cpu_clock
.elapsed_ms()
.unwrap_or(baseline_ms)
.saturating_sub(baseline_ms);
cpu_gauge.observe_depth(used as usize);
if used >= budget_ms {
fired_clone.store(true, Ordering::SeqCst);
isolate_handle.terminate_execution();
crate::session::signal_execution_abort(
&execution_abort,
crate::session::ExecutionAbortReason::CpuBudgetExceeded,
);
return;
}
}
}
}
})
.map_err(|error| {
format!(
"{CPU_BUDGET_GUARD_START_ERROR_CODE}: failed to spawn cpu-budget thread: {error}"
)
})?;
Ok(CpuBudgetGuard {
cancel_tx: Some(cancel_tx),
fired,
join_handle: Some(handle),
})
}
pub(crate) fn cancel(&mut self) {
self.cancel_tx.take();
if let Some(h) = self.join_handle.take() {
let _ = h.join();
}
}
#[cfg_attr(test, allow(dead_code))]
pub(crate) fn exceeded(&self) -> bool {
self.fired.load(Ordering::SeqCst)
}
}
#[cfg(unix)]
impl Drop for CpuBudgetGuard {
fn drop(&mut self) {
self.cancel();
}
}
#[cfg(not(unix))]
#[derive(Clone, Copy)]
pub(crate) struct ThreadCpuClock;
#[cfg(not(unix))]
pub(crate) fn current_thread_cpu_clock() -> Option<ThreadCpuClock> {
None
}
#[cfg(not(unix))]
impl CpuBudgetGuard {
pub(crate) fn new(
_budget_ms: u32,
_cpu_clock: ThreadCpuClock,
_isolate_handle: v8::IsolateHandle,
_execution_abort: crate::session::SharedExecutionAbort,
) -> Result<Self, String> {
Err(format!(
"{CPU_BUDGET_GUARD_START_ERROR_CODE}: per-thread CPU clock not supported on this platform"
))
}
pub(crate) fn cancel(&mut self) {}
#[cfg_attr(test, allow(dead_code))]
pub(crate) fn exceeded(&self) -> bool {
self.fired.load(Ordering::SeqCst)
}
}
pub struct TimeoutGuard {
cancel_tx: Option<crossbeam_channel::Sender<()>>,
fired: Arc<AtomicBool>,
join_handle: Option<thread::JoinHandle<()>>,
}
impl TimeoutGuard {
pub(crate) fn new(
timeout_ms: u32,
isolate_handle: v8::IsolateHandle,
abort_tx: crossbeam_channel::Sender<()>,
) -> Result<Self, String> {
Self::spawn(timeout_ms, isolate_handle, move || {
drop(abort_tx);
})
}
#[cfg_attr(test, allow(dead_code))]
pub(crate) fn with_execution_abort(
timeout_ms: u32,
isolate_handle: v8::IsolateHandle,
execution_abort: crate::session::SharedExecutionAbort,
) -> Result<Self, String> {
Self::spawn(timeout_ms, isolate_handle, move || {
crate::session::signal_execution_abort(
&execution_abort,
crate::session::ExecutionAbortReason::WallClockTimedOut,
);
})
}
fn spawn(
timeout_ms: u32,
isolate_handle: v8::IsolateHandle,
on_timeout: impl FnOnce() + Send + 'static,
) -> Result<Self, String> {
let (cancel_tx, cancel_rx) = crossbeam_channel::bounded::<()>(1);
let fired = Arc::new(AtomicBool::new(false));
let fired_clone = Arc::clone(&fired);
let wall_gauge = register_limit(TrackedLimit::V8WallClockMs, timeout_ms as usize);
let warn_at_ms =
timeout_ms as u64 * agentos_bridge::queue_tracker::WARN_FILL_PERCENT as u64 / 100;
let handle = thread::Builder::new()
.name("timeout".into())
.spawn(move || {
let warn_timer = crossbeam_channel::after(Duration::from_millis(warn_at_ms));
let timer = crossbeam_channel::after(Duration::from_millis(timeout_ms as u64));
loop {
crossbeam_channel::select! {
recv(warn_timer) -> _ => {
wall_gauge.observe_depth(warn_at_ms as usize);
}
recv(timer) -> _ => {
fired_clone.store(true, Ordering::SeqCst);
isolate_handle.terminate_execution();
on_timeout();
return;
}
recv(cancel_rx) -> _ => {
return;
}
}
}
})
.map_err(|error| {
format!("{TIMEOUT_GUARD_START_ERROR_CODE}: failed to spawn timeout thread: {error}")
})?;
Ok(TimeoutGuard {
cancel_tx: Some(cancel_tx),
fired,
join_handle: Some(handle),
})
}
pub fn cancel(&mut self) {
self.cancel_tx.take();
if let Some(h) = self.join_handle.take() {
let _ = h.join();
}
}
pub fn timed_out(&self) -> bool {
self.fired.load(Ordering::SeqCst)
}
}
impl Drop for TimeoutGuard {
fn drop(&mut self) {
self.cancel();
}
}
#[cfg(test)]
mod tests {
#[test]
fn timeout_guard_cancel_before_fire() {
let (abort_tx, abort_rx) = crossbeam_channel::bounded::<()>(0);
drop(abort_tx);
drop(abort_rx);
}
#[test]
fn timeout_guard_fires_on_expiry() {
}
}