    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89
    90
    91
    92
    93
    94
    95
    96
    97
    98
    99
    100
    101
    102
    103
    104

    //! Utility for tracking network requests that will be retried in the future.

    use core::cmp::Ordering;
    use std::collections::BinaryHeap;
    use std::time::{Duration, Instant};

    /// A tracker for network requests that have failed, and are awaiting to be
    /// retried in the future.
    pub struct SleepTracker<T> {
        /// This is a priority queue that tracks the time when the next sleeper
        /// should awaken (based on the [`Sleeper::wakeup`] property).
        heap: BinaryHeap<Sleeper<T>>,
    }

    /// An individual network request that is waiting to be retried in the future.
    struct Sleeper<T> {
        /// The time when this requests should be retried.
        wakeup: Instant,
        /// Information about the network request.
        data: T,
    }

    impl<T> PartialEq for Sleeper<T> {
        fn eq(&self, other: &Sleeper<T>) -> bool {
            self.wakeup == other.wakeup
        }
    }

    impl<T> PartialOrd for Sleeper<T> {
        fn partial_cmp(&self, other: &Sleeper<T>) -> Option<Ordering> {
            // This reverses the comparison so that the BinaryHeap tracks the
            // entry with the *lowest* wakeup time.
            Some(other.wakeup.cmp(&self.wakeup))
        }
    }

    impl<T> Eq for Sleeper<T> {}

    impl<T> Ord for Sleeper<T> {
        fn cmp(&self, other: &Sleeper<T>) -> Ordering {
            self.wakeup.cmp(&other.wakeup)
        }
    }

    impl<T> SleepTracker<T> {
        pub fn new() -> SleepTracker<T> {
            SleepTracker {
                heap: BinaryHeap::new(),
            }
        }

        /// Adds a new download that should be retried in the future.
        pub fn push(&mut self, sleep: u64, data: T) {
            self.heap.push(Sleeper {
                wakeup: Instant::now()
                    .checked_add(Duration::from_millis(sleep))
                    .expect("instant should not wrap"),
                data,
            });
        }

        pub fn len(&self) -> usize {
            self.heap.len()
        }

        /// Returns any downloads that are ready to go now.
        pub fn to_retry(&mut self) -> Vec<T> {
            let now = Instant::now();
            let mut result = Vec::new();
            while let Some(next) = self.heap.peek() {
                if next.wakeup < now {
                    result.push(self.heap.pop().unwrap().data);
                } else {
                    break;
                }
            }
            result
        }

        /// Returns the time when the next download is ready to go.
        ///
        /// Returns None if there are no sleepers remaining.
        pub fn time_to_next(&self) -> Option<Duration> {
            self.heap
                .peek()
                .map(|s| s.wakeup.saturating_duration_since(Instant::now()))
        }
    }

    #[test]
    fn returns_in_order() {
        let mut s = SleepTracker::new();
        s.push(30_000, 30_000);
        s.push(1, 1);
        assert_eq!(s.len(), 2);
        std::thread::sleep(Duration::from_millis(2));
        assert_eq!(s.to_retry(), &[1]);
        assert!(s.to_retry().is_empty());
        let next = s.time_to_next().expect("should be next");
        assert!(
            next < Duration::from_millis(30_000),
            "{next:?} should be less than 30s"
        );
    }