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, timer_change) = self.heap.schedule(timer, value, at);
61        self.apply_timer_change(bindings_ctx, timer_change);
62        prev_value
63    }
64
65    /// Like [`schedule_instant`] but does the instant math from current time.
66    ///
67    /// # Panics
68    ///
69    /// Panics if the current `BC::Instant` cannot be represented by adding
70    /// duration `after`.
71    pub fn schedule_after(
72        &mut self,
73        bindings_ctx: &mut BC,
74        timer: K,
75        value: V,
76        after: Duration,
77    ) -> Option<(BC::Instant, V)> {
78        let time = bindings_ctx.now().checked_add(after).unwrap();
79        self.schedule_instant(bindings_ctx, timer, value, time)
80    }
81
82    /// Pops an expired timer from the heap, if any.
83    pub fn pop(&mut self, bindings_ctx: &mut BC) -> Option<(K, V)> {
84        let (popped, timer_change) = self.heap.pop_if(|t| t <= bindings_ctx.now());
85        self.apply_timer_change(bindings_ctx, timer_change);
86        popped
87    }
88
89    /// Returns the scheduled instant and associated value for `timer`, if it's
90    /// scheduled.
91    pub fn get(&self, timer: &K) -> Option<(BC::Instant, &V)> {
92        self.heap.map.get(timer).map(|MapEntry { time, value, synced_with_heap: _ }| (*time, value))
93    }
94
95    /// Cancels `timer`, returning the scheduled instant and associated value if
96    /// any.
97    pub fn cancel(&mut self, bindings_ctx: &mut BC, timer: &K) -> Option<(BC::Instant, V)> {
98        let (scheduled, timer_change) = self.heap.cancel(timer);
99        self.apply_timer_change(bindings_ctx, timer_change);
100        scheduled
101    }
102
103    /// Gets an iterator over the installed timers.
104    pub fn iter(&self) -> impl Iterator<Item = (&K, &V, &BC::Instant)> {
105        self.heap
106            .map
107            .iter()
108            .map(|(k, MapEntry { time, value, synced_with_heap: _ })| (k, value, time))
109    }
110
111    fn apply_timer_change(
112        &mut self,
113        bindings_ctx: &mut BC,
114        timer_change: Option<TimerChange<BC::Instant>>,
115    ) {
116        let _: Option<BC::Instant> = match timer_change {
117            Some(TimerChange::Scheduled(time)) => {
118                bindings_ctx.schedule_timer_instant(time, &mut self.next_wakeup)
119            }
120            Some(TimerChange::Canceled) => bindings_ctx.cancel_timer(&mut self.next_wakeup),
121            None => None,
122        };
123    }
124
125    /// Removes all timers.
126    pub fn clear(&mut self, bindings_ctx: &mut BC) {
127        let Self { next_wakeup, heap } = self;
128        heap.clear();
129        let _: Option<BC::Instant> = bindings_ctx.cancel_timer(next_wakeup);
130    }
131
132    /// Returns true if there are no timers installed.
133    pub fn is_empty(&self) -> bool {
134        self.heap.map.is_empty()
135    }
136}
137
138/// Describes a change to the next wakeup time in [`KeyedHeap`].
139enum TimerChange<T> {
140    Canceled,
141    Scheduled(T),
142}
143
144/// A timer heap that is keyed on `K`.
145///
146/// This type is used to support [`LocalTimerHeap`].
147///
148/// It is self-healing: all operations maintain the invariant that the top of
149/// `heap` (if any) is a valid entry synchronized with `map`.
150#[derive(Debug)]
151struct KeyedHeap<K, V, T> {
152    // Implementation note: The map is the source of truth for the desired
153    // firing time for a timer `K`. The heap has a copy of the scheduled time
154    // that is *always* compared to the map before firing the timer.
155    //
156    // That allows timers to be rescheduled to a later time during `heal`, which
157    // reduces memory utilization for the heap by avoiding the stale entry.
158    map: HashMap<K, MapEntry<T, V>>,
159    heap: BinaryHeap<HeapEntry<T, K>>,
160}
161
162/// The lower bound for heap garbage collection.
163const GC_MIN_HEAP_SIZE: usize = 20;
164
165impl<K: Hash + Eq + Clone, V, T: Instant> KeyedHeap<K, V, T> {
166    fn new() -> Self {
167        Self { map: HashMap::new(), heap: BinaryHeap::new() }
168    }
169
170    /// Garbage collects stale entries from `heap`.
171    ///
172    /// Bounds the heap size to 2 times the map size (or `GC_MIN_HEAP_SIZE` for
173    /// small heaps). Returns `true` if garbage collection was performed.
174    fn maybe_gc(&mut self) -> bool {
175        let Self { map, heap } = self;
176        let threshold = GC_MIN_HEAP_SIZE.max(map.len() * 2);
177        if heap.len() > threshold {
178            // NB: Rebuilding the heap from scratch is expected to be more
179            // efficient than removing the stale entries. At least half the
180            // heap is stale.
181            heap.clear();
182            heap.extend(map.iter_mut().map(
183                |(key, MapEntry { time, value: _, synced_with_heap })| {
184                    *synced_with_heap = true;
185                    HeapEntry { time: *time, key: key.clone() }
186                },
187            ));
188            true
189        } else {
190            false
191        }
192    }
193
194    /// Schedules `key` with associated `value` at time `at`.
195    ///
196    /// Returns the previously associated value and firing time for `key` +
197    /// a [`TimerChange`] if the next timer to fire changed with this operation.
198    fn schedule(&mut self, key: K, value: V, at: T) -> (Option<(T, V)>, Option<TimerChange<T>>) {
199        let Self { map, heap } = self;
200        let was_front =
201            heap.peek().is_some_and(|HeapEntry { time: _, key: top_key }| top_key == &key);
202        let is_new_front =
203            heap.peek().map(|HeapEntry { time: top_time, key: _ }| at <= *top_time).unwrap_or(true);
204        let (heap_entry, prev) = match map.entry(key) {
205            hash_map::Entry::Occupied(mut o) => {
206                let MapEntry { time: prev_time, value: _, synced_with_heap } = o.get();
207                // Only create a new `HeapEntry` if the already scheduled time
208                // is later than the new time. The new `MapEntry` is synced
209                // with the heap if we created a new `HeapEntry` or if the two
210                // times match and the old `MapEntry` was synced with the heap.
211                let (synced_with_heap, heap_entry) = match at.cmp(prev_time) {
212                    core::cmp::Ordering::Less => {
213                        (true, Some(HeapEntry { time: at, key: o.key().clone() }))
214                    }
215                    core::cmp::Ordering::Equal => (*synced_with_heap, None),
216                    core::cmp::Ordering::Greater => (false, None),
217                };
218                let MapEntry { time, value, synced_with_heap: _ } =
219                    o.insert(MapEntry { time: at, value, synced_with_heap });
220                (heap_entry, Some((time, value)))
221            }
222            hash_map::Entry::Vacant(v) => {
223                let heap_entry = Some(HeapEntry { time: at, key: v.key().clone() });
224                let _: &mut MapEntry<_, _> =
225                    v.insert(MapEntry { time: at, value, synced_with_heap: true });
226                (heap_entry, None)
227            }
228        };
229        let did_gc = match heap_entry {
230            Some(heap_entry) => {
231                heap.push(heap_entry);
232                self.maybe_gc()
233            }
234            None => false,
235        };
236        if was_front && !is_new_front {
237            // NB: If GC ran, the heap is already healed.
238            let new_top = if did_gc {
239                self.heap.peek().map(|HeapEntry { time, key: _ }| *time)
240            } else {
241                self.heal()
242            }
243            .expect("heap cannot be empty after a `schedule` operation");
244            return (prev, Some(TimerChange::Scheduled(new_top)));
245        }
246        (prev, is_new_front.then_some(TimerChange::Scheduled(at)))
247    }
248
249    /// Cancels the timer with `key`.
250    ///
251    /// Returns the scheduled instant and value for `key` if it was scheduled +
252    /// a [`TimerChange`] if the next timer to fire changed with this operation.
253    fn cancel(&mut self, key: &K) -> (Option<(T, V)>, Option<TimerChange<T>>) {
254        let Self { heap, map } = self;
255        let Some(MapEntry { time, value, synced_with_heap: _ }) = map.remove(key) else {
256            return (None, None);
257        };
258        // The front of the heap will be changed if we're cancelling the top.
259        let was_front = heap.peek().is_some_and(|HeapEntry { time: _, key: top }| key == top);
260        let did_gc = self.maybe_gc();
261        let timer_change = was_front.then(|| {
262            // NB: If GC ran, the heap is already healed.
263            let new_top = if did_gc {
264                self.heap.peek().map(|HeapEntry { time, key: _ }| *time)
265            } else {
266                self.heal()
267            };
268            match new_top {
269                Some(new_top) => TimerChange::Scheduled(new_top),
270                None => TimerChange::Canceled,
271            }
272        });
273        (Some((time, value)), timer_change)
274    }
275
276    /// Pops the first valid entry if `f` returns `true`.
277    ///
278    /// Returns the popped value if one is found *and* `f` returns true + a
279    /// [`TimerChange`] if the next timer to fire changed with this operation.
280    ///
281    /// NB: This API is a bit wonky, but unfortunately we can't seem to be able
282    /// to express a type that would be equivalent to `BinaryHeap`'s `PeekMut`.
283    /// This is the next best thing.
284    fn pop_if<F: FnOnce(T) -> bool>(&mut self, f: F) -> (Option<(K, V)>, Option<TimerChange<T>>) {
285        let Self { heap, map } = self;
286        let Some(peek_mut) = heap.peek_mut() else {
287            return (None, None);
288        };
289        if !f(peek_mut.time) {
290            return (None, None);
291        }
292        let HeapEntry { time: heap_time, key } = binary_heap::PeekMut::pop(peek_mut);
293        let MapEntry { time: scheduled_for, value, synced_with_heap } =
294            map.remove(&key).expect("top of heap must be present in map");
295        // NB: `KeyedHeap` is self-healing, so ensure the top of the map is
296        // always a valid entry.
297        debug_assert_eq!(heap_time, scheduled_for);
298        debug_assert!(synced_with_heap);
299        let timer_change = match self.heal() {
300            Some(new_top) => TimerChange::Scheduled(new_top),
301            None => TimerChange::Canceled,
302        };
303        (Some((key, value)), Some(timer_change))
304    }
305
306    /// Heals the heap of stale entries, returning the firing time of the top
307    /// valid entry, if any.
308    fn heal(&mut self) -> Option<T> {
309        let Self { heap, map } = self;
310        loop {
311            let peek_mut = heap.peek_mut()?;
312            let HeapEntry { time: heap_time, key } = &*peek_mut;
313            // Always check the map state for the given key, since it's the
314            // source of truth for desired firing time.
315
316            // NB: We assume here that the key is cheaply cloned and that
317            // cloning it is faster than possibly hashing it more than once.
318            match map.entry(key.clone()) {
319                hash_map::Entry::Vacant(_) => {
320                    // This `HeapEntry` is stale. This may happen if either
321                    //   1) the original timer was canceled, or
322                    //   2) the original timer was rescheduled to an earlier
323                    //      time, and has since fired.
324                    // Pop and continue looking.
325                    let _: HeapEntry<_, _> = binary_heap::PeekMut::pop(peek_mut);
326                }
327                hash_map::Entry::Occupied(mut map_entry) => {
328                    let MapEntry { time: scheduled_for, value: _, synced_with_heap } =
329                        map_entry.get_mut();
330
331                    match heap_time.cmp(scheduled_for) {
332                        core::cmp::Ordering::Equal => {
333                            // Map and heap agree on firing time, this is the top of
334                            // the heap.
335                            *synced_with_heap = true;
336                            return Some(*scheduled_for);
337                        }
338                        core::cmp::Ordering::Less => {
339                            // When rescheduling a timer, we only touch the heap
340                            // if rescheduling to an earlier time. In this case
341                            // the map is telling us this is scheduled for
342                            // later. We check `synced_with_heap` to determine
343                            // whether this entry is stale and should be dropped
344                            // or whether the entry must be resynchronized and
345                            // put back in the heap.
346                            let HeapEntry { time: _, key } = binary_heap::PeekMut::pop(peek_mut);
347                            if !*synced_with_heap {
348                                heap.push(HeapEntry { time: *scheduled_for, key });
349                                *synced_with_heap = true;
350                            }
351                        }
352                        core::cmp::Ordering::Greater => {
353                            // Heap time greater than scheduled time is
354                            // effectively unobservable because any earlier
355                            // entry would've been popped from the heap already
356                            // and thus the entry would be missing from the map.
357                            // Even a catastrophic cancel => reschedule cycle
358                            // can't make us observe this branch given the heap
359                            // ordering properties.
360                            unreachable!(
361                                "observed heap time: {:?} later than the scheduled time {:?}",
362                                heap_time, scheduled_for
363                            );
364                        }
365                    }
366                }
367            }
368        }
369    }
370
371    fn clear(&mut self) {
372        let Self { map, heap } = self;
373        map.clear();
374        heap.clear();
375    }
376}
377
378/// The entry kept in [`LocalTimerHeap`]'s internal hash map.
379#[derive(Debug)]
380struct MapEntry<T, V> {
381    time: T,
382    value: V,
383    /// Whether `heap` contains a [`HeapEntry`] corresponding to this
384    /// [`MapEntry`] with a matching `time`.
385    synced_with_heap: bool,
386}
387
388/// A reusable struct to place a value and a timestamp in a [`BinaryHeap`].
389///
390/// Its `Ord` implementation is tuned to make [`BinaryHeap`] a min heap.
391#[derive(Debug)]
392struct HeapEntry<T, K> {
393    time: T,
394    key: K,
395}
396
397// Boilerplate to implement a heap entry.
398impl<T: Instant, K> PartialEq for HeapEntry<T, K> {
399    fn eq(&self, other: &Self) -> bool {
400        self.time == other.time
401    }
402}
403
404impl<T: Instant, K> Eq for HeapEntry<T, K> {}
405
406impl<T: Instant, K> Ord for HeapEntry<T, K> {
407    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
408        // Note that we flip the argument order here to make the `BinaryHeap` a
409        // min heap.
410        Ord::cmp(&other.time, &self.time)
411    }
412}
413
414impl<T: Instant, K> PartialOrd for HeapEntry<T, K> {
415    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
416        Some(Ord::cmp(self, other))
417    }
418}
419
420#[cfg(any(test, feature = "testutils"))]
421mod testutil {
422    use core::fmt::Debug;
423    use core::ops::RangeBounds;
424
425    use super::*;
426
427    impl<K, V, BC> LocalTimerHeap<K, V, BC>
428    where
429        K: Hash + Eq + Clone + Debug,
430        V: Debug + Eq + Clone + PartialEq,
431        BC: TimerContext,
432    {
433        /// Asserts installed timers with an iterator of `(key, value, instant)`
434        /// tuples.
435        #[track_caller]
436        pub fn assert_timers(&self, timers: impl IntoIterator<Item = (K, V, BC::Instant)>) {
437            let wanted = timers.into_iter().map(|(k, v, i)| (k, (v, i))).collect::<HashMap<_, _>>();
438            let actual = self
439                .heap
440                .map
441                .iter()
442                .map(|(k, MapEntry { value, time, synced_with_heap: _ })| {
443                    (k.clone(), (value.clone(), *time))
444                })
445                .collect::<HashMap<_, _>>();
446            assert_eq!(actual, wanted);
447        }
448
449        /// Like [`LocalTimerHeap::assert_timers`], but asserts based on a
450        /// duration after `bindings_ctx.now()`.
451        #[track_caller]
452        pub fn assert_timers_after(
453            &self,
454            bindings_ctx: &mut BC,
455            timers: impl IntoIterator<Item = (K, V, Duration)>,
456        ) {
457            let now = bindings_ctx.now();
458            self.assert_timers(timers.into_iter().map(|(k, v, d)| (k, v, now.panicking_add(d))))
459        }
460
461        /// Assets that the next time to fire has `key` and `value`.
462        #[track_caller]
463        pub fn assert_top(&mut self, key: &K, value: &V) {
464            // NB: `KeyedHeap` is self-healing, so ensure the top of the map is
465            // always a valid entry.
466            let top = self.heap.heap.peek().map(|HeapEntry { time: heap_time, key }| {
467                let MapEntry { time: map_time, value, synced_with_heap } =
468                    self.heap.map.get(key).expect("top of heap must be present in map");
469                assert_eq!(heap_time, map_time);
470                assert!(synced_with_heap);
471                (key, value)
472            });
473            assert_eq!(top, Some((key, value)));
474        }
475
476        /// Asserts that the given timer is installed with an instant at the
477        /// provided range.
478        #[track_caller]
479        pub fn assert_range<
480            'a,
481            R: RangeBounds<BC::Instant> + Debug,
482            I: IntoIterator<Item = (&'a K, R)>,
483        >(
484            &'a self,
485            expect: I,
486        ) {
487            for (timer, range) in expect {
488                let time = self
489                    .get(timer)
490                    .map(|(t, _)| t)
491                    .unwrap_or_else(|| panic!("timer {timer:?} not present"));
492                assert!(range.contains(&time), "timer {timer:?} is at {time:?} not in {range:?}");
493            }
494        }
495
496        /// Asserts that the given timer is installed with an instant at the
497        /// provided range, returning its information.
498        #[track_caller]
499        pub fn assert_range_single<'a, R: RangeBounds<BC::Instant> + Debug>(
500            &'a self,
501            timer: &K,
502            range: R,
503        ) -> (BC::Instant, &'a V) {
504            let (time, value) =
505                self.get(timer).unwrap_or_else(|| panic!("timer {timer:?} not present"));
506            assert!(range.contains(&time), "timer {timer:?} is at {time:?} not in {range:?}");
507            (time, value)
508        }
509    }
510}
511
512#[cfg(test)]
513mod tests {
514    use alloc::vec::Vec;
515
516    use crate::InstantContext;
517    use crate::testutil::{FakeAtomicInstant, FakeInstant, FakeInstantCtx};
518
519    use super::*;
520
521    #[derive(Default)]
522    struct FakeTimerCtx {
523        instant: FakeInstantCtx,
524    }
525
526    impl InstantBindingsTypes for FakeTimerCtx {
527        type Instant = FakeInstant;
528        type AtomicInstant = FakeAtomicInstant;
529    }
530
531    impl InstantContext for FakeTimerCtx {
532        fn now(&self) -> Self::Instant {
533            self.instant.now()
534        }
535    }
536
537    impl TimerBindingsTypes for FakeTimerCtx {
538        type Timer = FakeTimer;
539        type DispatchId = ();
540        type UniqueTimerId = !;
541    }
542
543    impl TimerContext for FakeTimerCtx {
544        fn new_timer(&mut self, (): Self::DispatchId) -> Self::Timer {
545            FakeTimer::default()
546        }
547
548        fn schedule_timer_instant(
549            &mut self,
550            time: Self::Instant,
551            timer: &mut Self::Timer,
552        ) -> Option<Self::Instant> {
553            timer.scheduled.replace(time)
554        }
555
556        fn cancel_timer(&mut self, timer: &mut Self::Timer) -> Option<Self::Instant> {
557            timer.scheduled.take()
558        }
559
560        fn scheduled_instant(&self, timer: &mut Self::Timer) -> Option<Self::Instant> {
561            timer.scheduled.clone()
562        }
563
564        fn unique_timer_id(&self, _: &Self::Timer) -> Self::UniqueTimerId {
565            unimplemented!()
566        }
567    }
568
569    #[derive(Default, Debug)]
570    struct FakeTimer {
571        scheduled: Option<FakeInstant>,
572    }
573
574    #[derive(Eq, PartialEq, Debug, Ord, PartialOrd, Copy, Clone, Hash)]
575    struct TimerId(usize);
576
577    type LocalTimerHeap = super::LocalTimerHeap<TimerId, (), FakeTimerCtx>;
578
579    impl LocalTimerHeap {
580        fn heap_len(&self) -> usize {
581            self.heap.heap.len()
582        }
583
584        #[track_caller]
585        fn assert_heap_entries<I: IntoIterator<Item = (FakeInstant, TimerId)>>(&self, i: I) {
586            let mut want = i.into_iter().collect::<Vec<_>>();
587            want.sort();
588            let mut got = self
589                .heap
590                .heap
591                .iter()
592                .map(|HeapEntry { time, key }| (*time, *key))
593                .collect::<Vec<_>>();
594            got.sort();
595            assert_eq!(got, want);
596        }
597
598        #[track_caller]
599        fn assert_map_entries<I: IntoIterator<Item = (FakeInstant, TimerId)>>(&self, i: I) {
600            let want = i.into_iter().map(|(t, k)| (k, t)).collect::<HashMap<_, _>>();
601            let got = self
602                .heap
603                .map
604                .iter()
605                .map(|(k, MapEntry { time, value: (), synced_with_heap: _ })| (*k, *time))
606                .collect::<HashMap<_, _>>();
607            assert_eq!(got, want);
608        }
609    }
610
611    const TIMER1: TimerId = TimerId(1);
612    const TIMER2: TimerId = TimerId(2);
613    const TIMER3: TimerId = TimerId(3);
614
615    const T1: FakeInstant = FakeInstant { offset: Duration::from_secs(1) };
616    const T2: FakeInstant = FakeInstant { offset: Duration::from_secs(2) };
617    const T3: FakeInstant = FakeInstant { offset: Duration::from_secs(3) };
618    const T4: FakeInstant = FakeInstant { offset: Duration::from_secs(4) };
619    const T5: FakeInstant = FakeInstant { offset: Duration::from_secs(5) };
620
621    #[test]
622    fn schedule_instant() {
623        let mut ctx = FakeTimerCtx::default();
624        let mut heap = LocalTimerHeap::new(&mut ctx, ());
625        assert_eq!(heap.next_wakeup.scheduled, None);
626        heap.assert_heap_entries([]);
627
628        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), None);
629        heap.assert_heap_entries([(T2, TIMER2)]);
630        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
631
632        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
633        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2)]);
634        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
635
636        assert_eq!(heap.schedule_instant(&mut ctx, TIMER3, (), T3), None);
637        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
638        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
639    }
640
641    #[test]
642    fn schedule_after() {
643        let mut ctx = FakeTimerCtx::default();
644        let mut heap = LocalTimerHeap::new(&mut ctx, ());
645        assert_eq!(heap.next_wakeup.scheduled, None);
646        let long_duration = Duration::from_secs(5);
647        let short_duration = Duration::from_secs(1);
648
649        let long_instant = ctx.now().checked_add(long_duration).unwrap();
650        let short_instant = ctx.now().checked_add(short_duration).unwrap();
651
652        assert_eq!(heap.schedule_after(&mut ctx, TIMER1, (), long_duration), None);
653        assert_eq!(heap.next_wakeup.scheduled, Some(long_instant));
654        heap.assert_heap_entries([(long_instant, TIMER1)]);
655        heap.assert_map_entries([(long_instant, TIMER1)]);
656
657        assert_eq!(
658            heap.schedule_after(&mut ctx, TIMER1, (), short_duration),
659            Some((long_instant, ()))
660        );
661        assert_eq!(heap.next_wakeup.scheduled, Some(short_instant));
662        heap.assert_heap_entries([(short_instant, TIMER1), (long_instant, TIMER1)]);
663        heap.assert_map_entries([(short_instant, TIMER1)]);
664    }
665
666    #[test]
667    fn cancel() {
668        let mut ctx = FakeTimerCtx::default();
669        let mut heap = LocalTimerHeap::new(&mut ctx, ());
670        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
671        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), None);
672        assert_eq!(heap.schedule_instant(&mut ctx, TIMER3, (), T3), None);
673        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
674        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
675
676        assert_eq!(heap.cancel(&mut ctx, &TIMER1), Some((T1, ())));
677        heap.assert_heap_entries([(T2, TIMER2), (T3, TIMER3)]);
678        heap.assert_map_entries([(T2, TIMER2), (T3, TIMER3)]);
679        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
680
681        assert_eq!(heap.cancel(&mut ctx, &TIMER1), None);
682
683        assert_eq!(heap.cancel(&mut ctx, &TIMER3), Some((T3, ())));
684        // Timer3 is still in the heap, hasn't had a chance to cleanup.
685        heap.assert_heap_entries([(T2, TIMER2), (T3, TIMER3)]);
686        heap.assert_map_entries([(T2, TIMER2)]);
687        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
688
689        assert_eq!(heap.cancel(&mut ctx, &TIMER2), Some((T2, ())));
690        heap.assert_heap_entries([]);
691        heap.assert_map_entries([]);
692        assert_eq!(heap.next_wakeup.scheduled, None);
693    }
694
695    #[test]
696    fn pop() {
697        let mut ctx = FakeTimerCtx::default();
698        let mut heap = LocalTimerHeap::new(&mut ctx, ());
699        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
700        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), None);
701        assert_eq!(heap.schedule_instant(&mut ctx, TIMER3, (), T3), None);
702        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
703        heap.assert_map_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
704        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
705
706        assert_eq!(heap.pop(&mut ctx), None);
707        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
708        heap.assert_map_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
709        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
710
711        ctx.instant.time = T1;
712        assert_eq!(heap.pop(&mut ctx), Some((TIMER1, ())));
713        heap.assert_heap_entries([(T2, TIMER2), (T3, TIMER3)]);
714        heap.assert_map_entries([(T2, TIMER2), (T3, TIMER3)]);
715        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
716        assert_eq!(heap.pop(&mut ctx), None);
717        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
718
719        ctx.instant.time = T3;
720        assert_eq!(heap.pop(&mut ctx), Some((TIMER2, ())));
721        heap.assert_heap_entries([(T3, TIMER3)]);
722        heap.assert_map_entries([(T3, TIMER3)]);
723
724        assert_eq!(heap.next_wakeup.scheduled, Some(T3));
725        assert_eq!(heap.pop(&mut ctx), Some((TIMER3, ())));
726        heap.assert_heap_entries([]);
727        heap.assert_map_entries([]);
728        assert_eq!(heap.next_wakeup.scheduled, None);
729
730        assert_eq!(heap.pop(&mut ctx), None);
731    }
732
733    #[test]
734    fn reschedule() {
735        let mut ctx = FakeTimerCtx::default();
736        let mut heap = LocalTimerHeap::new(&mut ctx, ());
737        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
738        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), None);
739        assert_eq!(heap.schedule_instant(&mut ctx, TIMER3, (), T3), None);
740        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
741        heap.assert_map_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
742        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
743
744        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T4), Some((T2, ())));
745        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER2), (T3, TIMER3)]);
746        heap.assert_map_entries([(T1, TIMER1), (T4, TIMER2), (T3, TIMER3)]);
747
748        ctx.instant.time = T4;
749        // Popping TIMER1 makes the heap entry update.
750        assert_eq!(heap.pop(&mut ctx), Some((TIMER1, ())));
751        heap.assert_heap_entries([(T4, TIMER2), (T3, TIMER3)]);
752        heap.assert_map_entries([(T4, TIMER2), (T3, TIMER3)]);
753        assert_eq!(heap.next_wakeup.scheduled, Some(T3));
754
755        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T2), Some((T4, ())));
756        heap.assert_heap_entries([(T2, TIMER2), (T4, TIMER2), (T3, TIMER3)]);
757        heap.assert_map_entries([(T2, TIMER2), (T3, TIMER3)]);
758        assert_eq!(heap.next_wakeup.scheduled, Some(T2));
759
760        assert_eq!(heap.pop(&mut ctx), Some((TIMER2, ())));
761        // Still has stale TIMER2 entry.
762        heap.assert_heap_entries([(T4, TIMER2), (T3, TIMER3)]);
763        heap.assert_map_entries([(T3, TIMER3)]);
764        assert_eq!(heap.next_wakeup.scheduled, Some(T3));
765
766        assert_eq!(heap.pop(&mut ctx), Some((TIMER3, ())));
767        heap.assert_heap_entries([]);
768        heap.assert_map_entries([]);
769        assert_eq!(heap.next_wakeup.scheduled, None);
770        assert_eq!(heap.pop(&mut ctx), None);
771    }
772
773    #[test]
774    fn reschedule_later_heals_stale_entries() {
775        let mut ctx = FakeTimerCtx::default();
776        let mut heap = LocalTimerHeap::new(&mut ctx, ());
777
778        // Schedule TIMER1 at T3, then reschedule earlier to T2 and T1.
779        // All three entries are in `heap`.
780        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T3), None);
781        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T2), Some((T3, ())));
782        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), Some((T2, ())));
783        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER1), (T3, TIMER1)]);
784        heap.assert_map_entries([(T1, TIMER1)]);
785
786        // Rescheduling TIMER1 to a later instant T4 heals all three earlier
787        // entries from `heap` and replaces them with a single entry at T4.
788        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T4), Some((T1, ())));
789        heap.assert_heap_entries([(T4, TIMER1)]);
790        heap.assert_map_entries([(T4, TIMER1)]);
791        assert_eq!(heap.next_wakeup.scheduled, Some(T4));
792    }
793
794    #[test]
795    fn reschedule_later_heals_stale_entries_with_multiple_timers() {
796        let mut ctx = FakeTimerCtx::default();
797        let mut heap = LocalTimerHeap::new(&mut ctx, ());
798
799        // Schedule TIMER1 at T4, and TIMER2 at T3. Then reschedule TIMER1
800        // earlier to T2 and T1.
801        // All four entries are in `heap`.
802        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T4), None);
803        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T3), None);
804        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T2), Some((T4, ())));
805        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), Some((T2, ())));
806        heap.assert_heap_entries([(T1, TIMER1), (T2, TIMER1), (T3, TIMER2), (T4, TIMER1)]);
807        heap.assert_map_entries([(T1, TIMER1), (T3, TIMER2)]);
808
809        // Rescheduling TIMER1 from T1 to T5 heals the entries up until the
810        // next valid entry (T3, TIMER2). Stale entries after it are not healed.
811        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T5), Some((T1, ())));
812        heap.assert_heap_entries([(T3, TIMER2), (T4, TIMER1), (T5, TIMER1)]);
813        heap.assert_map_entries([(T3, TIMER2), (T5, TIMER1)]);
814
815        // When TIMER2 fires at T3, `pop_if` heals the remaining stale entry
816        // (T4, TIMER1).
817        assert_eq!(heap.next_wakeup.scheduled, Some(T3));
818        ctx.instant.time = T3;
819        assert_eq!(heap.pop(&mut ctx), Some((TIMER2, ())));
820        heap.assert_heap_entries([(T5, TIMER1)]);
821        heap.assert_map_entries([(T5, TIMER1)]);
822    }
823
824    // Regression test for a bug where the timer heap would not reschedule the
825    // next wakeup when it has two timers for the exact same instant at the top.
826    #[test]
827    fn multiple_timers_same_instant() {
828        let mut ctx = FakeTimerCtx::default();
829        let mut heap = LocalTimerHeap::new(&mut ctx, ());
830        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
831        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T1), None);
832        assert_eq!(heap.next_wakeup.scheduled.take(), Some(T1));
833
834        ctx.instant.time = T1;
835
836        // Ordering is not guaranteed, just assert that we're getting timers.
837        assert!(heap.pop(&mut ctx).is_some());
838        assert_eq!(heap.next_wakeup.scheduled, Some(T1));
839        assert!(heap.pop(&mut ctx).is_some());
840        assert_eq!(heap.next_wakeup.scheduled, None);
841        assert_eq!(heap.pop(&mut ctx), None);
842    }
843
844    #[test]
845    fn clear() {
846        let mut ctx = FakeTimerCtx::default();
847        let mut heap = LocalTimerHeap::new(&mut ctx, ());
848        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), T1), None);
849        assert_eq!(heap.schedule_instant(&mut ctx, TIMER2, (), T1), None);
850        heap.clear(&mut ctx);
851        heap.assert_map_entries([]);
852        assert_eq!(heap.next_wakeup.scheduled, None);
853    }
854
855    #[test]
856    fn reschedule_single_timer_will_trigger_gc() {
857        let mut ctx = FakeTimerCtx::default();
858        let mut heap = LocalTimerHeap::new(&mut ctx, ());
859
860        let mut next_time = FakeInstant { offset: Duration::from_secs(1000) };
861        assert_eq!(heap.schedule_instant(&mut ctx, TIMER1, (), next_time), None);
862
863        // Reschedule the timer to earlier, which will generate stale entries.
864        // The GC won't run until the `GC_MIN_HEAP_SIZE` threshold has been
865        // exceeded.
866        for i in 1..GC_MIN_HEAP_SIZE + 1 {
867            let prev_time = next_time;
868            next_time = next_time.checked_sub(Duration::from_secs(1)).unwrap();
869            assert_eq!(
870                heap.schedule_instant(&mut ctx, TIMER1, (), next_time),
871                Some((prev_time, ()))
872            );
873            if i == GC_MIN_HEAP_SIZE {
874                // Final Iteration: GC should have been triggered.
875                assert_eq!(heap.heap_len(), 1);
876            } else {
877                // Intermediate Iteration: GC should not have been triggered.
878                assert_eq!(heap.heap_len(), 1 + i);
879            }
880        }
881    }
882
883    #[test]
884    fn reschedule_multiple_timers_will_trigger_gc() {
885        let mut ctx = FakeTimerCtx::default();
886        let mut heap = LocalTimerHeap::new(&mut ctx, ());
887
888        // Setup `GC_MIN_HEAP_SIZE` independent timers.
889        let num_timers = GC_MIN_HEAP_SIZE;
890        for i in 0..num_timers {
891            assert_eq!(heap.schedule_instant(&mut ctx, TimerId(i), (), T3), None);
892        }
893        assert_eq!(heap.heap_len(), num_timers);
894
895        // Reschedule every timer to earlier, which will generate stale entries.
896        // The GC won't run because the `2*num_timers` threshold hasn't been
897        // exceeded.
898        for i in 0..num_timers {
899            assert_eq!(heap.schedule_instant(&mut ctx, TimerId(i), (), T2), Some((T3, ())));
900        }
901        assert_eq!(heap.heap_len(), num_timers * 2);
902
903        // Now that we're at the threshold, scheduling a timer again will
904        // trigger the GC.
905        assert_eq!(heap.schedule_instant(&mut ctx, TimerId(0), (), T1), Some((T2, ())));
906        assert_eq!(heap.heap_len(), num_timers);
907    }
908
909    #[test]
910    fn cancel_will_trigger_gc() {
911        let mut ctx = FakeTimerCtx::default();
912        let mut heap = LocalTimerHeap::new(&mut ctx, ());
913
914        // Setup a single timer at T1, and then `GC_MIN_HEAP_SIZE` independent
915        // timers at T2.
916        let num_timers = GC_MIN_HEAP_SIZE;
917        assert_eq!(heap.schedule_instant(&mut ctx, TimerId(num_timers), (), T1), None);
918        for i in 0..num_timers {
919            assert_eq!(heap.schedule_instant(&mut ctx, TimerId(i), (), T2), None);
920        }
921        assert_eq!(heap.heap_len(), num_timers + 1);
922
923        // Cancel the timers at T2. Because of the existence of the timer at T1,
924        // the heap's front doesn't change, and the canceled timers will remain
925        // as stale entries.
926        //
927        // We should be able to cancel num_timers / 2 timers, before GC
928        // exceeds the `2*num_timers` threshold and is triggered.
929        let tipping_point = num_timers / 2;
930        for i in 0..tipping_point {
931            assert_eq!(heap.cancel(&mut ctx, &TimerId(i)), Some((T2, ())));
932            assert_eq!(heap.heap_len(), num_timers + 1);
933        }
934        assert_eq!(heap.cancel(&mut ctx, &TimerId(tipping_point)), Some((T2, ())));
935        assert_eq!(heap.heap_len(), num_timers - tipping_point);
936    }
937}
netstack3_base/time/local_timer_heap.rs

netstack3_base/time/
local_timer_heap.rs
