Skip to main content

ntex_util/time/
wheel.rs

1//! Hierarchical timer wheel backing ntex timers.
2//!
3//! The design follows the Linux kernel timer wheel (`kernel/time/timer.c`).
4//!
5//! # Clock
6//!
7//! The wheel counts time in *units* of 16 milliseconds (`1 << UNITS`). The
8//! wheel clock `elapsed` is the unit of the last processed expiry, and
9//! `elapsed_time` is the [`Instant`] that corresponds to it. Timer deadlines
10//! are converted to units relative to this pair.
11//!
12//! # Levels
13//!
14//! The wheel has `LVL_DEPTH` (8) levels of `LVL_SIZE` (64) buckets. Each level
15//! is 8 times coarser than the previous one, a timer is placed on the first
16//! level that can represent its delay:
17//!
18//! | Level | Granularity | Range            |
19//! |-------|-------------|------------------|
20//! | 0     | 16 ms       | 0 .. ~1 s        |
21//! | 1     | 128 ms      | ~1 s .. ~8 s     |
22//! | 2     | ~1 s        | ~8 s .. ~64 s    |
23//! | 3     | ~8 s        | ~64 s .. ~8.6 m  |
24//! | 4     | ~65 s       | ~8.6 m .. ~69 m  |
25//! | 5     | ~8.7 m      | ~69 m .. ~9.2 h  |
26//! | 6     | ~70 m       | ~9.2 h .. ~3 d   |
27//! | 7     | ~9.3 h      | ~3 d .. ~24.5 d  |
28//!
29//! Longer delays are clamped to the capacity of the wheel. The expiry is
30//! rounded up to the granularity of the level and the delay is measured from
31//! [`Instant::now()`], so a timer never fires early but may fire up to one
32//! granularity late. Timers are not cascaded to finer
33//! levels, which keeps insertion and removal `O(1)`.
34//!
35//! A bitmap per level tracks occupied buckets, the next expiry is found by
36//! scanning the bitmaps instead of the buckets.
37//!
38//! # Drivers
39//!
40//! Two tasks are spawned lazily on the current thread:
41//!
42//! * [`TimerDriver`] sleeps until the next occupied bucket expires and wakes
43//!   the timers stored in it. The clock advances to the scheduled time of the
44//!   bucket, not to the time the driver woke up, so a late wakeup does not
45//!   delay later timers. Buckets that are overdue are processed at once.
46//!   Wakers are invoked in place while the wheel is held; timer operations
47//!   are not reentrant. Callback timers are queued separately and invoked
48//!   after releasing the wheel.
49//! * [`LowresTimerDriver`] invalidates the cached [`system_time()`] value
50//!   every 150 milliseconds.
51//!
52//! Dropping the timer driver, i.e. when the runtime stops, stops the wheel and
53//! marks all timers as elapsed. Dropping the lowres driver invalidates the
54//! cached system time.
55//!
56//! A runtime does not always drop its pending tasks when it stops, the timer
57//! is also reset once the arbiter storage of the stopped system is cleared,
58//! so that the next runtime on the thread starts its own drivers. Drivers of
59//! an older generation exit without touching the state.
60use std::num::NonZeroUsize;
61use std::time::{Duration, Instant, SystemTime};
62use std::{cell::Cell, cmp, future::Future, mem, pin::Pin, rc::Rc, task, task::Poll};
63
64use futures_timer::Delay;
65use slab::Slab;
66
67use super::callbacks::Callbacks;
68use crate::task::LocalWaker;
69
70/// Resolution of the wheel clock, a unit is `1 << UNITS` milliseconds.
71const UNITS: u64 = 4;
72
73/// Each level is `LVL_CLK_DIV` times coarser than the previous one.
74const LVL_CLK_SHIFT: u64 = 3;
75const LVL_CLK_DIV: u64 = 1 << LVL_CLK_SHIFT;
76const LVL_CLK_MASK: u64 = LVL_CLK_DIV - 1;
77
78/// Number of buckets per level.
79const LVL_BITS: u64 = 6;
80const LVL_SIZE: u64 = 1 << LVL_BITS;
81const LVL_MASK: u64 = LVL_SIZE - 1;
82
83/// Number of levels.
84const LVL_DEPTH: u64 = 8;
85
86/// Total number of buckets.
87const WHEEL_SIZE: usize = (LVL_SIZE * LVL_DEPTH) as usize;
88
89/// Delays at or above the cutoff are clamped to `WHEEL_TIMEOUT_MAX`.
90const WHEEL_TIMEOUT_CUTOFF: u64 = lvl_start(LVL_DEPTH);
91const WHEEL_TIMEOUT_MAX: u64 = WHEEL_TIMEOUT_CUTOFF - lvl_gran(LVL_DEPTH - 1);
92
93/// Refresh interval of the cached time.
94const LOWRES_RESOLUTION: Duration = Duration::from_millis(150);
95
96/// Shift of the level clock relative to the wheel clock.
97const fn lvl_shift(lvl: u64) -> u64 {
98    lvl * LVL_CLK_SHIFT
99}
100
101/// Granularity of a level in units.
102const fn lvl_gran(lvl: u64) -> u64 {
103    1 << lvl_shift(lvl)
104}
105
106/// Smallest delay in units that is stored on level `lvl`, `lvl` must be at
107/// least 1.
108const fn lvl_start(lvl: u64) -> u64 {
109    (LVL_SIZE - 1) << ((lvl - 1) * LVL_CLK_SHIFT)
110}
111
112const fn to_units(millis: u64) -> u64 {
113    millis >> UNITS
114}
115
116const fn to_millis(units: u64) -> u64 {
117    units << UNITS
118}
119
120const fn as_millis(dur: Duration) -> u64 {
121    dur.as_secs() * 1_000 + (dur.subsec_millis() as u64)
122}
123
124/// Returns the start of the current runtime turn, see [`ntex_rt::now()`].
125///
126/// Tasks polled in the same turn observe the same instant on neon and compio,
127/// on tokio this is [`Instant::now()`].
128#[inline]
129pub fn now() -> Instant {
130    ntex_rt::now()
131}
132
133/// Returns a cached approximation of the current system time.
134///
135/// The cached value is refreshed at roughly 150 millisecond intervals.
136#[inline]
137pub fn system_time() -> SystemTime {
138    TIMER.with(Timer::system_time)
139}
140
141#[derive(Debug)]
142/// Handle to a timer registered with ntex's local timer wheel.
143///
144/// Dropping the handle cancels the timer. A handle may be reset and reused
145/// after it has elapsed.
146///
147/// Timer operations are not reentrant. Registered wakers must not synchronously
148/// access the timer wheel when cloned, woken or dropped.
149pub struct TimerHandle(NonZeroUsize);
150
151impl TimerHandle {
152    /// Registers a timer that elapses after `millis`.
153    pub fn new(millis: u64) -> Self {
154        TIMER.with(|t| t.add_timer(millis))
155    }
156
157    /// Restarts the timer with a new delay in milliseconds.
158    pub fn reset(&self, millis: u64) {
159        TIMER.with(|t| t.update_timer(self.0.get(), millis));
160    }
161
162    /// Completes the timer immediately and wakes its registered task.
163    pub fn elapse(&self) {
164        TIMER.with(|t| t.remove_timer(self.0.get()));
165    }
166
167    /// Returns `true` if this timer has elapsed.
168    pub fn is_elapsed(&self) -> bool {
169        TIMER.with(|t| t.with_wheel(|w| w.timers[self.0.get()].bucket.is_none()))
170    }
171
172    /// Polls until this timer has elapsed.
173    pub fn poll_elapsed(&self, cx: &mut task::Context<'_>) -> Poll<()> {
174        TIMER.with(|t| {
175            t.with_wheel(|w| {
176                let entry = &mut w.timers[self.0.get()];
177                if entry.bucket.is_none() {
178                    Poll::Ready(())
179                } else {
180                    if entry.task.as_ref().is_none_or(|w| !w.will_wake(cx.waker())) {
181                        entry.task = Some(cx.waker().clone());
182                    }
183                    Poll::Pending
184                }
185            })
186        })
187    }
188}
189
190impl Drop for TimerHandle {
191    fn drop(&mut self) {
192        // the wheel is already destroyed if the handle is dropped by
193        // another thread-local destructor
194        let _ = TIMER.try_with(|t| {
195            t.with_wheel(|w| {
196                w.unlink(self.0.get());
197                w.timers.remove(self.0.get());
198            });
199        });
200    }
201}
202
203bitflags::bitflags! {
204    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
205    struct Flags: u8 {
206        /// The timer driver task is spawned.
207        const DRIVER_STARTED = 0b0000_0001;
208        /// The cached time is populated and its refresh sleep is armed.
209        const LOWRES_TIMER   = 0b0000_1000;
210        /// The lowres driver task is spawned.
211        const LOWRES_DRIVER  = 0b0001_0000;
212        /// A timer was scheduled, `system_time()` caches the time.
213        const RUNNING        = 0b0010_0000;
214    }
215}
216
217thread_local! {
218    pub(super) static TIMER: Rc<Timer> = Rc::new(Timer::new());
219}
220
221/// Per-thread timer state, shared by the timer handles and the driver tasks.
222pub(super) struct Timer {
223    /// Wheel clock in units, the expiry that was processed last.
224    elapsed: Cell<u64>,
225    /// Instant that corresponds to `elapsed`, the scheduled time of the last
226    /// processed expiry. Set lazily when the wheel is idle.
227    elapsed_time: Cell<Option<Instant>>,
228    /// Expiry of the earliest occupied bucket, `u64::MAX` if the wheel is empty.
229    next_expiry: Cell<u64>,
230    flags: Cell<Flags>,
231    /// Incremented on reset, drivers of an older generation are stale.
232    pub(super) generation: Cell<u64>,
233    driver: LocalWaker,
234    lowres_stime: Cell<Option<SystemTime>>,
235    lowres_driver: LocalWaker,
236    /// Taken out of the cell for the duration of an operation.
237    wheel: Cell<Option<Box<Wheel>>>,
238}
239
240/// Timer storage.
241pub(super) struct Wheel {
242    /// Timer entries indexed by handle, slot 0 is reserved so handles are
243    /// non-zero.
244    pub(super) timers: Slab<TimerEntry>,
245    /// Handles of the timers stored in each bucket, `lvl * LVL_SIZE + offset`.
246    buckets: Box<[Slab<usize>]>,
247    /// Occupied buckets of each level, bit `n` is set if bucket `n` has timers.
248    occupied: [u64; LVL_DEPTH as usize],
249    /// Allocated only for callbacks, keeping ordinary entries waker-sized.
250    pub(super) callbacks: Option<Box<Callbacks>>,
251}
252
253struct WheelGuard<'a> {
254    slot: &'a Cell<Option<Box<Wheel>>>,
255    wheel: Option<Box<Wheel>>,
256}
257
258impl Drop for WheelGuard<'_> {
259    fn drop(&mut self) {
260        self.slot.set(self.wheel.take());
261    }
262}
263
264#[derive(Debug)]
265pub(super) struct TimerEntry {
266    /// Bucket index, `None` once the timer has elapsed.
267    pub(super) bucket: Option<u16>,
268    /// Key of the entry in the bucket.
269    pub(super) bucket_entry: usize,
270    pub(super) callback: bool,
271    pub(super) task: Option<task::Waker>,
272}
273
274struct ExpiredTimers<'a> {
275    timers: &'a mut Slab<TimerEntry>,
276    entries: slab::Drain<'a, usize>,
277}
278
279impl Drop for ExpiredTimers<'_> {
280    fn drop(&mut self) {
281        // A panicking wake must not leave entries pointing into a drained bucket.
282        for no in self.entries.by_ref() {
283            self.timers[no].bucket = None;
284        }
285    }
286}
287
288impl Timer {
289    fn new() -> Self {
290        let mut timers = Slab::default();
291        timers.insert(TimerEntry {
292            bucket: None,
293            bucket_entry: 0,
294            callback: false,
295            task: None,
296        });
297
298        Timer {
299            elapsed: Cell::new(0),
300            elapsed_time: Cell::new(None),
301            next_expiry: Cell::new(u64::MAX),
302            flags: Cell::new(Flags::empty()),
303            generation: Cell::new(0),
304            driver: LocalWaker::new(),
305            lowres_stime: Cell::new(None),
306            lowres_driver: LocalWaker::new(),
307            wheel: Cell::new(Some(Box::new(Wheel {
308                timers,
309                buckets: (0..WHEEL_SIZE).map(|_| Slab::new()).collect(),
310                occupied: [0; LVL_DEPTH as usize],
311                callbacks: None,
312            }))),
313        }
314    }
315
316    pub(super) fn with_wheel<F, R>(&self, f: F) -> R
317    where
318        F: FnOnce(&mut Wheel) -> R,
319    {
320        let mut guard = WheelGuard {
321            slot: &self.wheel,
322            wheel: Some(
323                self.wheel
324                    .take()
325                    .expect("timer wheel operations are not reentrant"),
326            ),
327        };
328        f(guard.wheel.as_deref_mut().unwrap())
329    }
330
331    fn insert_flags(&self, flags: Flags) -> Flags {
332        let mut f = self.flags.get();
333        f.insert(flags);
334        self.flags.set(f);
335        f
336    }
337
338    /// Returns the cached system time, the cache is populated only while the
339    /// wheel is running.
340    fn system_time(self: &Rc<Self>) -> SystemTime {
341        if let Some(cur) = self.lowres_stime.get() {
342            cur
343        } else {
344            let now = SystemTime::now();
345            if self.flags.get().contains(Flags::RUNNING) {
346                self.lowres_stime.set(Some(now));
347                self.refresh_lowres();
348            }
349            now
350        }
351    }
352
353    /// Arms the invalidation of the cached time.
354    fn refresh_lowres(self: &Rc<Self>) {
355        if self.flags.get().contains(Flags::LOWRES_DRIVER) {
356            self.lowres_driver.wake();
357        } else {
358            LowresTimerDriver::start(self);
359        }
360    }
361
362    /// Instant that corresponds to the wheel clock.
363    fn elapsed_time(&self) -> Instant {
364        if let Some(elapsed_time) = self.elapsed_time.get() {
365            elapsed_time
366        } else {
367            let elapsed_time = Instant::now();
368            self.elapsed_time.set(Some(elapsed_time));
369            elapsed_time
370        }
371    }
372
373    fn add_timer(self: &Rc<Self>, millis: u64) -> TimerHandle {
374        let no = if millis == 0 {
375            // elapsed immediately
376            self.with_wheel(|w| {
377                w.timers.insert(TimerEntry {
378                    bucket: None,
379                    bucket_entry: 0,
380                    callback: false,
381                    task: None,
382                })
383            })
384        } else {
385            let (idx, expiry) = self.calc_bucket(millis);
386            let no = self.with_wheel(|w| {
387                let no = w.timers.insert(TimerEntry {
388                    bucket: None,
389                    bucket_entry: 0,
390                    callback: false,
391                    task: None,
392                });
393                w.link(no, idx);
394                no
395            });
396            self.update_next_expiry(expiry);
397            no
398        };
399
400        // slot 0 is reserved in `Timer::new()`
401        TimerHandle(NonZeroUsize::new(no).unwrap())
402    }
403
404    /// Moves the timer to a new bucket, a zero delay elapses it and wakes its
405    /// task.
406    fn update_timer(self: &Rc<Self>, hnd: usize, millis: u64) {
407        if millis == 0 {
408            self.remove_timer(hnd);
409        } else {
410            let (idx, expiry) = self.calc_bucket(millis);
411            self.with_wheel(|w| w.relink(hnd, idx));
412            self.update_next_expiry(expiry);
413        }
414    }
415
416    /// Elapses the timer and wakes its task.
417    fn remove_timer(&self, hnd: usize) {
418        self.with_wheel(|w| {
419            if w.unlink(hnd)
420                && let Some(task) = w.timers[hnd].task.take()
421            {
422                task.wake();
423            }
424        });
425    }
426
427    /// Returns the bucket index and the bucket expiry for a delay starting now.
428    pub(super) fn calc_bucket(self: &Rc<Self>, millis: u64) -> (usize, u64) {
429        self.insert_flags(Flags::RUNNING);
430
431        // The delay is measured from the wheel clock. The cached time is not
432        // used, it goes stale while the thread is blocked and the timer would
433        // fire early by its age
434        let since = Instant::now().saturating_duration_since(self.elapsed_time());
435        let delta = to_units(as_millis(since) + millis);
436        self.calc_wheel_index(self.elapsed.get().wrapping_add(delta), delta)
437    }
438
439    /// Selects the level for a timer that expires at `expires`, `delta` units
440    /// from now.
441    fn calc_wheel_index(&self, expires: u64, delta: u64) -> (usize, u64) {
442        for lvl in 0..LVL_DEPTH {
443            if delta < lvl_start(lvl + 1) {
444                return Self::calc_index(expires, lvl);
445            }
446        }
447        // expire larger delays at the capacity limit of the wheel
448        Self::calc_index(
449            self.elapsed.get().wrapping_add(WHEEL_TIMEOUT_MAX),
450            LVL_DEPTH - 1,
451        )
452    }
453
454    /// Returns the bucket index and the bucket expiry on level `lvl`.
455    fn calc_index(expires: u64, lvl: u64) -> (usize, u64) {
456        // The timer must not fire early. Early expiry can happen because the
457        // timer is armed at the edge of a tick, or because the expiry is
458        // truncated to the level granularity, round up to prevent it.
459        let expires = (expires + lvl_gran(lvl)) >> lvl_shift(lvl);
460        (
461            (lvl * LVL_SIZE + (expires & LVL_MASK)) as usize,
462            expires << lvl_shift(lvl),
463        )
464    }
465
466    /// Wakes the driver if the new bucket expires before the current deadline.
467    pub(super) fn update_next_expiry(self: &Rc<Self>, expiry: u64) {
468        if expiry < self.next_expiry.get() {
469            self.next_expiry.set(expiry);
470            if self.flags.get().contains(Flags::DRIVER_STARTED) {
471                self.driver.wake();
472            } else {
473                TimerDriver::start(self);
474            }
475        }
476    }
477
478    /// Instant at which the bucket expiring at `expiry` is due.
479    fn expiry_time(&self, expiry: u64) -> Instant {
480        self.elapsed_time()
481            + Duration::from_millis(to_millis(expiry.saturating_sub(self.elapsed.get())))
482    }
483
484    /// Returns the expiry of the earliest occupied bucket.
485    fn next_pending_bucket(&self, wheel: &Wheel) -> Option<u64> {
486        let mut clk = self.elapsed.get();
487        let mut next = u64::MAX;
488
489        for lvl in 0..LVL_DEPTH {
490            let lvl_clk = clk & LVL_CLK_MASK;
491            let occupied = wheel.occupied[lvl as usize];
492
493            if occupied != 0 {
494                // distance to the next occupied bucket, wrapping around the level
495                let pos = u64::from(
496                    occupied
497                        .rotate_right((clk & LVL_MASK) as u32)
498                        .trailing_zeros(),
499                );
500                next = cmp::min(next, (clk + pos) << lvl_shift(lvl));
501
502                // The next level is reached once the clock of this level wraps
503                // to a multiple of `LVL_CLK_DIV`, an earlier bucket here cannot
504                // be preceded by one of the next level.
505                if pos <= ((LVL_CLK_DIV - lvl_clk) & LVL_CLK_MASK) {
506                    break;
507                }
508            }
509
510            // Clock of the next level. A partially elapsed tick of this level
511            // is rounded up, as the next level bucket it belongs to was already
512            // processed.
513            clk >>= LVL_CLK_SHIFT;
514            clk += u64::from(lvl_clk != 0);
515        }
516
517        if next < u64::MAX { Some(next) } else { None }
518    }
519
520    /// Removes `flags`, and `RUNNING` so the cached time is not populated and
521    /// no driver is spawned after the runtime stopped.
522    fn remove_flags(&self, flags: Flags) {
523        let mut f = self.flags.get();
524        f.remove(flags | Flags::RUNNING);
525        self.flags.set(f);
526    }
527
528    /// Invalidates the cached time, called when the lowres driver is dropped.
529    fn stop_lowres(&self) {
530        self.remove_flags(Flags::LOWRES_DRIVER | Flags::LOWRES_TIMER);
531        self.lowres_stime.set(None);
532    }
533
534    /// Stops the drivers of the current generation, called when the system
535    /// that runs them has stopped.
536    fn reset(&self) {
537        self.generation.set(self.generation.get().wrapping_add(1));
538        self.stop_wheel();
539        self.stop_lowres();
540        self.driver.take();
541        self.lowres_driver.take();
542    }
543
544    /// Marks all timers as elapsed and resets the wheel, called when the timer
545    /// driver is dropped. Tasks are not woken, the runtime is stopping.
546    fn stop_wheel(&self) {
547        self.remove_flags(Flags::DRIVER_STARTED);
548
549        self.with_wheel(|wheel| {
550            let Wheel {
551                timers,
552                buckets,
553                occupied,
554                callbacks,
555            } = wheel;
556            for bucket in buckets.iter_mut() {
557                for no in bucket.drain() {
558                    timers[no].bucket = None;
559                }
560            }
561            *occupied = [0; LVL_DEPTH as usize];
562            if let Some(callbacks) = callbacks {
563                callbacks.ready.clear();
564                for ticket in callbacks.entries.values() {
565                    if let Some(state) = ticket.state.upgrade() {
566                        state.pending.set(false);
567                    }
568                }
569            }
570
571            self.next_expiry.set(u64::MAX);
572            self.elapsed.set(0);
573            self.elapsed_time.set(None);
574        });
575    }
576}
577
578impl Wheel {
579    /// Level and occupied bit of a bucket.
580    fn bucket_bit(idx: usize) -> (usize, u64) {
581        (idx / LVL_SIZE as usize, 1 << (idx % LVL_SIZE as usize))
582    }
583
584    /// Stores the timer in bucket `idx`.
585    pub(super) fn link(&mut self, hnd: usize, idx: usize) {
586        let entry = &mut self.timers[hnd];
587        entry.bucket = Some(idx as u16);
588        entry.bucket_entry = self.buckets[idx].insert(hnd);
589
590        let (lvl, bit) = Self::bucket_bit(idx);
591        self.occupied[lvl] |= bit;
592    }
593
594    /// Removes the timer from its bucket and marks it as elapsed.
595    ///
596    /// Returns `false` if the timer has already elapsed.
597    pub(super) fn unlink(&mut self, hnd: usize) -> bool {
598        let entry = &mut self.timers[hnd];
599        if let Some(idx) = entry.bucket.take() {
600            let idx = idx as usize;
601            let bucket = &mut self.buckets[idx];
602            bucket.remove(entry.bucket_entry);
603            if bucket.is_empty() {
604                let (lvl, bit) = Self::bucket_bit(idx);
605                self.occupied[lvl] &= !bit;
606            }
607            true
608        } else {
609            false
610        }
611    }
612
613    /// Moves the timer to bucket `idx`.
614    pub(super) fn relink(&mut self, hnd: usize, idx: usize) {
615        if self.timers[hnd].bucket != Some(idx as u16) {
616            self.unlink(hnd);
617            self.link(hnd, idx);
618        }
619    }
620
621    /// Marks timers elapsed, waking tasks in place and queueing only callbacks.
622    fn execute_expired_timers(&mut self, mut clk: u64) {
623        for lvl in 0..LVL_DEPTH {
624            let idx = ((clk & LVL_MASK) + lvl * LVL_SIZE) as usize;
625            let bucket = &mut self.buckets[idx];
626            if !bucket.is_empty() {
627                let (lvl, bit) = Self::bucket_bit(idx);
628                self.occupied[lvl] &= !bit;
629                let mut expired = ExpiredTimers {
630                    timers: &mut self.timers,
631                    entries: bucket.drain(),
632                };
633                for no in expired.entries.by_ref() {
634                    let entry = &mut expired.timers[no];
635                    entry.bucket = None;
636                    if entry.callback {
637                        let callbacks = self.callbacks.as_mut().unwrap();
638                        callbacks.ready.push(callbacks.entries[&no].clone());
639                    } else if let Some(task) = entry.task.take() {
640                        task.wake();
641                    }
642                }
643            }
644
645            // The next level expires only when the clock is a multiple of its
646            // granularity
647            if (clk & LVL_CLK_MASK) != 0 {
648                break;
649            }
650            clk >>= LVL_CLK_SHIFT;
651        }
652    }
653}
654
655/// Resets the timer when dropped with the arbiter storage of a stopped system.
656#[derive(Default)]
657struct TimerReset;
658
659impl TimerReset {
660    fn register() {
661        ntex_rt::with_item::<TimerReset, _, _>(|_| ());
662    }
663}
664
665impl Drop for TimerReset {
666    fn drop(&mut self) {
667        // the timer is already destroyed if the thread is exiting
668        let _ = TIMER.try_with(|t| t.reset());
669    }
670}
671
672/// Task that sleeps until the next bucket expires and wakes its timers.
673struct TimerDriver {
674    timer: Rc<Timer>,
675    generation: u64,
676    sleep: Delay,
677    /// Deadline the sleep is armed for.
678    armed: Option<Instant>,
679}
680
681impl TimerDriver {
682    fn start(timer: &Rc<Timer>) {
683        timer.insert_flags(Flags::DRIVER_STARTED);
684        TimerReset::register();
685
686        let deadline = timer.expiry_time(timer.next_expiry.get());
687        crate::spawn_detached(TimerDriver {
688            timer: timer.clone(),
689            generation: timer.generation.get(),
690            sleep: Delay::new(deadline.saturating_duration_since(Instant::now())),
691            armed: Some(deadline),
692        });
693    }
694}
695
696impl Drop for TimerDriver {
697    fn drop(&mut self) {
698        if self.timer.generation.get() == self.generation {
699            self.timer.stop_wheel();
700        }
701    }
702}
703
704impl Future for TimerDriver {
705    type Output = ();
706
707    fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
708        let this = &mut *self;
709        let timer = &this.timer;
710        if timer.generation.get() != this.generation {
711            return Poll::Ready(());
712        }
713        timer.driver.register(cx.waker());
714
715        let mut now = Instant::now();
716
717        loop {
718            let expiry = timer.next_expiry.get();
719            if expiry == u64::MAX {
720                // the wheel is empty, a new timer wakes the driver
721                return Poll::Pending;
722            }
723
724            let deadline = timer.expiry_time(expiry);
725            if deadline > now {
726                if this.armed != Some(deadline) {
727                    this.armed = Some(deadline);
728                    this.sleep.reset(deadline.saturating_duration_since(now));
729                }
730                if Pin::new(&mut this.sleep).poll(cx).is_pending() {
731                    return Poll::Pending;
732                }
733                if deadline > now {
734                    // the sleep fired before the deadline, re-arm it
735                    this.armed = None;
736                    continue;
737                }
738            }
739
740            // Advance the clock to the scheduled time of the bucket, a late
741            // wakeup must not shift the timers that expire later
742            timer.elapsed.set(expiry);
743            timer.elapsed_time.set(Some(deadline));
744
745            let (next, callbacks) = timer.with_wheel(|w| {
746                w.execute_expired_timers(expiry);
747                let callbacks = w.callbacks.as_mut().and_then(|callbacks| {
748                    if callbacks.ready.is_empty() {
749                        None
750                    } else {
751                        Some(mem::take(&mut callbacks.ready))
752                    }
753                });
754                (timer.next_pending_bucket(w), callbacks)
755            });
756
757            if let Some(next) = next {
758                timer.next_expiry.set(next);
759            } else {
760                timer.next_expiry.set(u64::MAX);
761                timer.elapsed_time.set(None);
762            }
763
764            if let Some(mut callbacks) = callbacks {
765                for ticket in callbacks.drain(..) {
766                    if timer.generation.get() != this.generation {
767                        return Poll::Ready(());
768                    }
769                    if ticket.generation == this.generation
770                        && let Some(state) = ticket.state.upgrade()
771                        && state.arm.get() == ticket.arm
772                        && state.pending.replace(false)
773                    {
774                        (state.callback)();
775                    }
776                }
777                // A callback or its captures may have stopped this runtime.
778                if timer.generation.get() != this.generation {
779                    return Poll::Ready(());
780                }
781                timer.with_wheel(|wheel| wheel.callbacks.as_mut().unwrap().ready = callbacks);
782                now = Instant::now();
783            }
784        }
785    }
786}
787
788/// Task that invalidates the cached time every `LOWRES_RESOLUTION`.
789struct LowresTimerDriver {
790    timer: Rc<Timer>,
791    generation: u64,
792    sleep: Delay,
793}
794
795impl LowresTimerDriver {
796    fn start(timer: &Rc<Timer>) {
797        timer.insert_flags(Flags::LOWRES_DRIVER);
798        TimerReset::register();
799
800        crate::spawn_detached(LowresTimerDriver {
801            timer: timer.clone(),
802            generation: timer.generation.get(),
803            sleep: Delay::new(LOWRES_RESOLUTION),
804        });
805    }
806}
807
808impl Drop for LowresTimerDriver {
809    fn drop(&mut self) {
810        if self.timer.generation.get() == self.generation {
811            self.timer.stop_lowres();
812        }
813    }
814}
815
816impl Future for LowresTimerDriver {
817    type Output = ();
818
819    fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
820        let this = &mut *self;
821        let timer = &this.timer;
822        if timer.generation.get() != this.generation {
823            return Poll::Ready(());
824        }
825        timer.lowres_driver.register(cx.waker());
826
827        // the cache was populated, invalidate it after `LOWRES_RESOLUTION`
828        let mut flags = timer.flags.get();
829        if !flags.contains(Flags::LOWRES_TIMER) {
830            flags.insert(Flags::LOWRES_TIMER);
831            timer.flags.set(flags);
832            this.sleep.reset(LOWRES_RESOLUTION);
833        }
834
835        if Pin::new(&mut this.sleep).poll(cx).is_ready() {
836            timer.lowres_stime.set(None);
837            flags.remove(Flags::LOWRES_TIMER);
838            timer.flags.set(flags);
839        }
840        Poll::Pending
841    }
842}
843
844#[cfg(test)]
845mod tests {
846    use std::cell::RefCell;
847    use std::panic::{AssertUnwindSafe, catch_unwind};
848    use std::sync::Arc;
849    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
850    use std::task::{Context, Wake, Waker};
851
852    use super::*;
853    use crate::future::OnDropFn;
854    use crate::time::{CallbackTimer, Millis, interval, sleep};
855
856    #[derive(Default)]
857    struct InPlaceWaker {
858        wakes: AtomicUsize,
859        delay: Duration,
860        panic_on_wake: bool,
861    }
862
863    impl Wake for InPlaceWaker {
864        fn wake(self: Arc<Self>) {
865            TIMER.with(|timer| {
866                let wheel = timer.wheel.take();
867                let held = wheel.is_none();
868                timer.wheel.set(wheel);
869                assert!(held, "waker must run in place while the wheel is held");
870            });
871            let previous = self.wakes.fetch_add(1, Ordering::Relaxed);
872            assert!(!self.panic_on_wake, "notification");
873            if previous == 0 && !self.delay.is_zero() {
874                std::thread::sleep(self.delay);
875            }
876        }
877    }
878
879    fn timer_state() -> Rc<Timer> {
880        TIMER.with(|timer| {
881            timer.reset();
882            // These tests drive expiry manually, without an executor.
883            timer.insert_flags(Flags::DRIVER_STARTED);
884            timer.clone()
885        })
886    }
887
888    fn timer_driver(timer: &Rc<Timer>) -> TimerDriver {
889        TimerDriver {
890            timer: timer.clone(),
891            generation: timer.generation.get(),
892            sleep: Delay::new(Duration::ZERO),
893            armed: None,
894        }
895    }
896
897    fn make_due(timer: &Timer) {
898        timer.elapsed_time.set(Some(
899            Instant::now().checked_sub(Duration::from_secs(1)).unwrap(),
900        ));
901    }
902
903    #[ntex::test]
904    async fn test_callback_timer() {
905        let calls = Rc::new(Cell::new(0));
906        let count = calls.clone();
907        let callback = crate::time::CallbackTimer::new(Millis::ZERO, move || {
908            count.set(count.get() + 1);
909        });
910        assert!(!callback.is_elapsed());
911        assert_eq!(calls.get(), 0);
912        callback.reset(Millis(10_000));
913        callback.cancel();
914        assert!(callback.is_elapsed());
915        sleep(Millis(50)).await;
916        assert_eq!(calls.get(), 0);
917
918        for expected in 1..=2 {
919            callback.reset(Millis::ZERO);
920            assert!(!callback.is_elapsed());
921            sleep(Millis(50)).await;
922            assert!(callback.is_elapsed());
923            assert_eq!(calls.get(), expected);
924        }
925    }
926
927    #[test]
928    fn test_callback_stale_tickets() {
929        let timer = timer_state();
930        let calls = Rc::new(Cell::new(0));
931        let count = calls.clone();
932        let cancelled = CallbackTimer::new(Millis(1), move || count.set(count.get() + 1));
933        let count = calls.clone();
934        let reset = CallbackTimer::new(Millis(1), move || count.set(count.get() + 1));
935        let capture = Rc::new(Cell::new(0));
936        let weak = Rc::downgrade(&capture);
937        let dropped = CallbackTimer::new(Millis(1), move || capture.set(1));
938        let old_slot = dropped.hnd;
939        timer.with_wheel(|wheel| {
940            for handle in [&cancelled, &reset, &dropped] {
941                wheel.relink(handle.hnd.get(), 3);
942            }
943            wheel.execute_expired_timers(3);
944        });
945        assert!(cancelled.is_elapsed());
946        assert!(reset.is_elapsed());
947        assert!(dropped.is_elapsed());
948        assert_eq!(calls.get(), 0);
949        cancelled.cancel();
950        reset.reset(Millis(10_000));
951        drop(dropped);
952        assert!(
953            weak.upgrade().is_none(),
954            "queued tickets must not retain captures"
955        );
956        let count = calls.clone();
957        let replacement = CallbackTimer::new(Millis(10_000), move || count.set(count.get() + 100));
958        assert_eq!(replacement.hnd, old_slot);
959
960        make_due(&timer);
961        let mut driver = timer_driver(&timer);
962        assert!(
963            Pin::new(&mut driver)
964                .poll(&mut Context::from_waker(Waker::noop()))
965                .is_pending()
966        );
967        assert_eq!(calls.get(), 0);
968        assert!(!reset.is_elapsed());
969        assert!(!replacement.is_elapsed());
970        timer.elapsed_time.set(Some(
971            Instant::now().checked_sub(Duration::from_secs(20)).unwrap(),
972        ));
973        for _ in 0..2 {
974            assert!(
975                Pin::new(&mut driver)
976                    .poll(&mut Context::from_waker(Waker::noop()))
977                    .is_pending()
978            );
979            assert_eq!(calls.get(), 101);
980        }
981    }
982
983    #[test]
984    fn test_callback_self_rearm_and_cross_cancel() {
985        let timer = timer_state();
986        let calls = Rc::new(Cell::new(0));
987        let handles = Rc::new(RefCell::new(Vec::<CallbackTimer>::new()));
988        for index in 0..2 {
989            let count = calls.clone();
990            let weak = Rc::downgrade(&handles);
991            let callback = CallbackTimer::new(Millis(1), move || {
992                count.set(count.get() + 1);
993                let handles = weak.upgrade().unwrap();
994                let handles = handles.borrow();
995                assert!(handles[index].is_elapsed());
996                handles[1 - index].cancel();
997                handles[index].reset(Millis(10_000));
998            });
999            handles.borrow_mut().push(callback);
1000        }
1001        timer.with_wheel(|wheel| {
1002            for handle in handles.borrow().iter() {
1003                wheel.relink(handle.hnd.get(), 1);
1004            }
1005        });
1006        make_due(&timer);
1007        let mut driver = timer_driver(&timer);
1008        assert!(
1009            Pin::new(&mut driver)
1010                .poll(&mut Context::from_waker(Waker::noop()))
1011                .is_pending()
1012        );
1013        assert_eq!(calls.get(), 1);
1014        assert_eq!(
1015            handles.borrow().iter().filter(|h| !h.is_elapsed()).count(),
1016            1
1017        );
1018        timer.elapsed_time.set(Some(
1019            Instant::now().checked_sub(Duration::from_secs(20)).unwrap(),
1020        ));
1021        for _ in 0..2 {
1022            assert!(
1023                Pin::new(&mut driver)
1024                    .poll(&mut Context::from_waker(Waker::noop()))
1025                    .is_pending()
1026            );
1027            assert_eq!(calls.get(), 2);
1028        }
1029    }
1030
1031    #[test]
1032    fn test_callback_capture_drop_schedules_timers() {
1033        for active in [false, true] {
1034            let timer = timer_state();
1035            let released = Rc::new(Cell::new(false));
1036            let finished = Rc::new(Cell::new(false));
1037            let reset_calls = Rc::new(Cell::new(0));
1038            let new_calls = Rc::new(Cell::new(0));
1039            let count = reset_calls.clone();
1040            let reset = Rc::new(CallbackTimer::new(Millis(20_000), move || {
1041                count.set(count.get() + 1);
1042            }));
1043            let cancelled = Rc::new(CallbackTimer::new(Millis(20_000), || {
1044                panic!("cancelled callback");
1045            }));
1046            let replacement = Rc::new(RefCell::new(None::<CallbackTimer>));
1047            let capture = {
1048                let reset = reset.clone();
1049                let cancelled = cancelled.clone();
1050                let replacement = replacement.clone();
1051                let new_calls = new_calls.clone();
1052                let released = released.clone();
1053                let finished = finished.clone();
1054                OnDropFn::new(move || {
1055                    assert!(!active || finished.get());
1056                    let ordinary = TimerHandle::new(10_000);
1057                    ordinary.reset(0);
1058                    drop(ordinary);
1059                    cancelled.cancel();
1060                    reset.reset(Millis(10_000));
1061                    *replacement.borrow_mut() =
1062                        Some(CallbackTimer::new(Millis(10_000), move || {
1063                            new_calls.set(new_calls.get() + 1);
1064                        }));
1065                    released.set(true);
1066                })
1067            };
1068            let handle = Rc::new(RefCell::new(None::<CallbackTimer>));
1069            let weak = Rc::downgrade(&handle);
1070            let was_released = released.clone();
1071            *handle.borrow_mut() = Some(CallbackTimer::new(Millis(1), move || {
1072                let _capture = &capture;
1073                assert!(active, "dropped queued callback must not run");
1074                drop(weak.upgrade().unwrap().borrow_mut().take());
1075                assert!(!was_released.get());
1076                finished.set(true);
1077            }));
1078            let old_slot = handle.borrow().as_ref().unwrap().hnd;
1079            if !active {
1080                timer.with_wheel(|wheel| {
1081                    wheel.relink(old_slot.get(), 1);
1082                    wheel.execute_expired_timers(1);
1083                });
1084                drop(handle.borrow_mut().take());
1085                assert!(released.get());
1086            }
1087            make_due(&timer);
1088            let mut driver = timer_driver(&timer);
1089            assert!(
1090                Pin::new(&mut driver)
1091                    .poll(&mut Context::from_waker(Waker::noop()))
1092                    .is_pending()
1093            );
1094            assert!(released.get());
1095            assert!(handle.borrow().is_none());
1096            assert!(cancelled.is_elapsed());
1097            assert!(!reset.is_elapsed());
1098            assert_eq!(replacement.borrow().as_ref().unwrap().hnd, old_slot);
1099            assert!(!replacement.borrow().as_ref().unwrap().is_elapsed());
1100            assert_eq!((reset_calls.get(), new_calls.get()), (0, 0));
1101
1102            timer.elapsed_time.set(Some(
1103                Instant::now().checked_sub(Duration::from_secs(20)).unwrap(),
1104            ));
1105            for _ in 0..2 {
1106                assert!(
1107                    Pin::new(&mut driver)
1108                        .poll(&mut Context::from_waker(Waker::noop()))
1109                        .is_pending()
1110                );
1111                assert_eq!((reset_calls.get(), new_calls.get()), (1, 1));
1112            }
1113        }
1114    }
1115
1116    #[test]
1117    fn test_callback_capture_drop_resets_runtime() {
1118        let timer = timer_state();
1119        let new_calls = Rc::new(Cell::new(0));
1120        let replacement = Rc::new(RefCell::new(None::<CallbackTimer>));
1121        let capture = {
1122            let timer = timer.clone();
1123            let new_calls = new_calls.clone();
1124            let replacement = replacement.clone();
1125            OnDropFn::new(move || {
1126                timer.reset();
1127                timer.insert_flags(Flags::DRIVER_STARTED);
1128                let callback = CallbackTimer::new(Millis(1), move || {
1129                    new_calls.set(new_calls.get() + 1);
1130                });
1131                timer.with_wheel(|wheel| {
1132                    wheel.relink(callback.hnd.get(), 1);
1133                    wheel.execute_expired_timers(1);
1134                });
1135                *replacement.borrow_mut() = Some(callback);
1136            })
1137        };
1138        let handle = Rc::new(RefCell::new(None::<CallbackTimer>));
1139        let weak = Rc::downgrade(&handle);
1140        *handle.borrow_mut() = Some(CallbackTimer::new(Millis(1), move || {
1141            let _capture = &capture;
1142            drop(weak.upgrade().unwrap().borrow_mut().take());
1143        }));
1144        make_due(&timer);
1145        let mut old_driver = timer_driver(&timer);
1146        assert!(
1147            Pin::new(&mut old_driver)
1148                .poll(&mut Context::from_waker(Waker::noop()))
1149                .is_ready()
1150        );
1151        assert!(handle.borrow().is_none());
1152        assert!(replacement.borrow().as_ref().unwrap().is_elapsed());
1153        assert_eq!(new_calls.get(), 0);
1154        drop(old_driver);
1155
1156        make_due(&timer);
1157        let mut driver = timer_driver(&timer);
1158        for _ in 0..2 {
1159            assert!(
1160                Pin::new(&mut driver)
1161                    .poll(&mut Context::from_waker(Waker::noop()))
1162                    .is_pending()
1163            );
1164            assert_eq!(new_calls.get(), 1);
1165        }
1166    }
1167
1168    #[test]
1169    fn test_callback_refreshes_time() {
1170        let timer = timer_state();
1171        let calls = Rc::new(Cell::new(0));
1172        let handles = [1, 32].map(|millis| {
1173            let calls = calls.clone();
1174            CallbackTimer::new(Millis(millis), move || {
1175                let previous = calls.replace(calls.get() + 1);
1176                if previous == 0 {
1177                    std::thread::sleep(Duration::from_millis(64));
1178                }
1179            })
1180        });
1181        timer.with_wheel(|wheel| {
1182            wheel.relink(handles[0].hnd.get(), 1);
1183            wheel.relink(handles[1].hnd.get(), 3);
1184        });
1185        timer.elapsed_time.set(Some(
1186            Instant::now()
1187                .checked_sub(Duration::from_millis(32))
1188                .unwrap(),
1189        ));
1190        let mut driver = timer_driver(&timer);
1191        assert!(
1192            Pin::new(&mut driver)
1193                .poll(&mut Context::from_waker(Waker::noop()))
1194                .is_pending()
1195        );
1196        assert_eq!(calls.get(), 2);
1197        assert!(handles.iter().all(CallbackTimer::is_elapsed));
1198    }
1199
1200    #[test]
1201    fn test_callback_panic_after_rearm_cleanup() {
1202        let timer = timer_state();
1203        let callback = Rc::new_cyclic(|weak: &std::rc::Weak<CallbackTimer>| {
1204            let weak = weak.clone();
1205            CallbackTimer::new(Millis(1), move || {
1206                let callback = weak.upgrade().unwrap();
1207                callback.reset(Millis(10_000));
1208                panic!("callback");
1209            })
1210        });
1211        let later = CallbackTimer::new(Millis(10_000), || panic!("not due"));
1212        let ordinary = TimerHandle::new(10_000);
1213        make_due(&timer);
1214        let mut driver = timer_driver(&timer);
1215        assert!(
1216            catch_unwind(AssertUnwindSafe(|| {
1217                let _ = Pin::new(&mut driver).poll(&mut Context::from_waker(Waker::noop()));
1218            }))
1219            .is_err()
1220        );
1221        assert!(!callback.is_elapsed());
1222        drop(driver);
1223        assert!(callback.is_elapsed());
1224        assert!(!callback.state.pending.get());
1225        assert!(later.is_elapsed());
1226        assert!(ordinary.is_elapsed());
1227        assert!(!later.state.pending.get());
1228        timer.with_wheel(|wheel| {
1229            assert!(wheel.buckets.iter().all(Slab::is_empty));
1230            assert_eq!(wheel.occupied, [0; LVL_DEPTH as usize]);
1231            assert!(wheel.callbacks.as_ref().unwrap().ready.is_empty());
1232        });
1233    }
1234
1235    #[test]
1236    fn test_callback_runtime_reset_during_delivery() {
1237        let timer = timer_state();
1238        let calls = Rc::new(Cell::new(0));
1239        let new_calls = Rc::new(Cell::new(0));
1240        let replacement = Rc::new(RefCell::new(None::<CallbackTimer>));
1241        let handles = [0, 1].map(|_| {
1242            let timer = timer.clone();
1243            let calls = calls.clone();
1244            let new_calls = new_calls.clone();
1245            let replacement = replacement.clone();
1246            CallbackTimer::new(Millis(1), move || {
1247                calls.set(calls.get() + 1);
1248                timer.reset();
1249                timer.insert_flags(Flags::DRIVER_STARTED);
1250                let new_calls = new_calls.clone();
1251                *replacement.borrow_mut() = Some(CallbackTimer::new(Millis(1), move || {
1252                    new_calls.set(new_calls.get() + 1);
1253                }));
1254            })
1255        });
1256        timer.with_wheel(|wheel| {
1257            for handle in &handles {
1258                wheel.relink(handle.hnd.get(), 1);
1259            }
1260        });
1261        make_due(&timer);
1262        let mut old_driver = timer_driver(&timer);
1263        assert!(
1264            Pin::new(&mut old_driver)
1265                .poll(&mut Context::from_waker(Waker::noop()))
1266                .is_ready()
1267        );
1268        assert_eq!(calls.get(), 1);
1269        assert!(handles.iter().all(CallbackTimer::is_elapsed));
1270        drop(old_driver);
1271        assert!(!replacement.borrow().as_ref().unwrap().is_elapsed());
1272        make_due(&timer);
1273        let mut driver = timer_driver(&timer);
1274        assert!(
1275            Pin::new(&mut driver)
1276                .poll(&mut Context::from_waker(Waker::noop()))
1277                .is_pending()
1278        );
1279        assert_eq!(new_calls.get(), 1);
1280    }
1281
1282    #[test]
1283    fn test_callback_mixed_bucket_waker_panic() {
1284        for panic_position in 0..3 {
1285            let timer = timer_state();
1286            let calls = Rc::new(Cell::new(0));
1287            let callbacks = [0, 1].map(|_| {
1288                let calls = calls.clone();
1289                CallbackTimer::new(Millis(10_000), move || calls.set(calls.get() + 1))
1290            });
1291            let ordinary = TimerHandle::new(10_000);
1292            let task = Arc::new(InPlaceWaker {
1293                panic_on_wake: true,
1294                ..InPlaceWaker::default()
1295            });
1296            let waker = Waker::from(task.clone());
1297            assert!(
1298                ordinary
1299                    .poll_elapsed(&mut Context::from_waker(&waker))
1300                    .is_pending()
1301            );
1302            timer.with_wheel(|wheel| {
1303                // Control slab drain positions, not public callback delivery order.
1304                let mut entries = [
1305                    callbacks[0].hnd.get(),
1306                    ordinary.0.get(),
1307                    callbacks[1].hnd.get(),
1308                ];
1309                entries.swap(1, panic_position);
1310                for no in entries {
1311                    wheel.unlink(no);
1312                }
1313                wheel.buckets[3] = Slab::new();
1314                for no in entries {
1315                    wheel.link(no, 3);
1316                }
1317            });
1318            timer.next_expiry.set(3);
1319            make_due(&timer);
1320            let mut driver = timer_driver(&timer);
1321            assert!(
1322                catch_unwind(AssertUnwindSafe(|| {
1323                    let _ = Pin::new(&mut driver).poll(&mut Context::from_waker(Waker::noop()));
1324                }))
1325                .is_err()
1326            );
1327            assert_eq!(task.wakes.load(Ordering::Relaxed), 1);
1328            assert_eq!(calls.get(), 0);
1329            assert!(ordinary.is_elapsed());
1330            assert!(callbacks.iter().all(CallbackTimer::is_elapsed));
1331            timer.with_wheel(|wheel| {
1332                assert!(wheel.buckets.iter().all(Slab::is_empty));
1333                assert_eq!(wheel.occupied, [0; LVL_DEPTH as usize]);
1334                assert_eq!(
1335                    wheel.callbacks.as_ref().unwrap().ready.len(),
1336                    panic_position
1337                );
1338            });
1339            drop(driver);
1340            assert!(
1341                callbacks
1342                    .iter()
1343                    .all(|callback| !callback.state.pending.get())
1344            );
1345            timer.with_wheel(|wheel| {
1346                assert!(wheel.callbacks.as_ref().unwrap().ready.is_empty());
1347            });
1348
1349            timer.insert_flags(Flags::DRIVER_STARTED);
1350            callbacks[0].reset(Millis::ZERO);
1351            make_due(&timer);
1352            let mut driver = timer_driver(&timer);
1353            for _ in 0..2 {
1354                assert!(
1355                    Pin::new(&mut driver)
1356                        .poll(&mut Context::from_waker(Waker::noop()))
1357                        .is_pending()
1358                );
1359                assert_eq!(calls.get(), 1);
1360            }
1361        }
1362    }
1363
1364    #[test]
1365    fn test_callback_reuse_in_next_system() {
1366        std::thread::spawn(|| {
1367            let calls = Rc::new(Cell::new(0));
1368            let mut callback = None::<CallbackTimer>;
1369            for expected in 1..=2 {
1370                let calls = calls.clone();
1371                callback = Some(
1372                    ntex::rt::System::build()
1373                        .build(ntex::rt::DefaultRuntime)
1374                        .block_on(async move {
1375                            let callback = if let Some(callback) = callback {
1376                                assert!(callback.is_elapsed());
1377                                sleep(Millis(50)).await;
1378                                assert_eq!(calls.get(), expected - 1);
1379                                callback.reset(Millis::ZERO);
1380                                callback
1381                            } else {
1382                                let calls = calls.clone();
1383                                CallbackTimer::new(Millis::ZERO, move || {
1384                                    calls.set(calls.get() + 1);
1385                                })
1386                            };
1387                            sleep(Millis(50)).await;
1388                            assert_eq!(calls.get(), expected);
1389                            callback.reset(Millis(10_000));
1390                            callback
1391                        }),
1392                );
1393                assert!(callback.as_ref().unwrap().is_elapsed());
1394            }
1395        })
1396        .join()
1397        .unwrap();
1398    }
1399
1400    #[test]
1401    fn test_wheel_access_restored_on_unwind() {
1402        let timer = Timer::new();
1403        assert!(
1404            catch_unwind(AssertUnwindSafe(|| {
1405                timer.with_wheel(|wheel| {
1406                    wheel.timers.insert(entry());
1407                    panic!("wheel operation");
1408                });
1409            }))
1410            .is_err()
1411        );
1412        timer.with_wheel(|wheel| assert_eq!(wheel.timers.len(), 2));
1413    }
1414
1415    #[test]
1416    fn test_wake_in_place() {
1417        let timer = timer_state();
1418        let task = Arc::new(InPlaceWaker::default());
1419        let waker = Waker::from(task.clone());
1420        let mut cx = Context::from_waker(&waker);
1421        for reset in [false, true] {
1422            let handle = TimerHandle::new(10_000);
1423            assert!(handle.poll_elapsed(&mut cx).is_pending());
1424            if reset {
1425                handle.reset(0);
1426            } else {
1427                handle.elapse();
1428            }
1429            assert!(handle.is_elapsed());
1430        }
1431        let handles = [TimerHandle::new(1), TimerHandle::new(1)];
1432        for handle in &handles {
1433            assert!(handle.poll_elapsed(&mut cx).is_pending());
1434        }
1435        timer.elapsed_time.set(Some(
1436            Instant::now().checked_sub(Duration::from_secs(20)).unwrap(),
1437        ));
1438        let mut driver = timer_driver(&timer);
1439        assert!(
1440            Pin::new(&mut driver)
1441                .poll(&mut Context::from_waker(Waker::noop()))
1442                .is_pending()
1443        );
1444        assert!(handles.iter().all(TimerHandle::is_elapsed));
1445        assert_eq!(task.wakes.load(Ordering::Relaxed), 4);
1446        assert_eq!(Arc::strong_count(&task), 2);
1447    }
1448
1449    #[test]
1450    fn test_dispatch_uses_one_time_sample_for_wakers() {
1451        let timer = timer_state();
1452        let handles = [TimerHandle::new(1), TimerHandle::new(32)];
1453        timer.with_wheel(|wheel| wheel.relink(handles[1].0.get(), 3));
1454        let task = Arc::new(InPlaceWaker {
1455            delay: Duration::from_millis(64),
1456            ..InPlaceWaker::default()
1457        });
1458        let waker = Waker::from(task.clone());
1459        for handle in &handles {
1460            assert!(
1461                handle
1462                    .poll_elapsed(&mut Context::from_waker(&waker))
1463                    .is_pending()
1464            );
1465        }
1466        timer.elapsed_time.set(Some(
1467            Instant::now()
1468                .checked_sub(Duration::from_millis(32))
1469                .unwrap(),
1470        ));
1471        let mut driver = timer_driver(&timer);
1472        assert!(
1473            Pin::new(&mut driver)
1474                .poll(&mut Context::from_waker(Waker::noop()))
1475                .is_pending()
1476        );
1477        assert!(handles[0].is_elapsed());
1478        assert!(!handles[1].is_elapsed());
1479        assert_eq!(task.wakes.load(Ordering::Relaxed), 1);
1480        assert!(
1481            Pin::new(&mut driver)
1482                .poll(&mut Context::from_waker(Waker::noop()))
1483                .is_pending()
1484        );
1485        assert!(handles[1].is_elapsed());
1486        assert_eq!(task.wakes.load(Ordering::Relaxed), 2);
1487    }
1488
1489    #[test]
1490    fn test_dispatch_panic_cleans_expired_bucket() {
1491        let timer = timer_state();
1492        let handles = [TimerHandle::new(1), TimerHandle::new(1)];
1493        timer.with_wheel(|wheel| {
1494            let bucket = wheel.timers[handles[0].0.get()].bucket.unwrap();
1495            wheel.relink(handles[1].0.get(), usize::from(bucket));
1496        });
1497        let task = Arc::new(InPlaceWaker {
1498            panic_on_wake: true,
1499            ..InPlaceWaker::default()
1500        });
1501        let waker = Waker::from(task.clone());
1502        for handle in &handles {
1503            assert!(
1504                handle
1505                    .poll_elapsed(&mut Context::from_waker(&waker))
1506                    .is_pending()
1507            );
1508        }
1509        timer.elapsed_time.set(Some(
1510            Instant::now()
1511                .checked_sub(Duration::from_millis(64))
1512                .unwrap(),
1513        ));
1514        let mut driver = timer_driver(&timer);
1515        assert!(
1516            catch_unwind(AssertUnwindSafe(|| {
1517                let _ = Pin::new(&mut driver).poll(&mut Context::from_waker(Waker::noop()));
1518            }))
1519            .is_err()
1520        );
1521        assert_eq!(task.wakes.load(Ordering::Relaxed), 1);
1522        assert!(handles.iter().all(TimerHandle::is_elapsed));
1523        timer.with_wheel(|wheel| {
1524            assert!(wheel.buckets.iter().all(Slab::is_empty));
1525            assert_eq!(wheel.occupied, [0; LVL_DEPTH as usize]);
1526        });
1527        for handle in &handles {
1528            handle.reset(10_000);
1529            assert!(!handle.is_elapsed());
1530        }
1531        drop(driver);
1532        assert!(handles.iter().all(TimerHandle::is_elapsed));
1533        assert!(!timer.flags.get().contains(Flags::DRIVER_STARTED));
1534    }
1535
1536    /// A handle dropped by a thread-local destructor after the wheel is
1537    /// destroyed must not panic.
1538    #[test]
1539    fn test_drop_handle_after_wheel_destroyed() {
1540        use std::cell::RefCell;
1541
1542        thread_local! {
1543            static HOLDER: RefCell<Option<TimerHandle>> = const { RefCell::new(None) };
1544        }
1545
1546        let res = std::thread::spawn(|| {
1547            // register the holder destructor first, the wheel is destroyed
1548            // before it, destructors run in reverse registration order
1549            HOLDER.with(|h| h.borrow_mut().take());
1550            ntex::rt::System::build()
1551                .build(ntex::rt::DefaultRuntime)
1552                .block_on(async {
1553                    let hnd = TimerHandle::new(1000);
1554                    HOLDER.with(|h| *h.borrow_mut() = Some(hnd));
1555                });
1556        })
1557        .join();
1558        assert!(res.is_ok());
1559    }
1560
1561    #[test]
1562    fn test_drop_callback_after_wheel_destroyed() {
1563        thread_local! {
1564            static HOLDER: RefCell<Option<CallbackTimer>> = const { RefCell::new(None) };
1565        }
1566
1567        let drops = Arc::new(AtomicUsize::new(0));
1568        let destroyed = Arc::new(AtomicBool::new(false));
1569        let count = drops.clone();
1570        let was_destroyed = destroyed.clone();
1571        std::thread::spawn(move || {
1572            // Register this destructor before the timer's thread-local storage.
1573            HOLDER.with(|holder| holder.borrow_mut().take());
1574            ntex::rt::System::build()
1575                .build(ntex::rt::DefaultRuntime)
1576                .block_on(async move {
1577                    let ordinary = TimerHandle::new(10_000);
1578                    let capture = OnDropFn::new(move || {
1579                        was_destroyed.store(TIMER.try_with(|_| ()).is_err(), Ordering::Relaxed);
1580                        drop(ordinary);
1581                        count.fetch_add(1, Ordering::Relaxed);
1582                    });
1583                    let callback = CallbackTimer::new(Millis(10_000), move || {
1584                        let _capture = &capture;
1585                        panic!("callback after shutdown");
1586                    });
1587                    HOLDER.with(|holder| *holder.borrow_mut() = Some(callback));
1588                });
1589        })
1590        .join()
1591        .unwrap();
1592        assert!(destroyed.load(Ordering::Relaxed));
1593        assert_eq!(drops.load(Ordering::Relaxed), 1);
1594    }
1595
1596    /// The drivers of a stopped system may never be dropped, the next system
1597    /// on the same thread starts its own drivers.
1598    #[test]
1599    fn test_timer_in_next_system() {
1600        let res = std::thread::spawn(|| {
1601            for _ in 0..3 {
1602                let start = Instant::now();
1603                ntex::rt::System::build()
1604                    .build(ntex::rt::DefaultRuntime)
1605                    .block_on(async {
1606                        sleep(Millis(10)).await;
1607                        let _ = now();
1608                        // pending timer of the stopped system
1609                        let _hnd = sleep(Millis(10_000));
1610                        crate::spawn(async { sleep(Millis(10_000)).await });
1611                    });
1612                assert!(start.elapsed() < Duration::from_secs(5));
1613            }
1614        });
1615        let (tx, rx) = std::sync::mpsc::channel();
1616        std::thread::spawn(move || tx.send(res.join().is_ok()));
1617        assert_eq!(rx.recv_timeout(Duration::from_secs(10)), Ok(true));
1618    }
1619
1620    /// A stale driver does not stop the drivers of a newer generation.
1621    #[test]
1622    fn test_stale_driver_drop() {
1623        let timer = Rc::new(Timer::new());
1624        timer.insert_flags(Flags::RUNNING | Flags::DRIVER_STARTED | Flags::LOWRES_DRIVER);
1625        timer.reset();
1626        assert_eq!(timer.flags.get(), Flags::empty());
1627
1628        timer.insert_flags(Flags::RUNNING | Flags::DRIVER_STARTED | Flags::LOWRES_DRIVER);
1629        drop(TimerDriver {
1630            timer: timer.clone(),
1631            generation: 0,
1632            sleep: Delay::new(Duration::ZERO),
1633            armed: None,
1634        });
1635        drop(LowresTimerDriver {
1636            timer: timer.clone(),
1637            generation: 0,
1638            sleep: Delay::new(Duration::ZERO),
1639        });
1640        assert_eq!(
1641            timer.flags.get(),
1642            Flags::RUNNING | Flags::DRIVER_STARTED | Flags::LOWRES_DRIVER
1643        );
1644    }
1645
1646    fn entry() -> TimerEntry {
1647        TimerEntry {
1648            bucket: None,
1649            bucket_entry: 0,
1650            callback: false,
1651            task: None,
1652        }
1653    }
1654
1655    /// Bucket expiry is never earlier than the requested delay, and at most
1656    /// one level granularity later.
1657    #[test]
1658    fn test_bucket_expiry_bounds() {
1659        let timer = Timer::new();
1660        for elapsed in [0, 1, 7, 8, 63, 64, 511, 12_345, 1 << 30] {
1661            timer.elapsed.set(elapsed);
1662            let mut delta = 0;
1663            while delta < WHEEL_TIMEOUT_CUTOFF {
1664                let (idx, expiry) = timer.calc_wheel_index(elapsed + delta, delta);
1665                let lvl = (idx / LVL_SIZE as usize) as u64;
1666                assert!(expiry > elapsed + delta, "{elapsed} {delta} {expiry}");
1667                assert!(
1668                    expiry <= elapsed + delta + lvl_gran(lvl),
1669                    "{elapsed} {delta} {expiry}"
1670                );
1671                assert_eq!(expiry % lvl_gran(lvl), 0);
1672                assert_eq!(
1673                    ((expiry >> lvl_shift(lvl)) & LVL_MASK) as usize,
1674                    idx % LVL_SIZE as usize
1675                );
1676                delta = delta * 2 + 1;
1677            }
1678
1679            // clamped to the capacity of the wheel
1680            let (_, expiry) =
1681                timer.calc_wheel_index(elapsed + u64::from(u32::MAX), u64::from(u32::MAX));
1682            assert!(expiry <= elapsed + WHEEL_TIMEOUT_CUTOFF);
1683        }
1684    }
1685
1686    /// The next pending bucket is the earliest occupied one, and executing it
1687    /// wakes only its timers.
1688    #[test]
1689    fn test_next_pending_bucket() {
1690        let timer = Timer::new();
1691        timer.elapsed.set(1000);
1692
1693        timer.with_wheel(|w| {
1694            assert_eq!(timer.next_pending_bucket(w), None);
1695
1696            let mut expected = Vec::new();
1697            for delta in [5000, 30, 700, 100_000] {
1698                let (idx, expiry) = timer.calc_wheel_index(1000 + delta, delta);
1699                let no = w.timers.insert(entry());
1700                w.link(no, idx);
1701                expected.push((expiry, no));
1702            }
1703            expected.sort_unstable();
1704
1705            for (expiry, no) in expected {
1706                assert_eq!(timer.next_pending_bucket(w), Some(expiry));
1707                timer.elapsed.set(expiry);
1708                w.execute_expired_timers(expiry);
1709                assert!(w.timers[no].bucket.is_none());
1710            }
1711            assert_eq!(timer.next_pending_bucket(w), None);
1712            assert_eq!(w.occupied, [0; LVL_DEPTH as usize]);
1713        });
1714    }
1715
1716    #[test]
1717    fn test_unlink_relink() {
1718        let timer = Timer::new();
1719        timer.with_wheel(|w| {
1720            let a = w.timers.insert(entry());
1721            let b = w.timers.insert(entry());
1722            w.link(a, 3);
1723            w.link(b, 3);
1724            assert_eq!(w.occupied[0], 1 << 3);
1725
1726            assert!(w.unlink(a));
1727            assert!(!w.unlink(a));
1728            assert_eq!(w.occupied[0], 1 << 3);
1729
1730            w.relink(b, LVL_SIZE as usize + 5);
1731            assert_eq!(w.occupied[0], 0);
1732            assert_eq!(w.occupied[1], 1 << 5);
1733            assert_eq!(w.timers[b].bucket, Some(LVL_SIZE as u16 + 5));
1734
1735            assert!(w.unlink(b));
1736            assert_eq!(w.occupied, [0; LVL_DEPTH as usize]);
1737        });
1738    }
1739
1740    /// `reset(0)` elapses the timer and wakes the task waiting for it.
1741    #[ntex::test]
1742    async fn test_reset_zero_wakes_task() {
1743        let hnd = Rc::new(TimerHandle::new(10_000));
1744        let hnd2 = hnd.clone();
1745        crate::spawn(async move { hnd2.reset(0) });
1746
1747        // the sleep wakes the task if `reset(0)` does not
1748        let start = Instant::now();
1749        crate::future::select(
1750            std::future::poll_fn(|cx| hnd.poll_elapsed(cx)),
1751            sleep(Millis(500)),
1752        )
1753        .await;
1754        let elapsed = start.elapsed();
1755        assert!(elapsed < Duration::from_millis(100), "elapsed: {elapsed:?}");
1756    }
1757
1758    /// Dropping one driver does not reset the state of the other one.
1759    #[test]
1760    fn test_driver_drop_keeps_other_driver() {
1761        let timer = Rc::new(Timer::new());
1762        let no = timer.with_wheel(|w| {
1763            let no = w.timers.insert(entry());
1764            w.link(no, 3);
1765            no
1766        });
1767        timer.next_expiry.set(3);
1768        timer.lowres_stime.set(Some(SystemTime::now()));
1769        timer.insert_flags(
1770            Flags::RUNNING | Flags::DRIVER_STARTED | Flags::LOWRES_DRIVER | Flags::LOWRES_TIMER,
1771        );
1772
1773        drop(LowresTimerDriver {
1774            timer: timer.clone(),
1775            generation: 0,
1776            sleep: Delay::new(Duration::ZERO),
1777        });
1778        assert_eq!(timer.flags.get(), Flags::DRIVER_STARTED);
1779        assert!(timer.lowres_stime.get().is_none());
1780        // timers are still pending
1781        assert_eq!(timer.next_expiry.get(), 3);
1782        timer.with_wheel(|w| assert_eq!(w.timers[no].bucket, Some(3)));
1783
1784        timer.insert_flags(Flags::RUNNING | Flags::LOWRES_DRIVER);
1785        drop(TimerDriver {
1786            timer: timer.clone(),
1787            generation: 0,
1788            sleep: Delay::new(Duration::ZERO),
1789            armed: None,
1790        });
1791        assert_eq!(timer.flags.get(), Flags::LOWRES_DRIVER);
1792        assert_eq!(timer.next_expiry.get(), u64::MAX);
1793        timer.with_wheel(|w| {
1794            assert!(w.timers[no].bucket.is_none());
1795            assert_eq!(w.occupied, [0; LVL_DEPTH as usize]);
1796        });
1797    }
1798
1799    /// A short timer is measured from the current time, not from the cached
1800    /// time that went stale while the thread was blocked.
1801    #[ntex::test]
1802    async fn test_short_timer_after_blocking() {
1803        let _hnd = sleep(Millis(10_000));
1804        let _ = now();
1805
1806        // the lowres driver cannot run, the cached time goes stale
1807        std::thread::sleep(Duration::from_millis(100));
1808
1809        let start = Instant::now();
1810        sleep(Millis(50)).await;
1811        let elapsed = start.elapsed();
1812        assert!(elapsed >= Duration::from_millis(50), "elapsed: {elapsed:?}");
1813    }
1814
1815    /// `now()` does not go stale while the wheel is idle.
1816    #[ntex::test]
1817    async fn test_driver_cached_time_expires() {
1818        // the lowres timer of the cache populated at start fires first
1819        sleep(Millis(400)).await;
1820
1821        // the wheel is idle, the thread waits for a non-timer event
1822        let (tx, rx) = async_channel::bounded(1);
1823        std::thread::spawn(move || {
1824            std::thread::sleep(Duration::from_millis(600));
1825            let _ = tx.send_blocking(());
1826        });
1827        rx.recv().await.unwrap();
1828
1829        let stale = Instant::now().saturating_duration_since(now());
1830        assert!(stale < Duration::from_millis(300), "stale: {stale:?}");
1831    }
1832
1833    /// A late wakeup must not delay the timers that expire later.
1834    #[ntex::test]
1835    async fn test_late_wakeup_does_not_drift() {
1836        let start = Instant::now();
1837        let fut1 = sleep(Millis(50));
1838        let fut2 = sleep(Millis(600));
1839
1840        // block the thread, the driver wakes up ~450ms late for `fut1`.
1841        // `fut2` expires at ~620ms, with drift it would expire ~550ms after
1842        // the late wakeup, i.e. after ~1050ms. The bound leaves room for
1843        // scheduling latency on loaded machines.
1844        std::thread::sleep(Duration::from_millis(500));
1845        fut1.await;
1846        fut2.await;
1847
1848        let elapsed = start.elapsed();
1849        assert!(
1850            elapsed >= Duration::from_millis(600) && elapsed < Duration::from_millis(900),
1851            "elapsed: {elapsed:?}"
1852        );
1853    }
1854
1855    #[ntex::test]
1856    #[allow(unused_variables, clippy::used_underscore_binding)]
1857    async fn test_timer() {
1858        crate::spawn(async {
1859            let s = interval(Millis(25));
1860            loop {
1861                s.tick().await;
1862            }
1863        });
1864        let time = Instant::now();
1865        let fut1 = sleep(Millis(1000));
1866        let fut2 = sleep(Millis(200));
1867
1868        fut2.await;
1869        #[cfg(not(target_os = "macos"))]
1870        {
1871            let _elapsed = time.elapsed();
1872            assert!(
1873                _elapsed > Duration::from_millis(200) && _elapsed < Duration::from_millis(300),
1874                "elapsed: {_elapsed:?}"
1875            );
1876        }
1877
1878        fut1.await;
1879
1880        #[cfg(not(target_os = "macos"))]
1881        {
1882            let _elapsed = time.elapsed();
1883            assert!(
1884                _elapsed > Duration::from_secs(1) && _elapsed < Duration::from_millis(1200), // osx
1885                "elapsed: {_elapsed:?}",
1886            );
1887        }
1888
1889        let time = Instant::now();
1890        sleep(Millis(25)).await;
1891        #[cfg(not(target_os = "macos"))]
1892        {
1893            let _elapsed = time.elapsed();
1894            assert!(
1895                _elapsed > Duration::from_millis(20) && _elapsed < Duration::from_millis(50),
1896                "elapsed: {_elapsed:?}",
1897            );
1898        }
1899    }
1900}