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}
cargo/util/network/sleep.rs

cargo/util/network/
sleep.rs
