#
cargo/util/network/
sleep.rs

    1//! Utility for tracking network requests that will be retried in the future.
    2
    3use core::cmp::Ordering;
    4use std::collections::BinaryHeap;
    5use std::time::{Duration, Instant};
    6
    7/// A tracker for network requests that have failed, and are awaiting to be
    8/// retried in the future.
    9pub struct SleepTracker<T> {
    10    /// This is a priority queue that tracks the time when the next sleeper
    11    /// should awaken (based on the [`Sleeper::wakeup`] property).
    12    heap: BinaryHeap<Sleeper<T>>,
    13}
    14
    15/// An individual network request that is waiting to be retried in the future.
    16struct Sleeper<T> {
    17    /// The time when this requests should be retried.
    18    wakeup: Instant,
    19    /// Information about the network request.
    20    data: T,
    21}
    22
    23impl<T> PartialEq for Sleeper<T> {
    24    fn eq(&self, other: &Sleeper<T>) -> bool {
    25        self.wakeup == other.wakeup
    26    }
    27}
    28
    29impl<T> PartialOrd for Sleeper<T> {
    30    fn partial_cmp(&self, other: &Sleeper<T>) -> Option<Ordering> {
    31        // This reverses the comparison so that the BinaryHeap tracks the
    32        // entry with the *lowest* wakeup time.
    33        Some(other.wakeup.cmp(&self.wakeup))
    34    }
    35}
    36
    37impl<T> Eq for Sleeper<T> {}
    38
    39impl<T> Ord for Sleeper<T> {
    40    fn cmp(&self, other: &Sleeper<T>) -> Ordering {
    41        self.wakeup.cmp(&other.wakeup)
    42    }
    43}
    44
    45impl<T> SleepTracker<T> {
    46    pub fn new() -> SleepTracker<T> {
    47        SleepTracker {
    48            heap: BinaryHeap::new(),
    49        }
    50    }
    51
    52    /// Adds a new download that should be retried in the future.
    53    pub fn push(&mut self, sleep: u64, data: T) {
    54        self.heap.push(Sleeper {
    55            wakeup: Instant::now()
    56                .checked_add(Duration::from_millis(sleep))
    57                .expect("instant should not wrap"),
    58            data,
    59        });
    60    }
    61
    62    pub fn len(&self) -> usize {
    63        self.heap.len()
    64    }
    65
    66    /// Returns any downloads that are ready to go now.
    67    pub fn to_retry(&mut self) -> Vec<T> {
    68        let now = Instant::now();
    69        let mut result = Vec::new();
    70        while let Some(next) = self.heap.peek() {
    71            if next.wakeup < now {
    72                result.push(self.heap.pop().unwrap().data);
    73            } else {
    74                break;
    75            }
    76        }
    77        result
    78    }
    79
    80    /// Returns the time when the next download is ready to go.
    81    ///
    82    /// Returns None if there are no sleepers remaining.
    83    pub fn time_to_next(&self) -> Option<Duration> {
    84        self.heap
    85            .peek()
    86            .map(|s| s.wakeup.saturating_duration_since(Instant::now()))
    87    }
    88}
    89
    90#[test]
    91fn returns_in_order() {
    92    let mut s = SleepTracker::new();
    93    s.push(30_000, 30_000);
    94    s.push(1, 1);
    95    assert_eq!(s.len(), 2);
    96    std::thread::sleep(Duration::from_millis(2));
    97    assert_eq!(s.to_retry(), &[1]);
    98    assert!(s.to_retry().is_empty());
    99    let next = s.time_to_next().expect("should be next");
    100    assert!(
    101        next < Duration::from_millis(30_000),
    102        "{next:?} should be less than 30s"
    103    );
    104}