spire_core/frontier/facade/
retry.rs

1use std::cmp::{Ordering, Reverse};
2use std::collections::BinaryHeap;
3use std::fmt;
4use std::future::Future;
5use std::sync::Mutex;
6use std::time::Duration;
7
8use tokio::time::Instant;
9
10use crate::Result;
11use crate::crawler::Attempt;
12use crate::frontier::seed::Seed;
13use crate::frontier::{Ack, Admission, Frontier, NextSeed, Report};
14use crate::sync::lock;
15
16/// How often and how fast a [`Retry`] facade retries.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct RetryPolicy {
19    max_retries: u32,
20    base_delay: Duration,
21    max_delay: Duration,
22    jitter: u32,
23}
24
25impl RetryPolicy {
26    /// Retries up to `max_retries` times, waiting 500ms before the first
27    /// retry and doubling the delay up to 30s, shortened by up to 20% of
28    /// jitter.
29    #[must_use]
30    pub const fn new(max_retries: u32) -> Self {
31        Self {
32            max_retries,
33            base_delay: Duration::from_millis(500),
34            max_delay: Duration::from_secs(30),
35            jitter: 20,
36        }
37    }
38
39    /// Shortens each delay by up to `percent` of itself, so seeds failing
40    /// together are not retried together. Clamped to 100.
41    #[must_use]
42    pub const fn with_jitter(mut self, percent: u32) -> Self {
43        self.jitter = if percent > 100 { 100 } else { percent };
44        self
45    }
46
47    /// Sets the delay before the first retry.
48    #[must_use]
49    pub const fn with_base_delay(mut self, delay: Duration) -> Self {
50        self.base_delay = delay;
51        self
52    }
53
54    /// Sets the upper bound of the delay.
55    #[must_use]
56    pub const fn with_max_delay(mut self, delay: Duration) -> Self {
57        self.max_delay = delay;
58        self
59    }
60
61    /// Returns the delay before trying `attempt`, which is at least 1,
62    /// without jitter.
63    #[must_use]
64    pub fn delay(&self, attempt: Attempt) -> Duration {
65        let exponent = attempt.get().saturating_sub(1).min(31);
66        self.base_delay
67            .saturating_mul(1 << exponent)
68            .min(self.max_delay)
69    }
70
71    /// Returns the delay before `seed`'s attempt, with jitter derived from
72    /// its fingerprint and attempt, so it is stable for a given seed.
73    fn jittered(&self, seed: &Seed) -> Duration {
74        let delay = self.delay(seed.attempt());
75        let mix = seed
76            .fingerprint()
77            .get()
78            .wrapping_add(u64::from(seed.attempt().get()))
79            .wrapping_mul(0x9E37_79B9_7F4A_7C15);
80        let fraction = u32::try_from((mix >> 40) % 1000).unwrap_or_default();
81        let cut = delay * self.jitter / 100 * fraction / 1000;
82        delay.saturating_sub(cut)
83    }
84}
85
86impl Default for RetryPolicy {
87    fn default() -> Self {
88        Self::new(3)
89    }
90}
91
92/// A facade that re-queues seeds reported with [`Report::Retry`].
93///
94/// Retried seeds wait in the facade until their delay passes; they are not
95/// pushed again, so facades inside it never refuse them. Waiting seeds are
96/// held in memory; they are bounded by the requests that failed, not by the
97/// size of the frontier. Seeds out of
98/// retries are reported inwards as [`Report::Failed`] and refused with
99/// [`Ack::RetryRefused`].
100pub struct Retry<F> {
101    inner: F,
102    policy: RetryPolicy,
103    delayed: Mutex<Delayed>,
104}
105
106#[derive(Default)]
107struct Delayed {
108    heap: BinaryHeap<Reverse<Entry>>,
109    sequence: u64,
110}
111
112struct Entry {
113    at: Instant,
114    sequence: u64,
115    seed: Seed,
116}
117
118impl Ord for Entry {
119    fn cmp(&self, other: &Self) -> Ordering {
120        (self.at, self.sequence).cmp(&(other.at, other.sequence))
121    }
122}
123
124impl PartialOrd for Entry {
125    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
126        Some(self.cmp(other))
127    }
128}
129
130impl PartialEq for Entry {
131    fn eq(&self, other: &Self) -> bool {
132        self.cmp(other).is_eq()
133    }
134}
135
136impl Eq for Entry {}
137
138impl Delayed {
139    fn insert(&mut self, at: Instant, seed: Seed) {
140        self.sequence += 1;
141        self.heap.push(Reverse(Entry {
142            at,
143            sequence: self.sequence,
144            seed,
145        }));
146    }
147
148    fn pop_due(&mut self, now: Instant) -> Option<Seed> {
149        if self.earliest()? > now {
150            return None;
151        }
152        self.heap.pop().map(|Reverse(entry)| entry.seed)
153    }
154
155    fn earliest(&self) -> Option<Instant> {
156        self.heap.peek().map(|Reverse(entry)| entry.at)
157    }
158}
159
160impl<F> Retry<F> {
161    /// Wraps `inner`, retrying according to `policy`.
162    pub fn new(inner: F, policy: RetryPolicy) -> Self {
163        Self {
164            inner,
165            policy,
166            delayed: Mutex::default(),
167        }
168    }
169}
170
171impl<F: fmt::Debug> fmt::Debug for Retry<F> {
172    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
173        f.debug_struct("Retry")
174            .field("inner", &self.inner)
175            .field("policy", &self.policy)
176            .field("delayed", &lock(&self.delayed).heap.len())
177            .finish()
178    }
179}
180
181impl<F: Frontier> Frontier for Retry<F> {
182    fn push(&self, seed: Seed) -> impl Future<Output = Result<Admission>> + Send {
183        self.inner.push(seed)
184    }
185
186    async fn pop(&self) -> Result<NextSeed> {
187        let now = Instant::now();
188        if let Some(seed) = lock(&self.delayed).pop_due(now) {
189            return Ok(NextSeed::Ready(seed));
190        }
191
192        let next = self.inner.pop().await?;
193        let earliest = lock(&self.delayed).earliest();
194        Ok(match (next, earliest) {
195            (NextSeed::Wait(at), Some(earliest)) => NextSeed::Wait(at.min(earliest)),
196            (NextSeed::Empty, Some(earliest)) => NextSeed::Wait(earliest),
197            (next, _) => next,
198        })
199    }
200
201    async fn report(&self, report: Report) -> Result<Ack> {
202        let Report::Retry { seed, after } = report else {
203            return self.inner.report(report).await;
204        };
205
206        if seed.attempt().get() >= self.policy.max_retries {
207            tracing::debug!(
208                target: crate::TRACING_TARGET,
209                uri = %seed.uri(),
210                attempt = %seed.attempt(),
211                "retries exhausted",
212            );
213            self.inner.report(Report::Failed(seed)).await?;
214            return Ok(Ack::RetryRefused);
215        }
216
217        let seed = seed.next_attempt();
218        let delay = after.unwrap_or_else(|| self.policy.jittered(&seed));
219        lock(&self.delayed).insert(Instant::now() + delay, seed);
220        Ok(Ack::Applied)
221    }
222
223    fn size_hint(&self) -> Option<usize> {
224        let delayed = lock(&self.delayed).heap.len();
225        self.inner.size_hint().map(|len| len + delayed)
226    }
227
228    fn snapshot(&self) -> Option<Vec<Seed>> {
229        let mut seeds = self.inner.snapshot()?;
230        let delayed = lock(&self.delayed);
231        seeds.extend(delayed.heap.iter().map(|Reverse(entry)| entry.seed.clone()));
232        Some(seeds)
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use std::time::Duration;
239
240    use tokio::time::{Instant, advance};
241
242    use super::*;
243    use crate::frontier::test_util::{pop_path, seed};
244    use crate::frontier::{FrontierExt, InMemFrontier};
245
246    #[test]
247    fn policy_delay_doubles_and_caps() {
248        let policy = RetryPolicy::new(5)
249            .with_base_delay(Duration::from_secs(1))
250            .with_max_delay(Duration::from_secs(5))
251            .with_jitter(0);
252        assert_eq!(policy.delay(Attempt::new(1)), Duration::from_secs(1));
253        assert_eq!(policy.delay(Attempt::new(2)), Duration::from_secs(2));
254        assert_eq!(policy.delay(Attempt::new(3)), Duration::from_secs(4));
255        assert_eq!(policy.delay(Attempt::new(4)), Duration::from_secs(5));
256        assert_eq!(policy.delay(Attempt::new(40)), Duration::from_secs(5));
257    }
258
259    #[test]
260    fn jitter_spreads_and_shortens_delays() {
261        let policy = RetryPolicy::new(5).with_base_delay(Duration::from_secs(10));
262        let delays: Vec<_> = (0..20)
263            .map(|i| policy.jittered(&seed(&i.to_string()).next_attempt()))
264            .collect();
265
266        let full = Duration::from_secs(10);
267        assert!(
268            delays
269                .iter()
270                .all(|delay| *delay <= full && *delay >= full * 8 / 10)
271        );
272        assert!(delays.iter().any(|delay| *delay != delays[0]));
273        assert_eq!(policy.jittered(&seed("0").next_attempt()), delays[0]);
274    }
275
276    #[tokio::test(start_paused = true)]
277    async fn retried_seed_waits_for_its_delay() {
278        let frontier = InMemFrontier::fifo().retry(1);
279        let after = Some(Duration::from_secs(2));
280        frontier
281            .report(Report::Retry {
282                seed: seed("a"),
283                after,
284            })
285            .await
286            .unwrap();
287        assert_eq!(frontier.size_hint(), Some(1));
288
289        match frontier.pop().await.unwrap() {
290            NextSeed::Wait(at) => assert_eq!(at - Instant::now(), Duration::from_secs(2)),
291            next => panic!("expected wait, got {next:?}"),
292        }
293
294        advance(Duration::from_secs(2)).await;
295        let NextSeed::Ready(retried) = frontier.pop().await.unwrap() else {
296            panic!("expected the retried seed");
297        };
298        assert_eq!(retried.attempt(), Attempt::new(1));
299
300        // The single retry is used up.
301        frontier
302            .report(Report::Retry {
303                seed: retried,
304                after,
305            })
306            .await
307            .unwrap();
308        assert_eq!(pop_path(&frontier).await, None);
309        assert!(matches!(frontier.pop().await.unwrap(), NextSeed::Empty));
310    }
311
312    #[tokio::test(start_paused = true)]
313    async fn queued_seeds_are_served_while_retries_wait() {
314        let frontier = InMemFrontier::fifo().retry(3);
315        frontier.push(seed("fresh")).await.unwrap();
316        let retry = Report::Retry {
317            seed: seed("retried"),
318            after: Some(Duration::from_secs(1)),
319        };
320        frontier.report(retry).await.unwrap();
321
322        assert_eq!(pop_path(&frontier).await.as_deref(), Some("fresh"));
323        assert!(matches!(frontier.pop().await.unwrap(), NextSeed::Wait(_)));
324        advance(Duration::from_secs(1)).await;
325        assert_eq!(pop_path(&frontier).await.as_deref(), Some("retried"));
326    }
327}
spire_core/frontier/facade/retry.rs

spire_core/frontier/facade/
retry.rs
