1const LEVELS: usize = 4;
10const SLOTS: usize = 64;
11const NIL: u32 = u32::MAX;
12const OVERFLOW: usize = LEVELS * SLOTS;
14const SPAN_BITS: u32 = 6 * LEVELS as u32;
15const READY: usize = OVERFLOW + 1;
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
20pub struct TimerId {
21 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
42pub 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 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 pub fn now(&self) -> u64 {
74 self.now
75 }
76
77 pub fn len(&self) -> usize {
79 self.len
80 }
81
82 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 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 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 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 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 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 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 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 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); }
268 idx = next;
269 }
270 }
271 self.now = now;
272 }
273
274 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'); 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}