Skip to main content

dope_core/driver/schedule/timer/
mod.rs

1use std::{cell, cmp, marker, mem, num, time};
2
3use o3::cell::region;
4
5mod heap;
6mod overflow;
7mod registration;
8pub use registration::Registration;
9
10pub(super) enum Lane {}
11
12use crate::driver::{
13    self,
14    schedule::{self, credits, ready::completion},
15    settings,
16};
17
18const NIL: u32 = u32::MAX;
19
20/// An absolute monotonic deadline confined to one driver lifetime and thread.
21#[derive(Clone, Copy, PartialEq, Eq)]
22#[repr(transparent)]
23pub struct Deadline<'d> {
24    at: time::Instant,
25    _driver: marker::PhantomData<fn(&'d ()) -> &'d ()>,
26    _thread: marker::PhantomData<*mut ()>,
27}
28
29const _: () = {
30    assert!(mem::size_of::<Deadline<'static>>() == mem::size_of::<time::Instant>());
31    assert!(mem::align_of::<Deadline<'static>>() == mem::align_of::<time::Instant>());
32};
33
34impl<'d> Deadline<'d> {
35    pub(in crate::driver) const fn new(at: time::Instant) -> Self {
36        Self {
37            at,
38            _driver: marker::PhantomData,
39            _thread: marker::PhantomData,
40        }
41    }
42
43    pub fn checked_add(self, duration: time::Duration) -> Option<Self> {
44        self.at.checked_add(duration).map(Self::new)
45    }
46}
47
48impl PartialOrd for Deadline<'_> {
49    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
50        Some(self.cmp(other))
51    }
52}
53
54impl Ord for Deadline<'_> {
55    fn cmp(&self, other: &Self) -> cmp::Ordering {
56        self.at.cmp(&other.at)
57    }
58}
59
60#[derive(Clone, Copy, PartialEq, Eq)]
61struct Ticket {
62    slot: u32,
63    epoch: u32,
64    _thread: o3::ThreadBound,
65}
66
67#[derive(Clone, Copy, PartialEq, Eq)]
68enum State {
69    Free,
70    Pending,
71    Armed,
72    Fired,
73    Canceled,
74    Retired,
75}
76
77#[derive(Clone, Copy)]
78enum Phase {
79    Overflow,
80    Cancel,
81    Arm,
82    Heap,
83}
84
85impl Phase {
86    const COUNT: usize = 4;
87
88    const fn next(self) -> Self {
89        match self {
90            Self::Overflow => Self::Cancel,
91            Self::Cancel => Self::Arm,
92            Self::Arm => Self::Heap,
93            Self::Heap => Self::Overflow,
94        }
95    }
96}
97
98const _: () = assert!(mem::size_of::<State>() == 1);
99
100struct Slot<'d> {
101    epoch: cell::Cell<u32>,
102    state: cell::Cell<State>,
103    wake: completion::Slot<'d>,
104    deadline: cell::Cell<time::Instant>,
105    pending_min: cell::Cell<time::Instant>,
106    next: cell::Cell<u32>,
107    _driver: marker::PhantomData<fn(&'d ()) -> &'d ()>,
108}
109
110/// Address-stable timer storage shared by registrations in one driver scope.
111///
112/// ```compile_fail
113/// use dope_core::driver::schedule::timer::Timer;
114///
115/// fn require_unpin<T: Unpin>() {}
116/// require_unpin::<Timer<'static>>();
117/// ```
118pub struct Timer<'d> {
119    slots: Box<[Slot<'d>]>,
120    free: cell::Cell<u32>,
121    pending_arm: cell::Cell<u32>,
122    pending_cancel: cell::Cell<u32>,
123    phase: cell::Cell<Phase>,
124    overflow: overflow::Overflow<'d>,
125    heap: heap::Heap<'d>,
126    _pin: marker::PhantomPinned,
127}
128
129impl<'d> Timer<'d> {
130    pub fn deadline_after(&self, duration: time::Duration) -> Option<Deadline<'d>> {
131        time::Instant::now()
132            .checked_add(duration)
133            .map(Deadline::new)
134    }
135
136    pub(in crate::driver) fn with_capacity(
137        limit: settings::ScheduleCapacity,
138        token: &region::Token<'d>,
139    ) -> Self {
140        let mut capacity = num::NonZeroUsize::new(limit.get());
141        while let Some(current) = capacity {
142            use o3::collections::BoxSliceExt;
143
144            let cap = current.get();
145            let filler = time::Instant::now();
146            let slots = BoxSliceExt::try_box_with(cap, |index| Slot {
147                epoch: cell::Cell::new(0),
148                state: cell::Cell::new(State::Free),
149                wake: completion::Slot::empty(),
150                deadline: cell::Cell::new(filler),
151                pending_min: cell::Cell::new(filler),
152                next: cell::Cell::new(if index == 0 { NIL } else { index as u32 - 1 }),
153                _driver: marker::PhantomData,
154            });
155            if let Ok(slots) = slots
156                && let Ok(heap) = heap::Heap::try_new(token, current)
157            {
158                return Self {
159                    slots,
160                    free: cell::Cell::new(cap as u32 - 1),
161                    pending_arm: cell::Cell::new(NIL),
162                    pending_cancel: cell::Cell::new(NIL),
163                    phase: cell::Cell::new(Phase::Overflow),
164                    overflow: overflow::Overflow::new(),
165                    heap,
166                    _pin: marker::PhantomPinned,
167                };
168            }
169            capacity = num::NonZeroUsize::new(cap / 2);
170        }
171        Self {
172            slots: Box::default(),
173            free: cell::Cell::new(NIL),
174            pending_arm: cell::Cell::new(NIL),
175            pending_cancel: cell::Cell::new(NIL),
176            phase: cell::Cell::new(Phase::Overflow),
177            overflow: overflow::Overflow::new(),
178            heap: heap::Heap::empty(token),
179            _pin: marker::PhantomPinned,
180        }
181    }
182
183    fn waiter(&self) -> overflow::Waiter<'_, 'd> {
184        use crate::driver::schedule::timer::overflow::Waiter;
185        Waiter::new(self)
186    }
187
188    pub fn capacity(&self) -> usize {
189        self.slots.len()
190    }
191
192    fn has_free_slot(&self) -> bool {
193        self.free.get() != NIL
194    }
195
196    fn try_arm(
197        &self,
198        deadline: time::Instant,
199        wake: completion::Waker<'d>,
200    ) -> Result<Ticket, completion::Waker<'d>> {
201        use o3::ThreadBound;
202        let slot = self.free.get();
203        if slot == NIL {
204            return Err(wake);
205        }
206        let s = &self.slots[slot as usize];
207        self.free.set(s.next.get());
208        s.state.set(State::Pending);
209        s.wake.set(wake);
210        s.deadline.set(deadline);
211        let next = self.pending_arm.get();
212        let pending_min = if next == NIL {
213            deadline
214        } else {
215            deadline.min(self.slots[next as usize].pending_min.get())
216        };
217        s.pending_min.set(pending_min);
218        s.next.set(next);
219        self.pending_arm.set(slot);
220        Ok(Ticket {
221            slot,
222            epoch: s.epoch.get(),
223            _thread: ThreadBound::NEW,
224        })
225    }
226
227    fn is_fired(&self, ticket: Ticket) -> bool {
228        match self.slots.get(ticket.slot as usize) {
229            Some(s) => s.epoch.get() == ticket.epoch && s.state.get() == State::Fired,
230            None => false,
231        }
232    }
233
234    fn replace_waker(&self, ticket: Ticket, wake: completion::Waker<'d>) {
235        if let Some(s) = self.slots.get(ticket.slot as usize)
236            && s.epoch.get() == ticket.epoch
237            && matches!(s.state.get(), State::Pending | State::Armed)
238        {
239            s.wake.set(wake);
240        }
241    }
242
243    fn cancel(&self, ticket: Ticket) -> bool {
244        let Some(s) = self.slots.get(ticket.slot as usize) else {
245            return false;
246        };
247        if s.epoch.get() != ticket.epoch {
248            return false;
249        }
250        match s.state.get() {
251            State::Fired => self.release(ticket.slot),
252            State::Pending => s.state.set(State::Canceled),
253            State::Armed => {
254                s.state.set(State::Canceled);
255                s.next.set(self.pending_cancel.get());
256                self.pending_cancel.set(ticket.slot);
257            }
258            State::Free | State::Canceled | State::Retired => return false,
259        }
260        true
261    }
262
263    fn release(&self, slot: u32) {
264        let s = &self.slots[slot as usize];
265        s.wake.clear();
266        let Some(epoch) = s.epoch.get().checked_add(1) else {
267            s.state.set(State::Retired);
268            s.next.set(NIL);
269            return;
270        };
271        s.epoch.set(epoch);
272        s.state.set(State::Free);
273        s.next.set(self.free.get());
274        self.free.set(slot);
275        self.overflow.wake_min();
276    }
277
278    fn flush_cancel(
279        &self,
280        token: &mut region::Token<'d>,
281        budget: &mut credits::Budget<'_, 'd, Lane>,
282    ) {
283        let mut slot = self.pending_cancel.replace(NIL);
284        while slot != NIL && budget.take() {
285            let s = &self.slots[slot as usize];
286            let next = s.next.get();
287            if s.state.get() == State::Canceled {
288                self.heap.remove(token, slot as usize);
289                self.release(slot);
290            }
291            slot = next;
292        }
293        if slot != NIL {
294            self.pending_cancel.set(slot);
295        }
296    }
297
298    fn flush_arm(&self, token: &mut region::Token<'d>, budget: &mut credits::Budget<'_, 'd, Lane>) {
299        let mut slot = self.pending_arm.replace(NIL);
300        while slot != NIL && budget.take() {
301            let s = &self.slots[slot as usize];
302            let next = s.next.get();
303            match s.state.get() {
304                State::Canceled => self.release(slot),
305                State::Pending => {
306                    let key = heap::Key {
307                        deadline: s.deadline.get(),
308                        epoch: s.epoch.get(),
309                    };
310                    self.heap.insert(token, slot as usize, key);
311                    s.state.set(State::Armed);
312                }
313                State::Free | State::Armed | State::Fired | State::Retired => {}
314            }
315            slot = next;
316        }
317        if slot != NIL {
318            self.pending_arm.set(slot);
319        }
320    }
321
322    fn expire_heap(
323        &self,
324        token: &mut region::Token<'d>,
325        now: time::Instant,
326        budget: &mut credits::Budget<'_, 'd, Lane>,
327    ) {
328        while self.heap.peek(token).is_some_and(|top| top.deadline <= now) {
329            if !budget.take() {
330                break;
331            }
332            let Some((slot, key)) = self.heap.pop(token) else {
333                break;
334            };
335            let Some(s) = self.slots.get(slot) else {
336                continue;
337            };
338            if s.epoch.get() != key.epoch || s.state.get() != State::Armed {
339                continue;
340            }
341            s.state.set(State::Fired);
342            s.wake.wake();
343        }
344    }
345
346    pub fn earliest(&self, token: &region::Token<'d>) -> Option<time::Instant> {
347        let mut min = self.heap.peek(token).map(|key| key.deadline);
348        let pending = self.pending_arm.get();
349        if pending != NIL {
350            let deadline = self.slots[pending as usize].pending_min.get();
351            min = Some(min.map_or(deadline, |current| current.min(deadline)));
352        }
353        if let Some(deadline) = self.overflow.min_deadline() {
354            min = Some(min.map_or(deadline, |current| current.min(deadline)));
355        }
356        min
357    }
358
359    pub fn expire(
360        &self,
361        work: schedule::Timers<'_, 'd>,
362        driver: &mut driver::Context<'_, 'd>,
363        now: time::Instant,
364    ) {
365        self.expire_with(driver.region_token(), now, work);
366    }
367
368    fn expire_with(
369        &self,
370        token: &mut region::Token<'d>,
371        now: time::Instant,
372        work: schedule::Timers<'_, 'd>,
373    ) {
374        let mut budget = credits::Budget::from_timers(work);
375        let mut phase = self.phase.get();
376        for _ in 0..Phase::COUNT {
377            self.phase.set(phase.next());
378            if budget.remaining() == 0 {
379                return;
380            }
381            match phase {
382                Phase::Overflow => self.overflow.expire(now, &mut budget),
383                Phase::Cancel => self.flush_cancel(token, &mut budget),
384                Phase::Arm => self.flush_arm(token, &mut budget),
385                Phase::Heap => self.expire_heap(token, now, &mut budget),
386            }
387            phase = phase.next();
388        }
389    }
390}