Skip to main content

simoom_rt/
timer.rs

1//! A hierarchical timing wheel (Varghese and Lauck, SOSP 1987).
2//!
3//! Four levels of 64 slots; level `l` slots are `64^l` ticks wide, so the wheel spans `64^4`
4//! ticks (about 4.6 hours at the default 1 ms tick) and later deadlines wait in the last level
5//! and are re-filed as time approaches. Insert and cancel are O(1) (nodes live in a slab and form
6//! doubly linked lists per slot); finding the next deadline is a few bitmap operations; advancing
7//! touches only slots that hold timers, never every tick in between.
8
9const LEVELS: usize = 4;
10const SLOTS: usize = 64;
11const NIL: u32 = u32::MAX;
12/// The list of timers beyond the wheel's span, re-filed whenever the clock enters a new span.
13const OVERFLOW: usize = LEVELS * SLOTS;
14const SPAN_BITS: u32 = 6 * LEVELS as u32;
15/// Expired timers waiting to be popped, in deadline order; still cancellable.
16const READY: usize = OVERFLOW + 1;
17
18/// A handle to a scheduled timer, for [`TimerWheel::cancel`]. Stale handles are ignored.
19#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
20pub struct TimerId {
21    /// The node index plus one (so `Option<TimerId>` is as small as `TimerId`, K8).
22    idx: core::num::NonZeroU32,
23    gen: u32,
24}
25
26impl TimerId {
27    #[inline]
28    fn index(self) -> usize {
29        self.idx.get() as usize - 1
30    }
31}
32
33struct Node<T> {
34    deadline: u64,
35    prev: u32,
36    next: u32,
37    slot: u32,
38    gen: u32,
39    payload: Option<T>,
40}
41
42/// A hierarchical timing wheel holding payloads of type `T`. Time is in ticks; the owner
43/// decides what a tick is. See the [module docs](self).
44pub struct TimerWheel<T> {
45    now: u64,
46    nodes: Vec<Node<T>>,
47    free: Vec<u32>,
48    heads: [u32; LEVELS * SLOTS + 2],
49    ready_tail: u32,
50    occupied: [u64; LEVELS],
51    len: usize,
52}
53
54impl<T> std::fmt::Debug for TimerWheel<T> {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        f.debug_struct("TimerWheel").field("now", &self.now).field("len", &self.len).finish()
57    }
58}
59
60impl<T> Default for TimerWheel<T> {
61    fn default() -> Self {
62        TimerWheel::new(0)
63    }
64}
65
66impl<T> TimerWheel<T> {
67    /// An empty wheel whose clock reads `now`.
68    pub fn new(now: u64) -> TimerWheel<T> {
69        TimerWheel { now, nodes: Vec::new(), free: Vec::new(), heads: [NIL; LEVELS * SLOTS + 2], ready_tail: NIL, occupied: [0; LEVELS], len: 0 }
70    }
71
72    /// The wheel's current tick.
73    pub fn now(&self) -> u64 {
74        self.now
75    }
76
77    /// Scheduled timers.
78    pub fn len(&self) -> usize {
79        self.len
80    }
81
82    /// Whether no timer is scheduled.
83    pub fn is_empty(&self) -> bool {
84        self.len == 0
85    }
86
87    fn slot_for(&self, deadline: u64) -> u32 {
88        let d = deadline.max(self.now);
89        let diff = (d ^ self.now) | (SLOTS as u64 - 1);
90        let level = ((63 - diff.leading_zeros()) / 6) as usize;
91        if level >= LEVELS {
92            return OVERFLOW as u32;
93        }
94        let slot = (d >> (6 * level)) as usize & (SLOTS - 1);
95        (level * SLOTS + slot) as u32
96    }
97
98    fn link(&mut self, idx: u32) {
99        let slot = self.slot_for(self.nodes[idx as usize].deadline);
100        let head = self.heads[slot as usize];
101        {
102            let n = &mut self.nodes[idx as usize];
103            n.slot = slot;
104            n.prev = NIL;
105            n.next = head;
106        }
107        if head != NIL {
108            self.nodes[head as usize].prev = idx;
109        }
110        self.heads[slot as usize] = idx;
111        if (slot as usize) < OVERFLOW {
112            self.occupied[slot as usize / SLOTS] |= 1 << (slot as usize % SLOTS);
113        }
114    }
115
116    fn unlink(&mut self, idx: u32) {
117        let (prev, next, slot) = {
118            let n = &self.nodes[idx as usize];
119            (n.prev, n.next, n.slot)
120        };
121        if prev != NIL {
122            self.nodes[prev as usize].next = next;
123        } else {
124            self.heads[slot as usize] = next;
125            if next == NIL && (slot as usize) < OVERFLOW {
126                self.occupied[slot as usize / SLOTS] &= !(1 << (slot as usize % SLOTS));
127            }
128        }
129        if next != NIL {
130            self.nodes[next as usize].prev = prev;
131        } else if slot as usize == READY {
132            self.ready_tail = prev;
133        }
134    }
135
136    fn link_ready(&mut self, idx: u32) {
137        let tail = self.ready_tail;
138        {
139            let n = &mut self.nodes[idx as usize];
140            n.slot = READY as u32;
141            n.prev = tail;
142            n.next = NIL;
143        }
144        if tail != NIL {
145            self.nodes[tail as usize].next = idx;
146        } else {
147            self.heads[READY] = idx;
148        }
149        self.ready_tail = idx;
150    }
151
152    /// Schedules `payload` to expire at tick `deadline` (a deadline in the past expires at the
153    /// next [`TimerWheel::advance`]).
154    pub fn insert(&mut self, deadline: u64, payload: T) -> TimerId {
155        let idx = match self.free.pop() {
156            Some(i) => {
157                let n = &mut self.nodes[i as usize];
158                n.deadline = deadline;
159                n.payload = Some(payload);
160                i
161            }
162            None => {
163                self.nodes.push(Node { deadline, prev: NIL, next: NIL, slot: 0, gen: 0, payload: Some(payload) });
164                self.nodes.len() as u32 - 1
165            }
166        };
167        self.link(idx);
168        self.len += 1;
169        TimerId { idx: core::num::NonZeroU32::MIN.saturating_add(idx), gen: self.nodes[idx as usize].gen }
170    }
171
172    /// Whether a timer is still scheduled (not expired and popped, not cancelled).
173    pub fn is_pending(&self, id: TimerId) -> bool {
174        self.nodes.get(id.index()).is_some_and(|n| n.gen == id.gen && n.payload.is_some())
175    }
176
177    /// Cancels a timer and returns its payload, or `None` if it already expired or was
178    /// cancelled.
179    pub fn cancel(&mut self, id: TimerId) -> Option<T> {
180        let n = self.nodes.get(id.index())?;
181        if n.gen != id.gen || n.payload.is_none() {
182            return None;
183        }
184        let idx = id.index() as u32;
185        self.unlink(idx);
186        Some(self.release(idx))
187    }
188
189    fn release(&mut self, idx: u32) -> T {
190        let n = &mut self.nodes[idx as usize];
191        n.gen = n.gen.wrapping_add(1);
192        self.free.push(idx);
193        self.len -= 1;
194        n.payload.take().unwrap()
195    }
196
197    /// The earliest tick at which something may need doing: a deadline, or the start of a
198    /// slot whose timers must be re-filed. `None` when the wheel is empty.
199    pub fn next_deadline(&self) -> Option<u64> {
200        if self.heads[READY] != NIL {
201            return Some(self.now);
202        }
203        self.next_due().map(|(t, _)| t)
204    }
205
206    /// The next tick that needs work and the slot (or the overflow list) to take then.
207    fn next_due(&self) -> Option<(u64, usize)> {
208        let mut best: Option<(u64, usize)> = None;
209        for level in 0..LEVELS {
210            let occ = self.occupied[level];
211            if occ == 0 {
212                continue;
213            }
214            let shift = 6 * level as u32;
215            let cur = (self.now >> shift) as usize & (SLOTS - 1);
216            let ahead = occ.rotate_right(cur as u32).trailing_zeros() as usize;
217            let slot = (cur + ahead) % SLOTS;
218            let level_span = 1u64 << (shift + 6);
219            let mut t = (self.now & !(level_span - 1)) + ((slot as u64) << shift);
220            if slot < cur {
221                t += level_span;
222            }
223            // A past-due timer sits in the current level-0 slot, which reads as `now`.
224            best = Some((t.max(self.now), level * SLOTS + slot));
225            break;
226        }
227        if self.heads[OVERFLOW] != NIL {
228            let next_span = (self.now | ((1u64 << SPAN_BITS) - 1)).saturating_add(1);
229            if best.is_none_or(|(t, _)| next_span < t) {
230                best = Some((next_span, OVERFLOW));
231            }
232        }
233        best
234    }
235
236    /// Moves the clock to `now` and appends every expired payload to `out`, in deadline order.
237    pub fn advance(&mut self, now: u64, out: &mut Vec<T>) {
238        self.collect(now);
239        while let Some(p) = self.pop_ready() {
240            out.push(p);
241        }
242    }
243
244    /// Moves the clock to `now` and sets every expired timer aside, in deadline order, for
245    /// [`TimerWheel::pop_ready`]. Until popped they can still be cancelled, so a callback that
246    /// cancels a timer due in the same batch stops it from firing.
247    pub fn collect(&mut self, now: u64) {
248        if now < self.now {
249            return;
250        }
251        while let Some((t, slot)) = self.next_due() {
252            if t > now {
253                break;
254            }
255            self.now = t;
256            let mut idx = self.heads[slot];
257            self.heads[slot] = NIL;
258            if slot < OVERFLOW {
259                self.occupied[slot / SLOTS] &= !(1 << (slot % SLOTS));
260            }
261            while idx != NIL {
262                let next = self.nodes[idx as usize].next;
263                if self.nodes[idx as usize].deadline <= self.now {
264                    self.link_ready(idx);
265                } else {
266                    self.link(idx); // cascade to a finer level
267                }
268                idx = next;
269            }
270        }
271        self.now = now;
272    }
273
274    /// Takes the next expired timer set aside by [`TimerWheel::collect`]. Timers inserted after
275    /// the collect, even ones already due, wait for the next collect.
276    pub fn pop_ready(&mut self) -> Option<T> {
277        let idx = self.heads[READY];
278        if idx == NIL {
279            return None;
280        }
281        self.unlink(idx);
282        Some(self.release(idx))
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    fn drain(w: &mut TimerWheel<u64>, to: u64) -> Vec<u64> {
291        let mut out = Vec::new();
292        w.advance(to, &mut out);
293        out.sort();
294        out
295    }
296
297    #[test]
298    fn fires_at_the_deadline_never_early() {
299        let mut w = TimerWheel::new(0);
300        for d in [1u64, 5, 63, 64, 65, 100, 4095, 4096, 4097, 300_000, 20_000_000, 1u64 << 40] {
301            w.insert(d, d);
302        }
303        assert_eq!(w.next_deadline(), Some(1));
304        assert!(drain(&mut w, 0).is_empty());
305        assert_eq!(drain(&mut w, 1), vec![1]);
306        assert_eq!(drain(&mut w, 63), vec![5, 63]);
307        assert_eq!(drain(&mut w, 64), vec![64]);
308        assert_eq!(drain(&mut w, 4096), vec![65, 100, 4095, 4096]);
309        assert!(drain(&mut w, 299_999).iter().all(|&d| d <= 299_999));
310        assert_eq!(drain(&mut w, 300_000), vec![300_000]);
311        assert!(drain(&mut w, 19_999_999).is_empty());
312        assert_eq!(drain(&mut w, 20_000_000), vec![20_000_000]);
313        assert_eq!(w.len(), 1);
314        assert!(drain(&mut w, (1u64 << 40) - 1).is_empty());
315        assert_eq!(drain(&mut w, 1u64 << 40), vec![1u64 << 40]);
316        assert!(w.is_empty());
317        assert_eq!(w.next_deadline(), None);
318    }
319
320    #[test]
321    fn cancel_and_reuse() {
322        let mut w = TimerWheel::new(1000);
323        let a = w.insert(1010, 1);
324        let b = w.insert(1010, 2);
325        let c = w.insert(5000, 3);
326        assert_eq!(w.cancel(b), Some(2));
327        assert_eq!(w.cancel(b), None, "a stale id is ignored");
328        assert_eq!(w.cancel(c), Some(3));
329        let d = w.insert(1005, 4);
330        assert_ne!(d, b, "a reused node gets a new generation");
331        assert_eq!(drain(&mut w, 2000), vec![1, 4]);
332        assert_eq!(w.cancel(a), None);
333        assert!(w.is_empty());
334    }
335
336    #[test]
337    fn collected_timers_stay_cancellable_and_pop_in_order() {
338        let mut w = TimerWheel::new(0);
339        let a = w.insert(5, 'a');
340        let b = w.insert(15, 'b');
341        let c = w.insert(3, 'c');
342        w.collect(20);
343        assert_eq!(w.next_deadline(), Some(20));
344        assert_eq!(w.pop_ready(), Some('c'));
345        assert_eq!(w.cancel(b), Some('b'), "due but not popped: still cancellable");
346        let d = w.insert(1, 'd'); // due already, but waits for the next collect
347        assert_eq!(w.pop_ready(), Some('a'));
348        assert_eq!(w.pop_ready(), None);
349        assert_eq!(w.cancel(a), None);
350        w.collect(20);
351        assert_eq!(w.pop_ready(), Some('d'));
352        let _ = (c, d);
353        assert!(w.is_empty());
354    }
355
356    #[test]
357    fn past_deadlines_fire_on_the_next_advance() {
358        let mut w = TimerWheel::new(500);
359        w.insert(10, 7);
360        assert_eq!(w.next_deadline(), Some(500));
361        assert_eq!(drain(&mut w, 500), vec![7]);
362    }
363
364    #[test]
365    fn randomized_against_a_sorted_list() {
366        let mut seed = 0x9E37_79B9_7F4A_7C15u64;
367        let mut rnd = move || {
368            seed ^= seed << 13;
369            seed ^= seed >> 7;
370            seed ^= seed << 17;
371            seed
372        };
373        let mut w = TimerWheel::new(0);
374        let mut live: Vec<(u64, u64, TimerId)> = Vec::new();
375        let mut now = 0u64;
376        for step in 0..20_000u64 {
377            match rnd() % 4 {
378                0 | 1 => {
379                    let d = now + rnd() % [10, 1000, 100_000, 30_000_000][(rnd() % 4) as usize];
380                    let id = w.insert(d, step);
381                    live.push((d, step, id));
382                }
383                2 if !live.is_empty() => {
384                    let i = (rnd() % live.len() as u64) as usize;
385                    let (_, v, id) = live.swap_remove(i);
386                    assert_eq!(w.cancel(id), Some(v));
387                }
388                _ => {
389                    now += rnd() % 5000;
390                    let mut got = Vec::new();
391                    w.advance(now, &mut got);
392                    got.sort();
393                    let mut want: Vec<u64> = live.iter().filter(|x| x.0 <= now).map(|x| x.1).collect();
394                    want.sort();
395                    live.retain(|x| x.0 > now);
396                    assert_eq!(got, want, "at {now}");
397                    if let Some(next) = w.next_deadline() {
398                        assert!(next >= now);
399                        assert!(live.iter().all(|x| x.0 >= next), "nothing is due before next_deadline");
400                    }
401                }
402            }
403            assert_eq!(w.len(), live.len());
404        }
405    }
406
407    #[test]
408    fn insert_and_cancel_cost() {
409        let mut w = TimerWheel::new(0);
410        let n = 200_000u64;
411        let t = std::time::Instant::now();
412        let ids: Vec<TimerId> = (0..n).map(|i| w.insert(i % 60_000 + 1, i)).collect();
413        for id in &ids {
414            w.cancel(*id);
415        }
416        let per = t.elapsed().as_nanos() as f64 / n as f64;
417        eprintln!("timer insert+cancel: {per:.1} ns");
418    }
419}