Skip to main content

blitz_script/
job_budget.rs

1use std::cell::Cell;
2use std::rc::Rc;
3
4use web_time::{Duration, Instant};
5
6pub struct JobBudget {
7    duration: Cell<Duration>,
8    deadline: Cell<Option<Instant>>,
9}
10
11impl Default for JobBudget {
12    fn default() -> Self {
13        Self {
14            duration: Cell::new(Duration::from_millis(16)),
15            deadline: Cell::new(None),
16        }
17    }
18}
19
20impl JobBudget {
21    pub fn duration(&self) -> Duration {
22        self.duration.get()
23    }
24
25    pub fn set_duration(&self, duration: Duration) {
26        assert!(!duration.is_zero(), "the job budget must be positive");
27        self.duration.set(duration);
28    }
29
30    pub fn begin(self: &Rc<Self>) -> JobSlice {
31        let owns_deadline = self.deadline.get().is_none();
32        if owns_deadline {
33            self.deadline
34                .set(Some(Instant::now() + self.duration.get()));
35        }
36        JobSlice {
37            budget: Rc::clone(self),
38            owns_deadline,
39        }
40    }
41
42    pub fn remaining(&self) -> Duration {
43        self.deadline.get().map_or(self.duration.get(), |deadline| {
44            deadline.saturating_duration_since(Instant::now())
45        })
46    }
47
48    pub fn exhausted(&self) -> bool {
49        self.deadline
50            .get()
51            .is_some_and(|deadline| Instant::now() >= deadline)
52    }
53}
54
55/// Nested drains share the owning poll's deadline.
56pub struct JobSlice {
57    budget: Rc<JobBudget>,
58    owns_deadline: bool,
59}
60
61impl Drop for JobSlice {
62    fn drop(&mut self) {
63        if self.owns_deadline {
64            self.budget.deadline.set(None);
65        }
66    }
67}