sel4_async_time/
timer_queue.rs

Help
1//
2// Copyright 2023, Colias Group, LLC
3//
4// SPDX-License-Identifier: BSD-2-Clause
5//
6
7use alloc::collections::btree_map::{BTreeMap, Entry};
8
9use crate::SubKey;
10
11// We opt for a simple B-tree-based implementation rather than implementing a timer wheel. This is
12// good enough for now. Note that this approach is also good enough for `async-std`. If we ever do
13// actually need the scalability of something like a timer wheel, `tokio`'s implementation would be
14// a good place to start.
15
16// TODO
17// Add feature like `tokio::time::Interval`
18
19// NOTE
20// Once #![feature(btree_cursors)] stabilizes, revert back to using it for a simpler, more
21// lightweight, and more efficient (on the small scale) implementation. See git history for such an
22// implementation.
23
24pub struct TimerQueue<T, U, V> {
25    pending: BTreeMap<T, BTreeMap<U, V>>,
26}
27
28#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
29pub struct Key<T, U> {
30    absolute_expiry: T,
31    sub_key: U,
32}
33
34impl<T, U> Key<T, U> {
35    pub fn absolute_expiry(&self) -> &T {
36        &self.absolute_expiry
37    }
38}
39
40impl<T: Ord + Clone, U: SubKey + Clone, V> TimerQueue<T, U, V> {
41    pub fn new() -> Self {
42        Self {
43            pending: BTreeMap::new(),
44        }
45    }
46
47    pub fn peek_next_absolute_expiry(&self) -> Option<&T> {
48        let (absolute_expiry, _value) = self.pending.first_key_value()?;
49        Some(absolute_expiry)
50    }
51
52    pub fn insert(&mut self, absolute_expiry: T, value: V) -> Key<T, U> {
53        let sub_key = match self.pending.entry(absolute_expiry.clone()) {
54            Entry::Vacant(entry) => {
55                let sub_key = <U as SubKey>::min();
56                entry.insert(BTreeMap::new()).insert(sub_key.clone(), value);
57                sub_key
58            }
59            Entry::Occupied(mut entry) => {
60                let sub_map = entry.get_mut();
61                let sub_key = sub_map
62                    .last_entry()
63                    .unwrap()
64                    .key()
65                    .succ()
66                    .expect("too many timers for one instant");
67                sub_map.insert(sub_key.clone(), value);
68                sub_key
69            }
70        };
71
72        Key {
73            absolute_expiry,
74            sub_key,
75        }
76    }
77
78    pub fn try_remove(&mut self, key: &Key<T, U>) -> Option<Expired<T, U, V>> {
79        match self
80            .pending
81            .entry(key.absolute_expiry().clone() /* HACK silly clone */)
82        {
83            Entry::Vacant(_) => None,
84            Entry::Occupied(mut entry) => {
85                let sub_map = entry.get_mut();
86                let (sub_key, value) = sub_map.remove_entry(&key.sub_key)?;
87                if sub_map.is_empty() {
88                    entry.remove();
89                }
90                Some(Expired {
91                    key: Key {
92                        absolute_expiry: key.absolute_expiry().clone(),
93                        sub_key,
94                    },
95                    value,
96                })
97            }
98        }
99    }
100
101    pub fn remove(&mut self, key: &Key<T, U>) -> Expired<T, U, V> {
102        self.try_remove(key).unwrap()
103    }
104
105    pub fn iter_expired(&mut self, now: T) -> IterExpired<'_, T, U, V> {
106        IterExpired { now, queue: self }
107    }
108}
109
110pub struct Expired<T, U, V> {
111    #[allow(dead_code)]
112    key: Key<T, U>,
113    value: V,
114}
115
116impl<T, U, V> Expired<T, U, V> {
117    #[allow(dead_code)]
118    pub fn key(&self) -> &Key<T, U> {
119        &self.key
120    }
121
122    #[allow(dead_code)]
123    pub fn absolute_expiry(&self) -> &T {
124        self.key().absolute_expiry()
125    }
126
127    pub fn value(&self) -> &V {
128        &self.value
129    }
130}
131
132pub struct IterExpired<'a, T, U, V> {
133    now: T,
134    queue: &'a mut TimerQueue<T, U, V>,
135}
136
137impl<T: Ord + Clone, U: Ord, V> Iterator for IterExpired<'_, T, U, V> {
138    type Item = Expired<T, U, V>;
139
140    fn next(&mut self) -> Option<Self::Item> {
141        let mut entry = self.queue.pending.first_entry()?;
142        if entry.key() > &self.now {
143            return None;
144        }
145        let sub_map = entry.get_mut();
146        let (sub_key, value) = sub_map.pop_first().unwrap();
147        let absolute_expiry = if sub_map.is_empty() {
148            entry.remove_entry().0
149        } else {
150            entry.key().clone()
151        };
152        Some(Expired {
153            key: Key {
154                absolute_expiry,
155                sub_key,
156            },
157            value,
158        })
159    }
160}
sel4_async_time/timer_queue.rs

sel4_async_time/
timer_queue.rs
