Nothing to see here, move along meow
1use crate::header::NONE_SENTINEL;
2use crate::intrusive_list::{IntrusiveList, ListLinks};
3
4pub struct TimerTag;
5
6pub trait TimerEntryStorage: ListLinks<TimerTag> {
7 fn tw_slot(&self, id: u32) -> u16;
8 fn set_tw_slot(&mut self, id: u32, slot: u16);
9 fn tw_deadline(&self, id: u32) -> u64;
10 fn set_tw_deadline(&mut self, id: u32, deadline: u64);
11}
12
13const SLOTS_PER_LEVEL: usize = 64;
14const NUM_LEVELS: usize = 4;
15const TOTAL_SLOTS: usize = SLOTS_PER_LEVEL * NUM_LEVELS;
16
17const LEVEL_GRANULARITY: [u64; NUM_LEVELS] = [1, 64, 4096, 262144];
18
19pub struct TimerWheel {
20 slots: [IntrusiveList; TOTAL_SLOTS],
21 current_tick: u64,
22 entry_count: u32,
23}
24
25impl TimerWheel {
26 pub const fn new() -> Self {
27 Self {
28 slots: [IntrusiveList::empty(); TOTAL_SLOTS],
29 current_tick: 0,
30 entry_count: 0,
31 }
32 }
33
34 pub const fn current_tick(&self) -> u64 {
35 self.current_tick
36 }
37
38 pub fn seed_tick(&mut self, tick: u64) {
39 self.current_tick = tick;
40 }
41
42 fn compute_slot(current: u64, deadline: u64) -> usize {
43 let delta = deadline.saturating_sub(current);
44 let level = match delta {
45 0..64 => 0,
46 64..4096 => 1,
47 4096..262144 => 2,
48 _ => 3,
49 };
50 let slot_within = ((deadline / LEVEL_GRANULARITY[level]) % 64) as usize;
51 level * SLOTS_PER_LEVEL + slot_within
52 }
53
54 fn append_to_slot(&mut self, slot: usize, id: u32, store: &mut impl TimerEntryStorage) {
55 store.set_tw_slot(id, slot as u16);
56 self.slots[slot].append(id, store);
57 self.entry_count = self.entry_count.saturating_add(1);
58 }
59
60 pub fn insert(&mut self, id: u32, deadline: u64, store: &mut impl TimerEntryStorage) {
61 let effective = match deadline <= self.current_tick {
62 true => self.current_tick + 1,
63 false => deadline,
64 };
65 store.set_tw_deadline(id, effective);
66 let slot = Self::compute_slot(self.current_tick, effective);
67 self.append_to_slot(slot, id, store);
68 }
69
70 pub fn cancel(&mut self, id: u32, store: &mut impl TimerEntryStorage) {
71 let slot = store.tw_slot(id) as usize;
72 if self.slots[slot].unlink(id, store) {
73 self.entry_count = self.entry_count.saturating_sub(1);
74 }
75 }
76
77 pub fn advance(
78 &mut self,
79 new_tick: u64,
80 store: &mut impl TimerEntryStorage,
81 mut on_expire: impl FnMut(u32),
82 ) {
83 if self.entry_count == 0 {
84 self.current_tick = new_tick;
85 return;
86 }
87
88 let start = self.current_tick + 1;
89 (start..=new_tick).for_each(|tick| {
90 self.current_tick = tick;
91
92 (1..NUM_LEVELS).rev().for_each(|level| {
93 let gran = LEVEL_GRANULARITY[level];
94 if tick % gran == 0 {
95 let slot_idx = level * SLOTS_PER_LEVEL + ((tick / gran) % 64) as usize;
96 self.redistribute_slot(slot_idx, store, &mut on_expire);
97 }
98 });
99
100 let level0_slot = (tick % 64) as usize;
101 self.redistribute_slot(level0_slot, store, &mut on_expire);
102 });
103 }
104
105 fn redistribute_slot(
106 &mut self,
107 slot: usize,
108 store: &mut impl TimerEntryStorage,
109 on_expire: &mut impl FnMut(u32),
110 ) {
111 let mut cursor = self.slots[slot].take_all();
112 let current_tick = self.current_tick;
113 let entry_count = &mut self.entry_count;
114 let slots = &mut self.slots;
115
116 core::iter::from_fn(|| {
117 (cursor != NONE_SENTINEL).then(|| {
118 let id = cursor;
119 cursor = store.link_next(id);
120 store.set_link_next(id, NONE_SENTINEL);
121 store.set_link_prev(id, NONE_SENTINEL);
122 *entry_count = entry_count.saturating_sub(1);
123
124 let deadline = store.tw_deadline(id);
125 match deadline <= current_tick {
126 true => on_expire(id),
127 false => {
128 let new_slot = Self::compute_slot(current_tick, deadline);
129 store.set_tw_slot(id, new_slot as u16);
130 slots[new_slot].append(id, store);
131 *entry_count = entry_count.saturating_add(1);
132 }
133 }
134 })
135 })
136 .count();
137 }
138}