dope_core/driver/timer/
mod.rs

1use std::cell::Cell;
2use std::marker::PhantomData;
3use std::mem::size_of;
4use std::pin::Pin;
5use std::time::Instant;
6
7mod heap;
8mod registration;
9mod starved;
10
11use heap::{Heap, Key};
12use o3::cell::RegionToken;
13use o3::marker::ThreadBound;
14pub use registration::Registration;
15use starved::Waiter;
16use starved::{Entry, Queue};
17
18use super::DriverRef;
19use super::ready::{CompletionSlot, CompletionWaker};
20
21const NIL: u32 = u32::MAX;
22
23#[derive(Clone, Copy, PartialEq, Eq)]
24struct Ticket {
25    slot: u32,
26    epoch: u32,
27    _thread: ThreadBound,
28}
29
30#[derive(Clone, Copy, PartialEq, Eq)]
31enum State {
32    Free,
33    Pending,
34    Armed,
35    Fired,
36    Canceled,
37    Retired,
38}
39
40const _: () = assert!(size_of::<State>() == 1);
41
42struct Slot<'d> {
43    epoch: Cell<u32>,
44    state: Cell<State>,
45    wake: CompletionSlot<'d>,
46    deadline: Cell<Instant>,
47    next: Cell<u32>,
48    _driver: PhantomData<fn(&'d ()) -> &'d ()>,
49}
50
51pub struct Timer<'d> {
52    slots: Box<[Slot<'d>]>,
53    free: Cell<u32>,
54    pending_arm: Cell<u32>,
55    pending_cancel: Cell<u32>,
56    waiters: Queue<'d>,
57    heap: Heap<'d>,
58}
59
60impl<'d> Timer<'d> {
61    pub(super) fn with_capacity(
62        cap: usize,
63        token: &RegionToken<'d>,
64        _driver: DriverRef<'d>,
65    ) -> Self {
66        assert!(cap <= u32::MAX as usize);
67        let filler = Instant::now();
68        Self {
69            slots: (0..cap)
70                .map(|index| Slot {
71                    epoch: Cell::new(0),
72                    state: Cell::new(State::Free),
73                    wake: CompletionSlot::empty(),
74                    deadline: Cell::new(filler),
75                    next: Cell::new(if index == 0 { NIL } else { index as u32 - 1 }),
76                    _driver: PhantomData,
77                })
78                .collect(),
79            free: Cell::new(if cap == 0 { NIL } else { cap as u32 - 1 }),
80            pending_arm: Cell::new(NIL),
81            pending_cancel: Cell::new(NIL),
82            waiters: Queue::new(),
83            heap: Heap::new(token, cap),
84        }
85    }
86
87    fn waiter(&self) -> Waiter<'_, 'd> {
88        Waiter::new(self)
89    }
90
91    pub fn capacity(&self) -> usize {
92        self.slots.len()
93    }
94
95    fn register_waiter(
96        &self,
97        waiter: Pin<&Entry<'d>>,
98        deadline: Instant,
99        wake: CompletionWaker<'d>,
100    ) {
101        self.waiters.register(waiter, deadline, wake);
102    }
103
104    fn unregister_waiter(&self, waiter: Pin<&Entry<'d>>) {
105        if self.waiters.unregister(waiter) && self.free.get() != NIL {
106            self.waiters.wake_min();
107        }
108    }
109
110    fn try_arm(
111        &self,
112        deadline: Instant,
113        wake: CompletionWaker<'d>,
114    ) -> Result<Ticket, CompletionWaker<'d>> {
115        let slot = self.free.get();
116        if slot == NIL {
117            return Err(wake);
118        }
119        let s = &self.slots[slot as usize];
120        self.free.set(s.next.get());
121        s.state.set(State::Pending);
122        s.wake.set(wake);
123        s.deadline.set(deadline);
124        s.next.set(self.pending_arm.get());
125        self.pending_arm.set(slot);
126        Ok(Ticket {
127            slot,
128            epoch: s.epoch.get(),
129            _thread: ThreadBound::NEW,
130        })
131    }
132
133    fn is_fired(&self, ticket: Ticket) -> bool {
134        match self.slots.get(ticket.slot as usize) {
135            Some(s) => s.epoch.get() == ticket.epoch && s.state.get() == State::Fired,
136            None => false,
137        }
138    }
139
140    fn replace_waker(&self, ticket: Ticket, wake: CompletionWaker<'d>) {
141        if let Some(s) = self.slots.get(ticket.slot as usize)
142            && s.epoch.get() == ticket.epoch
143            && matches!(s.state.get(), State::Pending | State::Armed)
144        {
145            s.wake.set(wake);
146        }
147    }
148
149    fn cancel(&self, ticket: Ticket) -> bool {
150        let Some(s) = self.slots.get(ticket.slot as usize) else {
151            return false;
152        };
153        if s.epoch.get() != ticket.epoch {
154            return false;
155        }
156        match s.state.get() {
157            State::Fired => self.release(ticket.slot),
158            State::Pending => s.state.set(State::Canceled),
159            State::Armed => {
160                s.state.set(State::Canceled);
161                s.next.set(self.pending_cancel.get());
162                self.pending_cancel.set(ticket.slot);
163            }
164            State::Free | State::Canceled | State::Retired => return false,
165        }
166        true
167    }
168
169    fn release(&self, slot: u32) {
170        let s = &self.slots[slot as usize];
171        s.wake.clear();
172        let Some(epoch) = s.epoch.get().checked_add(1) else {
173            s.state.set(State::Retired);
174            s.next.set(NIL);
175            return;
176        };
177        s.epoch.set(epoch);
178        s.state.set(State::Free);
179        s.next.set(self.free.get());
180        self.free.set(slot);
181        self.waiters.wake_min();
182    }
183
184    pub fn flush(&self, token: &mut RegionToken<'d>) {
185        let mut slot = self.pending_cancel.replace(NIL);
186        while slot != NIL {
187            let s = &self.slots[slot as usize];
188            let next = s.next.get();
189            if s.state.get() == State::Canceled {
190                self.heap.remove(token, slot as usize);
191                self.release(slot);
192            }
193            slot = next;
194        }
195        let mut slot = self.pending_arm.replace(NIL);
196        while slot != NIL {
197            let s = &self.slots[slot as usize];
198            let next = s.next.get();
199            match s.state.get() {
200                State::Canceled => self.release(slot),
201                State::Pending => {
202                    let key = Key {
203                        deadline: s.deadline.get(),
204                        epoch: s.epoch.get(),
205                    };
206                    self.heap.insert(token, slot as usize, key);
207                    s.state.set(State::Armed);
208                }
209                State::Free | State::Armed | State::Fired | State::Retired => {}
210            }
211            slot = next;
212        }
213    }
214
215    pub fn earliest(&self, token: &RegionToken<'d>) -> Option<Instant> {
216        let mut min = self.heap.peek(token).map(|key| key.deadline);
217        let mut slot = self.pending_arm.get();
218        while slot != NIL {
219            let s = &self.slots[slot as usize];
220            if s.state.get() == State::Pending {
221                let deadline = s.deadline.get();
222                min = Some(min.map_or(deadline, |current| current.min(deadline)));
223            }
224            slot = s.next.get();
225        }
226        if let Some(deadline) = self.waiters.min_deadline() {
227            min = Some(min.map_or(deadline, |current| current.min(deadline)));
228        }
229        min
230    }
231
232    pub fn expire(&self, token: &mut RegionToken<'d>, now: Instant) {
233        self.waiters.expire(now);
234        self.flush(token);
235        while self.heap.peek(token).is_some_and(|top| top.deadline <= now) {
236            let Some((slot, key)) = self.heap.pop(token) else {
237                break;
238            };
239            let Some(s) = self.slots.get(slot) else {
240                continue;
241            };
242            if s.epoch.get() != key.epoch || s.state.get() != State::Armed {
243                continue;
244            }
245            s.state.set(State::Fired);
246            s.wake.wake();
247        }
248    }
249}
dope_core/driver/timer/mod.rs

dope_core/driver/timer/
mod.rs
