#[cfg(target_arch = "wasm32")]
use std::hint;
#[cfg(not(target_arch = "wasm32"))]
use std::thread;
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
use std::time::Instant;
use std::{
cell::Cell,
error::Error,
fmt,
sync::atomic::{AtomicUsize, Ordering},
time::Duration,
};
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
use web_time::Instant;
pub const OOM_EXIT_CODE: i32 = 65;
pub static LIVE_MEMORY: AtomicUsize = AtomicUsize::new(0);
pub static BASELINE_MEMORY: AtomicUsize = AtomicUsize::new(usize::MAX);
pub fn allocate_into_baseline<T>(build: impl FnOnce() -> T) -> T {
let before = LIVE_MEMORY.load(Ordering::Relaxed);
let value = build();
let cost = LIVE_MEMORY.load(Ordering::Relaxed).saturating_sub(before);
let _ = BASELINE_MEMORY.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |baseline| {
(baseline != usize::MAX).then(|| baseline.saturating_add(cost))
});
value
}
const MEMORY_LIMIT_HEADROOM: usize = 4 * 1024 * 1024;
const TYPE_CHECK_MEMORY_LIMIT_HEADROOM: usize = 32 * 1024 * 1024;
#[must_use]
pub fn memory_limit_with_headroom(max_memory: Option<usize>, type_check: bool) -> Option<usize> {
max_memory.map(|bytes| {
let headroom = if type_check {
TYPE_CHECK_MEMORY_LIMIT_HEADROOM
} else {
MEMORY_LIMIT_HEADROOM
};
bytes.saturating_add(headroom)
})
}
pub const LARGE_RESULT_THRESHOLD: usize = 100_000;
#[derive(Debug, Clone)]
pub enum ResourceError {
Time {
scope: TimeLimitScope,
limit: Duration,
elapsed: Duration,
},
Memory { limit: usize, used: usize },
Recursion { limit: usize, depth: usize },
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TimeLimitScope {
Feed,
Turn,
}
impl TimeLimitScope {
fn prefix(self) -> &'static str {
match self {
Self::Feed => "feed ",
Self::Turn => "turn ",
}
}
}
impl fmt::Display for ResourceError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Time { scope, limit, elapsed } => {
write!(f, "{}time limit exceeded: {elapsed:?} > {limit:?}", scope.prefix())
}
Self::Memory { limit, used } => {
write!(f, "memory limit exceeded: {used} bytes > {limit} bytes")
}
Self::Recursion { .. } => {
write!(f, "maximum recursion depth exceeded")
}
}
}
}
impl Error for ResourceError {}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ResourceLimits {
pub max_feed_duration: Option<Duration>,
pub max_turn_duration: Option<Duration>,
pub max_memory: Option<usize>,
pub gc_interval: Option<usize>,
pub max_recursion_depth: usize,
pub max_suspensions: usize,
pub max_total_sleep: Option<Duration>,
}
pub const DEFAULT_MAX_RECURSION_DEPTH: usize = 1000;
pub const DEFAULT_MAX_SUSPENSIONS: usize = 1000;
impl Default for ResourceLimits {
fn default() -> Self {
Self {
max_feed_duration: None,
max_turn_duration: None,
max_memory: None,
gc_interval: None,
max_recursion_depth: DEFAULT_MAX_RECURSION_DEPTH,
max_suspensions: DEFAULT_MAX_SUSPENSIONS,
max_total_sleep: None,
}
}
}
impl ResourceLimits {
#[must_use]
pub fn max_feed_duration(mut self, limit: Duration) -> Self {
self.max_feed_duration = Some(limit);
self
}
#[must_use]
pub fn max_turn_duration(mut self, limit: Duration) -> Self {
self.max_turn_duration = Some(limit);
self
}
#[must_use]
pub fn max_memory(mut self, limit: usize) -> Self {
self.max_memory = Some(limit);
self
}
#[must_use]
pub fn gc_interval(mut self, interval: usize) -> Self {
self.gc_interval = Some(interval);
self
}
#[must_use]
pub fn max_recursion_depth(mut self, limit: usize) -> Self {
self.max_recursion_depth = limit;
self
}
#[must_use]
pub fn max_suspensions(mut self, limit: usize) -> Self {
self.max_suspensions = limit;
self
}
#[must_use]
pub fn max_total_sleep(mut self, limit: Duration) -> Self {
self.max_total_sleep = Some(limit);
self
}
}
#[derive(Debug, serde::Serialize, serde::Deserialize)]
pub struct ResourceTracker {
limits: ResourceLimits,
total_execution_time: Cell<Duration>,
feed_execution_time: Cell<Duration>,
#[serde(skip)]
turn_execution_time: Cell<Duration>,
#[serde(skip)]
running_since: Cell<Option<Instant>>,
recursion_limit_override: Cell<Option<usize>>,
}
impl Default for ResourceTracker {
fn default() -> Self {
Self::new(ResourceLimits::default())
}
}
impl ResourceTracker {
#[must_use]
pub fn new(limits: ResourceLimits) -> Self {
Self {
limits,
total_execution_time: Cell::new(Duration::ZERO),
feed_execution_time: Cell::new(Duration::ZERO),
turn_execution_time: Cell::new(Duration::ZERO),
running_since: Cell::new(None),
recursion_limit_override: Cell::new(None),
}
}
#[inline]
fn active_recursion_limit(&self) -> usize {
self.recursion_limit_override
.get()
.unwrap_or(self.limits.max_recursion_depth)
}
#[must_use]
pub fn elapsed(&self) -> Duration {
self.total_execution_time.get() + self.running_window()
}
#[must_use]
pub fn feed_elapsed(&self) -> Duration {
self.feed_execution_time.get() + self.running_window()
}
#[must_use]
pub fn turn_elapsed(&self) -> Duration {
self.turn_execution_time.get() + self.running_window()
}
#[inline]
fn running_window(&self) -> Duration {
self.running_since.get().map_or(Duration::ZERO, |t| t.elapsed())
}
#[must_use]
pub fn max_feed_duration(&self) -> Option<Duration> {
self.limits.max_feed_duration
}
#[must_use]
pub fn max_turn_duration(&self) -> Option<Duration> {
self.limits.max_turn_duration
}
#[must_use]
pub fn max_memory(&self) -> Option<usize> {
self.limits.max_memory
}
#[must_use]
pub fn max_suspensions(&self) -> usize {
self.limits.max_suspensions
}
#[must_use]
pub fn max_total_sleep(&self) -> Option<Duration> {
self.limits.max_total_sleep
}
#[must_use]
pub fn has_memory_time_limit(&self) -> bool {
self.limits.max_memory.is_some() || self.has_time_limit()
}
#[must_use]
#[inline]
pub fn has_time_limit(&self) -> bool {
self.limits.max_feed_duration.is_some() || self.limits.max_turn_duration.is_some()
}
pub fn set_max_feed_duration(&mut self, duration: Duration) {
self.limits.max_feed_duration = Some(duration);
self.on_feed_start();
}
pub fn set_max_turn_duration(&mut self, duration: Duration) {
self.limits.max_turn_duration = Some(duration);
self.on_turn_start();
}
#[inline]
pub fn check_allocation(&self, additional: usize) -> Result<(), ResourceError> {
if let Some(limit) = self.limits.max_memory {
let used = probe_memory().saturating_add(additional);
if used > limit {
return Err(ResourceError::Memory { limit, used });
}
}
Ok(())
}
#[inline]
pub fn check_memory_time(&self) -> Result<(), ResourceError> {
if let Some(limit) = self.limits.max_memory {
let used = probe_memory();
if used > limit {
return Err(ResourceError::Memory { limit, used });
}
}
self.check_time()
}
#[inline]
pub fn check_time(&self) -> Result<(), ResourceError> {
if !self.has_time_limit() {
return Ok(());
}
let running = self.running_window();
check_budget(
TimeLimitScope::Feed,
self.limits.max_feed_duration,
self.feed_execution_time.get() + running,
)
.and_then(|()| {
check_budget(
TimeLimitScope::Turn,
self.limits.max_turn_duration,
self.turn_execution_time.get() + running,
)
})
}
pub const LOOP_CHECK_INTERVAL: usize = 64;
#[inline]
pub fn check_time_every(&self, i: usize) -> Result<(), ResourceError> {
if i % Self::LOOP_CHECK_INTERVAL == Self::LOOP_CHECK_INTERVAL - 1 {
self.check_time()
} else {
Ok(())
}
}
#[inline]
pub fn check_memory_time_every(&self, i: usize) -> Result<(), ResourceError> {
if i % Self::LOOP_CHECK_INTERVAL == Self::LOOP_CHECK_INTERVAL - 1 {
self.check_memory_time()
} else {
Ok(())
}
}
#[inline]
pub fn check_growth(&self, len: usize, capacity: usize, elem_size: usize) -> Result<(), ResourceError> {
self.check_pending_allocation(Self::growth_bytes(len, capacity, elem_size))
}
#[inline]
#[must_use]
pub fn growth_bytes(len: usize, capacity: usize, elem_size: usize) -> usize {
if len < capacity {
0
} else {
let min_non_zero_capacity = match elem_size {
1 => 8,
2..=1024 => 4,
_ => 1,
};
let new_capacity = capacity
.saturating_mul(2)
.max(len.saturating_add(1))
.max(min_non_zero_capacity);
new_capacity.saturating_sub(capacity).saturating_mul(elem_size)
}
}
#[inline]
pub fn check_pending_allocation(&self, additional: usize) -> Result<(), ResourceError> {
if additional == 0 {
Ok(())
} else {
self.check_allocation(additional)
}
}
#[inline]
pub fn check_recursion_depth(&self, current_depth: usize) -> Result<(), ResourceError> {
let limit = self.active_recursion_limit();
if current_depth >= limit {
return Err(ResourceError::Recursion {
limit,
depth: current_depth + 1,
});
}
Ok(())
}
#[inline]
pub fn check_large_result(&self, estimated_bytes: usize) -> Result<(), ResourceError> {
self.check_allocation(estimated_bytes)
}
#[must_use]
#[inline]
pub fn gc_interval(&self) -> Option<usize> {
self.limits.gc_interval
}
pub fn on_execution_start(&self) {
debug_assert!(
self.running_since.get().is_none(),
"nested on_execution_start: VM-internal re-entry must use the raw run loop, not run_external"
);
self.running_since.set(Some(Instant::now()));
}
pub fn on_execution_stop(&self) {
if let Some(started) = self.running_since.take() {
let window = started.elapsed();
for clock in [
&self.total_execution_time,
&self.feed_execution_time,
&self.turn_execution_time,
] {
clock.set(clock.get() + window);
}
}
}
pub fn on_feed_start(&self) {
self.feed_execution_time.set(Duration::ZERO);
self.on_turn_start();
}
pub fn on_turn_start(&self) {
self.turn_execution_time.set(Duration::ZERO);
}
pub fn sandbox_sleep(&self, duration: Duration) {
let was_running = self.running_since.get().is_some();
self.on_execution_stop();
block_for(duration);
if was_running {
self.on_execution_start();
}
}
#[cfg(feature = "test-hooks")]
pub fn lower_recursion_limit(&self, new_limit: usize) -> Result<(), usize> {
let limit = self.active_recursion_limit();
if new_limit > limit {
return Err(limit);
}
self.recursion_limit_override.set(Some(new_limit));
Ok(())
}
}
#[inline]
fn check_budget(scope: TimeLimitScope, limit: Option<Duration>, elapsed: Duration) -> Result<(), ResourceError> {
match limit {
Some(limit) if elapsed > limit => Err(ResourceError::Time { scope, limit, elapsed }),
_ => Ok(()),
}
}
fn probe_memory() -> usize {
LIVE_MEMORY
.load(Ordering::Relaxed)
.saturating_sub(BASELINE_MEMORY.load(Ordering::Relaxed))
}
#[cfg(not(target_arch = "wasm32"))]
fn block_for(duration: Duration) {
thread::sleep(duration);
}
#[cfg(target_arch = "wasm32")]
fn block_for(duration: Duration) {
let started = Instant::now();
while started.elapsed() < duration {
hint::spin_loop();
}
}