netstack3_base/time/
local_timer_heap.rs

1// Copyright 2024 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! A local timer heap for use in netstack3 core.
6
7use alloc::collections::{BinaryHeap, binary_heap};
8use core::hash::Hash;
9use core::time::Duration;
10
11use netstack3_hashmap::{HashMap, hash_map};
12
13use crate::{CoreTimerContext, Instant, InstantBindingsTypes, TimerBindingsTypes, TimerContext};
14
15/// A local timer heap that keeps timers for core modules.
16///
17/// `LocalTimerHeap` manages its wakeups through a [`TimerContext`].
18///
19/// `K` is the key that timers are keyed on. `V` is optional sidecar data to be
20/// kept with each timer.
21///
22/// Note that to provide fast timer deletion, `LocalTimerHeap` requires `K` to
23/// be `Clone` and the implementation assumes that the clone is cheap.
24#[derive(Debug)]
25pub struct LocalTimerHeap<K, V, BT: TimerBindingsTypes + InstantBindingsTypes> {
26    next_wakeup: BT::Timer,
27    heap: KeyedHeap<K, V, BT::Instant>,
28}
29
30impl<K, V, BC> LocalTimerHeap<K, V, BC>
31where
32    K: Hash + Eq + Clone,
33    BC: TimerContext,
34{
35    /// Creates a new `LocalTimerHeap` with wakeup dispatch ID `dispatch_id`.
36    pub fn new(bindings_ctx: &mut BC, dispatch_id: BC::DispatchId) -> Self {
37        let next_wakeup = bindings_ctx.new_timer(dispatch_id);
38        Self { next_wakeup, heap: KeyedHeap::new() }
39    }
40
41    /// Like [`new`] but uses `CC` to covert the `dispatch_id` to match the type required by `BC`.
42    pub fn new_with_context<D, CC: CoreTimerContext<D, BC>>(
43        bindings_ctx: &mut BC,
44        dispatch_id: D,
45    ) -> Self {
46        Self::new(bindings_ctx, CC::convert_timer(dispatch_id))
47    }
48
49    /// Schedules `timer` with `value` at or after `at`.
50    ///
51    /// If `timer` was already scheduled, returns the previously scheduled
52    /// instant and the associated value.
53    pub fn schedule_instant(
54        &mut self,
55        bindings_ctx: &mut BC,
56        timer: K,
57        value: V,
58        at: BC::Instant,
59    ) -> Option<(BC::Instant, V)> {
60        let (prev_value, dirty) = self.heap.schedule(timer, value, at);
61        if dirty {
62            self.heal_and_reschedule(bindings_ctx);
63        }
64        prev_value
65    }
66
67    /// Like [`schedule_instant`] but does the instant math from current time.
68    ///
69    /// # Panics
70    ///
71    /// Panics if the current `BC::Instant` cannot be represented by adding
72    /// duration `after`.
73    pub fn schedule_after(
74        &mut self,
75        bindings_ctx: &mut BC,
76        timer: K,
77        value: V,
78        after: Duration,
79    ) -> Option<(BC::Instant, V)> {
80        let time = bindings_ctx.now().checked_add(after).unwrap();
81        self.schedule_instant(bindings_ctx, timer, value, time)
82    }
83
84    /// Pops an expired timer from the heap, if any.
85    pub fn pop(&mut self, bindings_ctx: &mut BC) -> Option<(K, V)> {
86        let Self { next_wakeup: _, heap } = self;
87        let (popped, dirty) = heap.pop_if(|t| t <= bindings_ctx.now());
88        if dirty {
89            self.heal_and_reschedule(bindings_ctx);
90        }
91        popped
92    }
93
94    /// Returns the scheduled instant and associated value for `timer`, if it's
95    /// scheduled.
96    pub fn get(&self, timer: &K) -> Option<(BC::Instant, &V)> {
97        self.heap.map.get(timer).map(|MapEntry { time, value }| (*time, value))
98    }
99
100    /// Cancels `timer`, returning the scheduled instant and associated value if
101    /// any.
102    pub fn cancel(&mut self, bindings_ctx: &mut BC, timer: &K) -> Option<(BC::Instant, V)> {
103        let (scheduled, dirty) = self.heap.cancel(timer);
104        if dirty {
105            self.heal_and_reschedule(bindings_ctx);
106        }
107        scheduled
108    }
109
110    /// Gets an iterator over the installed timers.
111    pub fn iter(&self) -> impl Iterator<Item = (&K, &V, &BC::Instant)> {
112        self.heap.map.iter().map(|(k, MapEntry { time, value })| (k, value, time))
113    }
114
115    fn heal_and_reschedule(&mut self, bindings_ctx: &mut BC) {
116        let Self { next_wakeup, heap } = self;
117        let mut new_top = None;
118        let _ = heap.pop_if(|t| {
119            // Extract the next time to fire, but don't pop it from the heap.
120            // This is equivalent to peeking, but it also "heals" the keyed heap
121            // by getting rid of stale entries.
122            new_top = Some(t);
123            false
124        });
125        let _: Option<BC::Instant> = match new_top {
126            Some(time) => bindings_ctx.schedule_timer_instant(time, next_wakeup),
127            None => bindings_ctx.cancel_timer(next_wakeup),
128        };
129    }
130
131    /// Removes all timers.
132    pub fn clear(&mut self, bindings_ctx: &mut BC) {
133        let Self { next_wakeup, heap } = self;
134        heap.clear();
135        let _: Option<BC::Instant> = bindings_ctx.cancel_timer(next_wakeup);
136    }
137
138    /// Returns true if there are no timers installed.
139    pub fn is_empty(&self) -> bool {
140        self.heap.map.is_empty()
141    }
142}
143
144/// A timer heap that is keyed on `K`.
145///
146/// This type is used to support [`LocalTimerHeap`].
147#[derive(Debug)]
148struct KeyedHeap<K, V, T> {
149    // Implementation note: The map is the source of truth for the desired
150    // firing time for a timer `K`. The heap has a copy of the scheduled time
151    // that is *always* compared to the map before firing the timer.
152    //
153    // That allows timers to be rescheduled to a later time during `pop`, which
154    // reduces memory utilization for the heap by avoiding the stale entry.
155    map: HashMap<K, MapEntry<T, V>>,
156    heap: BinaryHeap<HeapEntry<T, K>>,
157}
158
159impl<K: Hash + Eq + Clone, V, T: Instant> KeyedHeap<K, V, T> {
160    fn new() -> Self {
161        Self { map: HashMap::new(), heap: BinaryHeap::new() }
162    }
163
164    /// Schedules `key` with associated `value` at time `at`.
165    ///
166    /// Returns the previously associated value and firing time for `key` +
167    /// a boolean indicating whether the top of the heap (i.e. the next timer to
168    /// fire) changed with this operation.
169    fn schedule(&mut self, key: K, value: V, at: T) -> (Option<(T, V)>, bool) {
170        let Self { map, heap } = self;
171        // The top of the heap is changed if any of the following is true:
172        // - There were previously no entries in the heap.
173        // - The current top of the heap is the key being changed here.
174        // - The scheduled value is earlier than the current top of the heap.
175        let dirty = heap
176            .peek()
177            .map(|HeapEntry { time, key: top_key }| top_key == &key || at < *time)
178            .unwrap_or(true);
179        let (heap_entry, prev) = match map.entry(key) {
180            hash_map::Entry::Occupied(mut o) => {
181                let MapEntry { time, value } = o.insert(MapEntry { time: at, value });
182                // Only create a new entry if the already scheduled time is
183                // later than the new value.
184                let heap_entry = (at < time).then(|| HeapEntry { time: at, key: o.key().clone() });
185                (heap_entry, Some((time, value)))
186            }
187            hash_map::Entry::Vacant(v) => {
188                let heap_entry = Some(HeapEntry { time: at, key: v.key().clone() });
189                let _: &mut MapEntry<_, _> = v.insert(MapEntry { time: at, value });
190                (heap_entry, None)
191            }
192        };
193        if let Some(heap_entry) = heap_entry {
194            heap.push(heap_entry);
195        }
196        (prev, dirty)
197    }
198
199    /// Cancels the timer with `key`.
200    ///
201    /// Returns the scheduled instant and value for `key` if it was scheduled +
202    /// a boolean indicating whether the top of the heap (i.e. the next timer to
203    /// fire) changed with this operation.
204    fn cancel(&mut self, key: &K) -> (Option<(T, V)>, bool) {
205        let Self { heap, map } = self;
206        // The front of the heap will be changed if we're cancelling the top.
207        let was_front = heap.peek().is_some_and(|HeapEntry { time: _, key: top }| key == top);
208        let prev = map.remove(key).map(|MapEntry { time, value }| (time, value));
209        (prev, was_front)
210    }
211
212    /// Heals the heap of stale entries, popping the first valid entry if `f`
213    /// returns `true``.
214    ///
215    /// Returns the popped value if one is found *and* `f` returns true + a
216    /// boolean that is `true` iff the internal heap has changed.
217    ///
218    /// NB: This API is a bit wonky, but unfortunately we can't seem to be able
219    /// to express a type that would be equivalent to `BinaryHeap`'s `PeekMut`.
220    /// This is the next best thing.
221    fn pop_if<F: FnOnce(T) -> bool>(&mut self, f: F) -> (Option<(K, V)>, bool) {
222        let mut changed_heap = false;
223        let popped = loop {
224            let Self { heap, map } = self;
225            let Some(peek_mut) = heap.peek_mut() else {
226                break None;
227            };
228            let HeapEntry { time: heap_time, key } = &*peek_mut;
229            // Always check the map state for the given key, since it's the
230            // source of truth for desired firing time.
231
232            // NB: We assume here that the key is cheaply cloned and that
233            // cloning it is faster than possibly hashing it more than once.
234            match map.entry(key.clone()) {
235                hash_map::Entry::Vacant(_) => {
236                    // Timer has been canceled. Pop and continue looking.
237                    let _: HeapEntry<_, _> = binary_heap::PeekMut::pop(peek_mut);
238                    changed_heap = true;
239                }
240                hash_map::Entry::Occupied(map_entry) => {
241                    let MapEntry { time: scheduled_for, value: _ } = map_entry.get();
242
243                    match heap_time.cmp(scheduled_for) {
244                        core::cmp::Ordering::Equal => {
245                            // Map and heap agree on firing time, this is the top of
246                            // the heap.
247                            break f(*scheduled_for).then(|| {
248                                let HeapEntry { time: _, key } =
249                                    binary_heap::PeekMut::pop(peek_mut);
250                                changed_heap = true;
251                                let MapEntry { time: _, value } = map_entry.remove();
252                                (key, value)
253                            });
254                        }
255                        core::cmp::Ordering::Less => {
256                            // When rescheduling a timer, we only touch the heap
257                            // if rescheduling to an earlier time. In this case
258                            // the map is telling us this is scheduled for
259                            // later, so we must put it back in the heap.
260                            let HeapEntry { time: _, key } = binary_heap::PeekMut::pop(peek_mut);
261                            heap.push(HeapEntry { time: *scheduled_for, key });
262                            changed_heap = true;
263                        }
264                        core::cmp::Ordering::Greater => {
265                            // Heap time greater than scheduled time is
266                            // effectively unobservable because any earlier
267                            // entry would've been popped from the heap already
268                            // and thus the entry would be missing from the map.
269                            // Even a catastrophic cancel => reschedule cycle
270                            // can't make us observe this branch given the heap
271                            // ordering properties.
272                            unreachable!(
273                                "observed heap time: {:?} later than the scheduled time {:?}",
274                                heap_time, scheduled_for
275                            );
276                        }
277                    }
278                }
279            }
280        };
281        (popped, changed_heap)
282    }
283
284    fn clear(&mut self) {
285        let Self { map, heap } = self;
286        map.clear();
287        heap.clear();
288    }
289}
290
291/// The entry kept in [`LocalTimerHeap`]'s internal hash map.
292#[derive(Debug, Eq, PartialEq)]
293struct MapEntry<T, V> {
294    time: T,
295    value: V,
296}
297
298/// A reusable struct to place a value and a timestamp in a [`BinaryHeap`].
299///
300/// Its `Ord` implementation is tuned to make [`BinaryHeap`] a min heap.
301#[derive(Debug)]
302struct HeapEntry<T, K> {
303    time: T,
304    key: K,
305}
306
307// Boilerplate to implement a heap entry.
308impl<T: Instant, K> PartialEq for HeapEntry<T, K> {
309    fn eq(&self, other: &Self) -> bool {
310        self.time == other.time
311    }
312}
313
314impl<T: Instant, K> Eq for HeapEntry<T, K> {}
315
316impl<T: Instant, K> Ord for HeapEntry<T, K> {
317    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
318        // Note that we flip the argument order here to make the `BinaryHeap` a
319        // min heap.
320        Ord::cmp(&other.time, &self.time)
321    }
322}
323
324impl<T: Instant, K> PartialOrd for HeapEntry<T, K> {
325    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
326        Some(Ord::cmp(self, other))
327    }
328}
329
330#[cfg(any(test, feature = "testutils"))]
331mod testutil {
332    use core::fmt::Debug;
333    use core::ops::RangeBounds;
334
335    use super::*;
336
337    impl<K, V, BC> LocalTimerHeap<K, V, BC>
338    where
339        K: Hash + Eq + Clone + Debug,
340        V: Debug + Eq + PartialEq,
341        BC: TimerContext,
342    {
343        /// Asserts installed timers with an iterator of `(key, value, instant)`
344        /// tuples.
345        #[track_caller]
346        pub fn assert_timers(&self, timers: impl IntoIterator<Item = (K, V, BC::Instant)>) {
347            let map = timers
348                .into_iter()
349                .map(|(k, value, time)| (k, MapEntry { value, time }))
350                .collect::<HashMap<_, _>>();
351            assert_eq!(&self.heap.map, &map);
352        }
353
354        /// Like [`LocalTimerHeap::assert_timers`], but asserts based on a
355        /// duration after `bindings_ctx.now()`.
356        #[track_caller]
357        pub fn assert_timers_after(
358            &self,
359            bindings_ctx: &mut BC,
360            timers: impl IntoIterator<Item = (K, V, Duration)>,
361        ) {
362            let now = bindings_ctx.now();
363            self.assert_timers(timers.into_iter().map(|(k, v, d)| (k, v, now.panicking_add(d))))
364        }
365
366        /// Assets that the next time to fire has `key` and `value`.
367        #[track_caller]
368        pub fn assert_top(&mut self, key: &K, value: &V) {
369            // NB: We can't know that the top of the heap holds a valid entry,
370            // so we need to do the slow thing and look in the map for this
371            // assertion.
372            let top = self
373                .heap
374                .map
375                .iter()
376                .min_by_key(|(_key, MapEntry { time, .. })| time)
377                .map(|(key, MapEntry { time: _, value })| (key, value));
378            assert_eq!(top, Some((key, value)));
379        }
380
381        /// Asserts that the given timer is installed with an instant at the
382        /// provided range.
383        #[track_caller]
384        pub fn assert_range<
385            'a,
386            R: RangeBounds<BC::Instant> + Debug,
387            I: IntoIterator<Item = (&'a K, R)>,
388        >(
389            &'a self,
390            expect: I,
391        ) {
392            for (timer, range) in expect {
393                let time = self
394                    .get(timer)
395                    .map(|(t, _)| t)
396                    .unwrap_or_else(|| panic!("timer {timer:?} not present"));
397                assert!(range.contains(&time), "timer {timer:?} is at {time:?} not in {range:?}");
398            }
399        }
400
401        /// Asserts that the given timer is installed with an instant at the
402        /// provided range, returning its information.
403        #[track_caller]
404        pub fn assert_range_single<'a, R: RangeBounds<BC::Instant> + Debug>(
405            &'a self,
406            timer: &K,
407            range: R,
408        ) -> (BC::Instant, &'a V) {
409            let (time, value) =
410                self.get(timer).unwrap_or_else(|| panic!("timer {timer:?} not present"));
411            assert!(range.contains(&time), "timer {timer:?} is at {time:?} not in {range:?}");
412            (time, value)
413        }
414    }
415}
416
417#[cfg(test)]
418mod tests {
419    use alloc::vec::Vec;
420
421    use crate::InstantContext;
422    use crate::testutil::{FakeAtomicInstant, FakeInstant, FakeInstantCtx};
423
424    use super::*;
425
426    #[derive(Default)]
427    struct FakeTimerCtx {
428        instant: FakeInstantCtx,
429    }
430
431    impl InstantBindingsTypes for FakeTimerCtx {
432        type Instant = FakeInstant;
433        type AtomicInstant = FakeAtomicInstant;
434    }
435
436    impl InstantContext for FakeTimerCtx {
437        fn now(&self) -> Self::Instant {
438            self.instant.now()
439        }
440    }
441
442    impl TimerBindingsTypes for FakeTimerCtx {
443        type Timer = FakeTimer;
444        type DispatchId = ();
445        type UniqueTimerId = !;
446    }
447
448    impl TimerContext for FakeTimerCtx {
449        fn new_timer(&mut self, (): Self::DispatchId) -> Self::Timer {
450            FakeTimer::default()
451        }
452
453        fn schedule_timer_instant(
454            &mut self,
455            time: Self::Instant,
456            timer: &mut Self::Timer,
457        ) -> Option<Self::Instant> {
458            timer.scheduled.replace(time)
459        }
460
461        fn cancel_timer(&mut self, timer: &mut Self::Timer) -> Option<Self::Instant> {
462            timer.scheduled.take()
463        }
464
465        fn scheduled_instant(&self, timer: &mut Self::Timer) -> Option<Self::Instant> {
466            timer.scheduled.clone()
467        }
468
469        fn unique_timer_id(&self, _: &Self::Timer) -> Self::UniqueTimerId {
470            unimplemented!()
471        }
472    }
473
474    #[derive(Default, Debug)]
475    struct FakeTimer {
476        scheduled: Option<FakeInstant>,
477    }
478
479    #[derive(Eq, PartialEq, Debug, Ord, PartialOrd, Copy, Clone, Hash)]
480    struct TimerId(usize);
481
482    type LocalTimerHeap = super::LocalTimerHeap<TimerId, (), FakeTimerCtx>;
483
484    impl LocalTimerHeap {
485        #[track_caller]
486        fn assert_heap_entries<I: IntoIterator<Item = (FakeInstant, TimerId)>>(&self, i: I) {
487            let mut want = i.into_iter().collect::<Vec<_>>();
488            want.sort();
489            let mut got = self
490                .heap
491                .heap
492                .iter()
493                .map(|HeapEntry { time, key }| (*time, *key))
494                .collect::<Vec<_>>();
495            got.sort();
496            assert_eq!(got, want);
497        }
498
499        #[track_caller]
500        fn assert_map_entries<I: IntoIterator<Item = (FakeInstant, TimerId)>>(&self, i: I) {
501            let want = i.into_iter().map(|(t, k)| (k, t)).collect::<HashMap<_, _>>();
502            let got = self
503                .heap
504                .map
505                .iter()
506                .map(|(k, MapEntry { time, value: () })| (*k, *time))
507                .collect::<HashMap<_, _>>();
508            assert_eq!(got, want);
509        }
510    }
511
512    const TIMER1: TimerId = TimerId(1);
513    const TIMER2: TimerId = TimerId(2);
514    const TIMER3: TimerId = TimerId(3);
515
516    const T1: FakeInstant = FakeInstant { offset: Duration::from_secs(1) };
517    const T2: FakeInstant = FakeInstant { offset: Duration::from_secs(2) };
518    const T3: FakeInstant = FakeInstant { offset: Duration::from_secs(3) };
519    const T4: FakeInstant = FakeInstant { offset: Duration::from_secs(4) };
520
521    #[test]
522    fn schedule_instant() {
523        let mut ctx = FakeTimerCtx::default();
524        let mut heap = LocalTimerHeap::new(&mut ctx, ());
525        assert_eq!(heap.next_wakeup.scheduled, None);
526        heap.assert_heap_entries([]);
527
528        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), None);
529        heap.assert_heap_entries([(T2, TIMER2)]);
530        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
531
532        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
533        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2)]);
534        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
535
536        assert_eq!(heap.schedule_instant(&mut ctx, TIMER3, (), T3), None);
537        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
538        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
539    }
540
541    #[test]
542    fn schedule_after() {
543        let mut ctx = FakeTimerCtx::default();
544        let mut heap = LocalTimerHeap::new(&mut ctx, ());
545        assert_eq!(heap.next_wakeup.scheduled, None);
546        let long_duration = Duration::from_secs(5);
547        let short_duration = Duration::from_secs(1);
548
549        let long_instant = ctx.now().checked_add(long_duration).unwrap();
550        let short_instant = ctx.now().checked_add(short_duration).unwrap();
551
552        assert_eq!(heap.schedule_after(&mut ctx, TIMER1, (), long_duration), None);
553        assert_eq!(heap.next_wakeup.scheduled, Some(long_instant));
554        heap.assert_heap_entries([(long_instant, TIMER1)]);
555        heap.assert_map_entries([(long_instant, TIMER1)]);
556
557        assert_eq!(
558            heap.schedule_after(&mut ctx, TIMER1, (), short_duration),
559            Some((long_instant, ()))
560        );
561        assert_eq!(heap.next_wakeup.scheduled, Some(short_instant));
562        heap.assert_heap_entries([(short_instant, TIMER1), (long_instant, TIMER1)]);
563        heap.assert_map_entries([(short_instant, TIMER1)]);
564    }
565
566    #[test]
567    fn cancel() {
568        let mut ctx = FakeTimerCtx::default();
569        let mut heap = LocalTimerHeap::new(&mut ctx, ());
570        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
571        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), None);
572        assert_eq!(heap.schedule_instant(&mut ctx, TIMER3, (), T3), None);
573        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
574        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
575
576        assert_eq!(heap.cancel(&mut ctx, &TIMER1), Some((T1, ())));
577        heap.assert_heap_entries([(T2, TIMER2), (T3, TIMER3)]);
578        heap.assert_map_entries([(T2, TIMER2), (T3, TIMER3)]);
579        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
580
581        assert_eq!(heap.cancel(&mut ctx, &TIMER1), None);
582
583        assert_eq!(heap.cancel(&mut ctx, &TIMER3), Some((T3, ())));
584        // Timer3 is still in the heap, hasn't had a chance to cleanup.
585        heap.assert_heap_entries([(T2, TIMER2), (T3, TIMER3)]);
586        heap.assert_map_entries([(T2, TIMER2)]);
587        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
588
589        assert_eq!(heap.cancel(&mut ctx, &TIMER2), Some((T2, ())));
590        heap.assert_heap_entries([]);
591        heap.assert_map_entries([]);
592        assert_eq!(heap.next_wakeup.scheduled, None);
593    }
594
595    #[test]
596    fn pop() {
597        let mut ctx = FakeTimerCtx::default();
598        let mut heap = LocalTimerHeap::new(&mut ctx, ());
599        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
600        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), None);
601        assert_eq!(heap.schedule_instant(&mut ctx, TIMER3, (), T3), None);
602        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
603        heap.assert_map_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
604        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
605
606        assert_eq!(heap.pop(&mut ctx), None);
607        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
608        heap.assert_map_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
609        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
610
611        ctx.instant.time = T1;
612        assert_eq!(heap.pop(&mut ctx), Some((TIMER1, ())));
613        heap.assert_heap_entries([(T2, TIMER2), (T3, TIMER3)]);
614        heap.assert_map_entries([(T2, TIMER2), (T3, TIMER3)]);
615        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
616        assert_eq!(heap.pop(&mut ctx), None);
617        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
618
619        ctx.instant.time = T3;
620        assert_eq!(heap.pop(&mut ctx), Some((TIMER2, ())));
621        heap.assert_heap_entries([(T3, TIMER3)]);
622        heap.assert_map_entries([(T3, TIMER3)]);
623
624        assert_eq!(heap.next_wakeup.scheduled, Some(T3));
625        assert_eq!(heap.pop(&mut ctx), Some((TIMER3, ())));
626        heap.assert_heap_entries([]);
627        heap.assert_map_entries([]);
628        assert_eq!(heap.next_wakeup.scheduled, None);
629
630        assert_eq!(heap.pop(&mut ctx), None);
631    }
632
633    #[test]
634    fn reschedule() {
635        let mut ctx = FakeTimerCtx::default();
636        let mut heap = LocalTimerHeap::new(&mut ctx, ());
637        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
638        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), None);
639        assert_eq!(heap.schedule_instant(&mut ctx, TIMER3, (), T3), None);
640        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
641        heap.assert_map_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
642        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
643
644        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T4), Some((T2, ())));
645        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
646        heap.assert_map_entries([(T1, TIMER1), (T4, TIMER2), (T3, TIMER3)]);
647
648        ctx.instant.time = T4;
649        // Popping TIMER1 makes the heap entry update.
650        assert_eq!(heap.pop(&mut ctx), Some((TIMER1, ())));
651        heap.assert_heap_entries([(T4, TIMER2), (T3, TIMER3)]);
652        heap.assert_map_entries([(T4, TIMER2), (T3, TIMER3)]);
653        assert_eq!(heap.next_wakeup.scheduled, Some(T3));
654
655        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), Some((T4, ())));
656        heap.assert_heap_entries([(T2, TIMER2), (T4, TIMER2), (T3, TIMER3)]);
657        heap.assert_map_entries([(T2, TIMER2), (T3, TIMER3)]);
658        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
659
660        assert_eq!(heap.pop(&mut ctx), Some((TIMER2, ())));
661        // Still has stale TIMER2 entry.
662        heap.assert_heap_entries([(T4, TIMER2), (T3, TIMER3)]);
663        heap.assert_map_entries([(T3, TIMER3)]);
664        assert_eq!(heap.next_wakeup.scheduled, Some(T3));
665
666        assert_eq!(heap.pop(&mut ctx), Some((TIMER3, ())));
667        heap.assert_heap_entries([]);
668        heap.assert_map_entries([]);
669        assert_eq!(heap.next_wakeup.scheduled, None);
670        assert_eq!(heap.pop(&mut ctx), None);
671    }
672
673    // Regression test for a bug where the timer heap would not reschedule the
674    // next wakeup when it has two timers for the exact same instant at the top.
675    #[test]
676    fn multiple_timers_same_instant() {
677        let mut ctx = FakeTimerCtx::default();
678        let mut heap = LocalTimerHeap::new(&mut ctx, ());
679        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
680        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T1), None);
681        assert_eq!(heap.next_wakeup.scheduled.take(), Some(T1));
682
683        ctx.instant.time = T1;
684
685        // Ordering is not guaranteed, just assert that we're getting timers.
686        assert!(heap.pop(&mut ctx).is_some());
687        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
688        assert!(heap.pop(&mut ctx).is_some());
689        assert_eq!(heap.next_wakeup.scheduled, None);
690        assert_eq!(heap.pop(&mut ctx), None);
691    }
692
693    #[test]
694    fn clear() {
695        let mut ctx = FakeTimerCtx::default();
696        let mut heap = LocalTimerHeap::new(&mut ctx, ());
697        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
698        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T1), None);
699        heap.clear(&mut ctx);
700        heap.assert_map_entries([]);
701        assert_eq!(heap.next_wakeup.scheduled, None);
702    }
703}
netstack3_base/time/local_timer_heap.rs

netstack3_base/time/
local_timer_heap.rs
