Skip to main content

boa_engine/vm/
execution_budget.rs

1use std::future::{Future, poll_fn};
2use std::sync::atomic::{AtomicBool, AtomicU16, Ordering};
3use std::sync::{Arc, Mutex, RwLock};
4use std::task::{Poll, Waker};
5use std::time::Duration;
6
7#[derive(Debug, Default)]
8pub(crate) struct ExecutionBudget {
9    duration: RwLock<Option<Duration>>,
10    active: RwLock<Option<Arc<Deadline>>>,
11    checks_remaining: AtomicU16,
12}
13
14#[derive(Debug, Default)]
15struct Deadline {
16    expired: AtomicBool,
17    cancelled: AtomicBool,
18    waker: Mutex<Option<Waker>>,
19}
20
21impl ExecutionBudget {
22    pub(crate) fn set_duration(&self, duration: Option<Duration>) {
23        assert!(
24            duration.is_none_or(|duration| !duration.is_zero()),
25            "the execution time limit must be positive"
26        );
27        *self.duration.write().expect("execution duration poisoned") = duration;
28    }
29
30    pub(crate) fn begin(self: &Arc<Self>) -> ExecutionScope {
31        let schedule = {
32            let mut active = self.active.write().expect("execution budget poisoned");
33            let duration = *self.duration.read().expect("execution duration poisoned");
34            if active.is_none() {
35                duration.map(|duration| {
36                    let deadline = Arc::new(Deadline::default());
37                    *active = Some(Arc::clone(&deadline));
38                    self.checks_remaining.store(0, Ordering::Relaxed);
39                    (deadline, duration)
40                })
41            } else {
42                None
43            }
44        };
45        let owns_deadline = schedule.is_some();
46        if let Some((deadline, duration)) = schedule {
47            drop(nagoya::spawn(async move {
48                let mut sleep = std::pin::pin!(nagoya::sleep(duration));
49                let elapsed = poll_fn(|cx| {
50                    if deadline.cancelled.load(Ordering::Acquire) {
51                        return Poll::Ready(false);
52                    }
53                    *deadline.waker.lock().expect("execution waker poisoned") =
54                        Some(cx.waker().clone());
55                    if deadline.cancelled.load(Ordering::Acquire) {
56                        return Poll::Ready(false);
57                    }
58                    sleep.as_mut().poll(cx).map(|()| true)
59                })
60                .await;
61                if elapsed && !deadline.cancelled.load(Ordering::Acquire) {
62                    deadline.expired.store(true, Ordering::Release);
63                }
64            }));
65        }
66        ExecutionScope {
67            budget: Arc::clone(self),
68            owns_deadline,
69        }
70    }
71
72    pub(crate) fn expired(&self) -> bool {
73        self.active
74            .read()
75            .expect("execution budget poisoned")
76            .as_ref()
77            .is_some_and(|deadline| deadline.expired.load(Ordering::Acquire))
78    }
79
80    #[inline]
81    pub(crate) fn should_interrupt(&self) -> bool {
82        let remaining = self.checks_remaining.load(Ordering::Relaxed);
83        if remaining != 0 {
84            self.checks_remaining
85                .store(remaining - 1, Ordering::Relaxed);
86            return false;
87        }
88        self.checks_remaining.store(1023, Ordering::Relaxed);
89        self.expired()
90    }
91}
92
93/// Nested host and VM entries share the first scope's cancellation signal.
94#[derive(Debug)]
95#[must_use = "retain the scope for the execution boundary"]
96pub struct ExecutionScope {
97    budget: Arc<ExecutionBudget>,
98    owns_deadline: bool,
99}
100
101impl Drop for ExecutionScope {
102    fn drop(&mut self) {
103        if !self.owns_deadline {
104            return;
105        }
106        let deadline = self
107            .budget
108            .active
109            .write()
110            .expect("execution budget poisoned")
111            .take();
112        if let Some(deadline) = deadline {
113            deadline.cancelled.store(true, Ordering::Release);
114            let waker = deadline
115                .waker
116                .lock()
117                .expect("execution waker poisoned")
118                .take();
119            if let Some(waker) = waker {
120                waker.wake();
121            }
122        }
123    }
124}