boa_engine/vm/
execution_budget.rs1use 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#[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}