use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
pub const DEFAULT_TIMEOUT: Duration = Duration::from_mins(5);
pub const DEFAULT_BACKSTOP_GRACE: Duration = Duration::from_secs(1);
pub const DEFAULT_MEMORY_LIMIT: usize = 256 * 1024 * 1024;
pub const DEFAULT_STACK_SIZE: usize = 1024 * 1024;
pub const DEFAULT_GC_THRESHOLD: usize = 64 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Limits {
pub memory: usize,
pub stack: usize,
pub gc_threshold: usize,
pub timeout: Duration,
pub backstop_grace: Duration,
pub backstop_poisons: bool,
}
impl Default for Limits {
fn default() -> Self {
Self {
memory: DEFAULT_MEMORY_LIMIT,
stack: DEFAULT_STACK_SIZE,
gc_threshold: DEFAULT_GC_THRESHOLD,
timeout: DEFAULT_TIMEOUT,
backstop_grace: DEFAULT_BACKSTOP_GRACE,
backstop_poisons: true,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct RunOptions {
pub timeout: Option<Duration>,
pub memory: Option<usize>,
pub stack: Option<usize>,
pub gc_threshold: Option<usize>,
}
pub trait PauseClock: Send + Sync {
fn parked_now(&self) -> Duration;
}
#[derive(Debug, Default, Clone, Copy)]
pub struct NeverParked;
impl PauseClock for NeverParked {
fn parked_now(&self) -> Duration {
Duration::ZERO
}
}
pub(crate) struct AppliedLimits {
pub memory: AtomicUsize,
pub stack: AtomicUsize,
pub gc: AtomicUsize,
}
impl AppliedLimits {
pub fn new(limits: &Limits) -> Self {
Self {
memory: AtomicUsize::new(limits.memory),
stack: AtomicUsize::new(limits.stack),
gc: AtomicUsize::new(limits.gc_threshold),
}
}
}
pub(crate) struct TimeoutState {
epoch: Instant,
earliest_ms: AtomicU64,
armed: std::sync::Mutex<Vec<(u64, u64, u64)>>,
next_token: AtomicU64,
pub timed_out: AtomicBool,
clock: Arc<dyn PauseClock>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ArmToken(u64);
const HOST_TOKEN: u64 = 0;
impl TimeoutState {
pub const DISARMED: u64 = u64::MAX;
pub fn new(clock: Arc<dyn PauseClock>) -> Self {
Self {
epoch: Instant::now(),
earliest_ms: AtomicU64::new(Self::DISARMED),
armed: std::sync::Mutex::new(Vec::new()),
next_token: AtomicU64::new(1),
timed_out: AtomicBool::new(false),
clock,
}
}
fn deadline_ms(&self, deadline: Instant) -> u64 {
u64::try_from(deadline.saturating_duration_since(self.epoch).as_millis())
.unwrap_or(Self::DISARMED - 1)
.min(Self::DISARMED - 1)
}
fn armed(&self) -> std::sync::MutexGuard<'_, Vec<(u64, u64, u64)>> {
self.armed.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
}
fn recompute(&self, armed: &[(u64, u64, u64)]) {
let earliest = armed.iter().map(|(_, d, _)| *d).min().unwrap_or(Self::DISARMED);
self.earliest_ms.store(earliest, Ordering::Relaxed);
}
pub fn arm(&self, deadline: Instant) -> ArmToken {
let token = self.next_token.fetch_add(1, Ordering::Relaxed);
let mut armed = self.armed();
armed.push((token, self.deadline_ms(deadline), self.parked_ms()));
self.recompute(&armed);
ArmToken(token)
}
pub fn disarm(&self, token: ArmToken) {
let mut armed = self.armed();
armed.retain(|(t, _, _)| *t != token.0);
self.recompute(&armed);
}
pub fn arm_host(&self, deadline: Instant) {
let mut armed = self.armed();
armed.retain(|(t, _, _)| *t != HOST_TOKEN);
armed.push((HOST_TOKEN, self.deadline_ms(deadline), self.parked_ms()));
self.recompute(&armed);
}
pub fn disarm_host(&self) {
self.disarm(ArmToken(HOST_TOKEN));
}
fn parked_ms(&self) -> u64 {
u64::try_from(self.clock.parked_now().as_millis()).unwrap_or(u64::MAX)
}
pub fn expired(&self) -> bool {
let earliest = self.earliest_ms.load(Ordering::Relaxed);
if earliest == Self::DISARMED {
return false;
}
let elapsed = u64::try_from(self.epoch.elapsed().as_millis()).unwrap_or(u64::MAX);
if elapsed < earliest {
return false;
}
let parked_now = self.parked_ms();
let armed = self.armed();
armed.iter().any(|(_, deadline, parked_at_arm)| {
let parked = parked_now.saturating_sub(*parked_at_arm);
elapsed >= deadline.saturating_add(parked)
})
}
pub fn clock(&self) -> &Arc<dyn PauseClock> {
&self.clock
}
}
#[derive(Clone)]
pub struct Deadline(pub(crate) Arc<TimeoutState>);
impl Deadline {
pub fn arm(&self, timeout: Duration) {
self.0.arm_host(Instant::now() + timeout);
}
pub fn disarm(&self) {
self.0.disarm_host();
}
#[must_use]
pub fn force_halted(&self) -> bool {
self.0.timed_out.load(Ordering::Relaxed)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Timedout;
pub async fn run_within<F: std::future::Future>(
clock: &Arc<dyn PauseClock>,
limit: Duration,
fut: F,
) -> Result<F::Output, Timedout> {
let started = Instant::now();
let parked_before = clock.parked_now();
let deadline_now = || started + limit + clock.parked_now().saturating_sub(parked_before);
let mut fut = std::pin::pin!(fut);
loop {
let deadline = deadline_now();
tokio::select! {
out = &mut fut => return Ok(out),
() = tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)) => {
if Instant::now() >= deadline_now() {
return Err(Timedout);
}
},
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn run_within_lets_a_quick_future_through() {
let clock: Arc<dyn PauseClock> = Arc::new(NeverParked);
let out = run_within(&clock, Duration::from_secs(5), async { 7 }).await;
assert_eq!(out, Ok(7));
}
#[tokio::test]
async fn run_within_times_out_a_parked_future() {
let clock: Arc<dyn PauseClock> = Arc::new(NeverParked);
let out = run_within(&clock, Duration::from_millis(20), std::future::pending::<()>()).await;
assert_eq!(out, Err(Timedout));
}
#[test]
fn disarmed_deadline_never_expires() {
let state = TimeoutState::new(Arc::new(NeverParked));
assert!(!state.expired());
let token = state.arm(Instant::now().checked_sub(Duration::from_secs(1)).unwrap());
assert!(state.expired());
state.disarm(token);
assert!(!state.expired());
}
#[test]
fn the_earliest_of_several_budgets_wins_and_survives_the_others_ending() {
let state = TimeoutState::new(Arc::new(NeverParked));
let far = state.arm(Instant::now() + Duration::from_mins(1));
let near = state.arm(Instant::now().checked_sub(Duration::from_secs(1)).unwrap());
assert!(state.expired());
state.disarm(far);
assert!(state.expired());
state.disarm(near);
assert!(!state.expired());
state.arm_host(Instant::now().checked_sub(Duration::from_secs(1)).unwrap());
state.arm_host(Instant::now() + Duration::from_mins(1));
assert!(!state.expired());
state.disarm_host();
}
}