bitwarden_access_connector/executor/
retry.rs

1//! Retry helpers for target-side and server-side rotation steps.
2
3use std::{convert::Infallible, future::Future, time::Duration};
4
5use tokio::time::Instant;
6
7use crate::error::Retry;
8
9/// Configuration for the retry helpers.
10#[derive(Debug, Clone)]
11pub(crate) struct RetryCfg {
12    /// Total tries including the first, so the default of 5 allows at most 4 backoff sleeps. Zero
13    /// counts as one.
14    pub(crate) max_retry_attempts: u32,
15
16    /// The sleep before the n-th retry is `retry_base_delay * 2^(n-1)`, capped at 32x. Defaults to
17    /// one second.
18    pub(crate) retry_base_delay: Duration,
19}
20
21impl Default for RetryCfg {
22    fn default() -> Self {
23        Self {
24            max_retry_attempts: 5,
25            retry_base_delay: Duration::from_secs(1),
26        }
27    }
28}
29
30/// The outcome of a [`with_retries_gated`] call.
31#[derive(Debug)]
32pub(crate) enum GatedOutcome<T, E, A> {
33    Ok(T),
34    /// The gate aborted execution (session lost, execute_by expired, cancelled).
35    Aborted(A),
36    /// A fatal error, or the last transient error once the tries ran out.
37    Failed(E),
38}
39
40/// Retry `op` with exponential backoff. [`Retry::Permanent`] returns at once;
41/// [`Retry::Transient`] retries until the tries run out or `deadline` passes, with sleeps
42/// truncated to it.
43pub(crate) async fn with_retries<F, Fut, T, E>(
44    cfg: &RetryCfg,
45    deadline: Option<Instant>,
46    op: F,
47) -> Result<T, E>
48where
49    F: FnMut() -> Fut,
50    Fut: Future<Output = Result<T, (Retry, E)>>,
51{
52    let outcome = retry_loop(cfg, deadline, || async { Ok::<(), Infallible>(()) }, op).await;
53
54    match outcome {
55        GatedOutcome::Ok(v) => Ok(v),
56        GatedOutcome::Failed(e) => Err(e),
57        GatedOutcome::Aborted(never) => match never {},
58    }
59}
60
61/// Like [`with_retries`], but awaits `gate()` before every try; an abort returns
62/// [`GatedOutcome::Aborted`]. Session loss or lease expiry stops target-side steps between tries,
63/// never mid-call. Sleeps are not capped to a deadline: the gate checks execute_by before the next
64/// try.
65pub(crate) async fn with_retries_gated<G, GFut, F, Fut, T, E, A>(
66    cfg: &RetryCfg,
67    gate: G,
68    op: F,
69) -> GatedOutcome<T, E, A>
70where
71    G: FnMut() -> GFut,
72    GFut: Future<Output = Result<(), A>>,
73    F: FnMut() -> Fut,
74    Fut: Future<Output = Result<T, (Retry, E)>>,
75{
76    retry_loop(cfg, None, gate, op).await
77}
78
79/// The shared retry loop: gate before each try, then `op`, sleeping between transient failures.
80async fn retry_loop<G, GFut, F, Fut, T, E, A>(
81    cfg: &RetryCfg,
82    deadline: Option<Instant>,
83    mut gate: G,
84    mut op: F,
85) -> GatedOutcome<T, E, A>
86where
87    G: FnMut() -> GFut,
88    GFut: Future<Output = Result<(), A>>,
89    F: FnMut() -> Fut,
90    Fut: Future<Output = Result<T, (Retry, E)>>,
91{
92    let max_tries = cfg.max_retry_attempts.max(1);
93    let base = cfg.retry_base_delay;
94    let mut last_err: Option<E> = None;
95
96    for attempt in 0..max_tries {
97        if let Err(abort) = gate().await {
98            return GatedOutcome::Aborted(abort);
99        }
100
101        match op().await {
102            Ok(v) => return GatedOutcome::Ok(v),
103            Err((Retry::Permanent, e)) => return GatedOutcome::Failed(e),
104            Err((Retry::Transient, e)) => {
105                last_err = Some(e);
106
107                if attempt + 1 >= max_tries {
108                    break;
109                }
110
111                let Some(sleep) = cap_to_deadline(exponential_delay(base, attempt), deadline)
112                else {
113                    break;
114                };
115                tokio::time::sleep(sleep).await;
116            }
117        }
118    }
119
120    // At least one transient error was recorded.
121    #[allow(clippy::unwrap_used)]
122    GatedOutcome::Failed(last_err.unwrap())
123}
124
125/// `base * 2^attempt` for a 0-indexed attempt, capped at 32x to avoid overflow.
126fn exponential_delay(base: Duration, attempt: u32) -> Duration {
127    let shift = attempt.min(5);
128    base * (1u32 << shift)
129}
130
131/// Truncate `delay` so it does not push past `deadline`, or `None` once the deadline has passed.
132fn cap_to_deadline(delay: Duration, deadline: Option<Instant>) -> Option<Duration> {
133    match deadline {
134        None => Some(delay),
135        Some(dl) => {
136            let now = Instant::now();
137            if now >= dl {
138                None
139            } else {
140                Some(delay.min(dl - now))
141            }
142        }
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use std::{
149        sync::{Arc, Mutex},
150        time::Duration,
151    };
152
153    use tokio::time::Instant;
154
155    use super::*;
156    use crate::error::Retry;
157
158    #[tokio::test(start_paused = true)]
159    async fn retry_calls_op_max_retry_attempts_times_on_transient() {
160        let calls = Arc::new(Mutex::new(0u32));
161        let cfg = RetryCfg {
162            max_retry_attempts: 5,
163            retry_base_delay: Duration::from_secs(1),
164        };
165
166        let calls_clone = Arc::clone(&calls);
167        let result = with_retries::<_, _, (), String>(&cfg, None, || {
168            let c = Arc::clone(&calls_clone);
169            async move {
170                *c.lock().unwrap() += 1;
171                Err((Retry::Transient, "transient".to_string()))
172            }
173        })
174        .await;
175
176        assert!(result.is_err());
177        assert_eq!(*calls.lock().unwrap(), 5, "should try exactly 5 times");
178    }
179
180    #[tokio::test(start_paused = true)]
181    async fn retry_zero_base_delay_still_retries() {
182        let calls = Arc::new(Mutex::new(0u32));
183        let cfg = RetryCfg {
184            max_retry_attempts: 3,
185            retry_base_delay: Duration::ZERO,
186        };
187
188        let calls_clone = Arc::clone(&calls);
189        let result = with_retries::<_, _, (), String>(&cfg, None, || {
190            let c = Arc::clone(&calls_clone);
191            async move {
192                *c.lock().unwrap() += 1;
193                Err((Retry::Transient, "transient".to_string()))
194            }
195        })
196        .await;
197
198        assert!(result.is_err());
199        assert_eq!(
200            *calls.lock().unwrap(),
201            3,
202            "zero delay should not stop retries"
203        );
204    }
205
206    #[tokio::test(start_paused = true)]
207    async fn retry_fatal_short_circuits_immediately() {
208        let calls = Arc::new(Mutex::new(0u32));
209        let cfg = RetryCfg {
210            max_retry_attempts: 5,
211            retry_base_delay: Duration::from_secs(1),
212        };
213
214        let calls_clone = Arc::clone(&calls);
215        let result = with_retries::<_, _, (), String>(&cfg, None, || {
216            let c = Arc::clone(&calls_clone);
217            async move {
218                *c.lock().unwrap() += 1;
219                Err((Retry::Permanent, "fatal".to_string()))
220            }
221        })
222        .await;
223
224        assert!(result.is_err());
225        assert_eq!(*calls.lock().unwrap(), 1, "fatal should stop after 1 call");
226    }
227
228    #[tokio::test(start_paused = true)]
229    async fn retry_succeeds_on_first_try() {
230        let cfg = RetryCfg::default();
231        let result = with_retries::<_, _, i32, String>(&cfg, None, || async { Ok(42) }).await;
232        assert_eq!(result, Ok(42));
233    }
234
235    #[tokio::test(start_paused = true)]
236    async fn retry_succeeds_on_third_try() {
237        let calls = Arc::new(Mutex::new(0u32));
238        let cfg = RetryCfg {
239            max_retry_attempts: 5,
240            retry_base_delay: Duration::from_secs(1),
241        };
242
243        let calls_clone = Arc::clone(&calls);
244        let result = with_retries::<_, _, i32, String>(&cfg, None, || {
245            let c = Arc::clone(&calls_clone);
246            async move {
247                let mut n = c.lock().unwrap();
248                *n += 1;
249                if *n >= 3 {
250                    Ok(99)
251                } else {
252                    Err((Retry::Transient, "wait".to_string()))
253                }
254            }
255        })
256        .await;
257
258        assert_eq!(result, Ok(99));
259        assert_eq!(*calls.lock().unwrap(), 3);
260    }
261
262    #[tokio::test(start_paused = true)]
263    async fn retry_deadline_in_past_stops_after_first_transient() {
264        let cfg = RetryCfg {
265            max_retry_attempts: 5,
266            retry_base_delay: Duration::from_secs(1),
267        };
268        let past = Instant::now().checked_sub(Duration::from_secs(1));
269        // Skip this test on subtraction underflow.
270        let Some(past_deadline) = past else {
271            return;
272        };
273
274        let calls = Arc::new(Mutex::new(0u32));
275        let calls_clone = Arc::clone(&calls);
276        let result = with_retries::<_, _, (), String>(&cfg, Some(past_deadline), || {
277            let c = Arc::clone(&calls_clone);
278            async move {
279                *c.lock().unwrap() += 1;
280                Err((Retry::Transient, "t".to_string()))
281            }
282        })
283        .await;
284
285        assert!(result.is_err());
286        assert_eq!(
287            *calls.lock().unwrap(),
288            1,
289            "past deadline should stop after first try"
290        );
291    }
292
293    #[tokio::test(start_paused = true)]
294    async fn gated_gate_called_before_each_try() {
295        let gate_calls = Arc::new(Mutex::new(0u32));
296        let op_calls = Arc::new(Mutex::new(0u32));
297        let cfg = RetryCfg {
298            max_retry_attempts: 3,
299            retry_base_delay: Duration::from_millis(10),
300        };
301
302        let gc = Arc::clone(&gate_calls);
303        let oc = Arc::clone(&op_calls);
304        let result = with_retries_gated::<_, _, _, _, (), String, ()>(
305            &cfg,
306            || {
307                let gc = Arc::clone(&gc);
308                async move {
309                    *gc.lock().unwrap() += 1;
310                    Ok(())
311                }
312            },
313            || {
314                let oc = Arc::clone(&oc);
315                async move {
316                    *oc.lock().unwrap() += 1;
317                    Err((Retry::Transient, "t".to_string()))
318                }
319            },
320        )
321        .await;
322
323        assert!(matches!(result, GatedOutcome::Failed(_)));
324        assert_eq!(
325            *gate_calls.lock().unwrap(),
326            3,
327            "gate called before each try"
328        );
329        assert_eq!(*op_calls.lock().unwrap(), 3);
330    }
331
332    #[tokio::test(start_paused = true)]
333    async fn gated_abort_between_tries_stops_before_next_op() {
334        let gate_calls = Arc::new(Mutex::new(0u32));
335        let op_calls = Arc::new(Mutex::new(0u32));
336        let cfg = RetryCfg {
337            max_retry_attempts: 5,
338            retry_base_delay: Duration::from_millis(10),
339        };
340
341        let gc = Arc::clone(&gate_calls);
342        let oc = Arc::clone(&op_calls);
343        let result = with_retries_gated::<_, _, _, _, (), String, &'static str>(
344            &cfg,
345            || {
346                let gc = Arc::clone(&gc);
347                async move {
348                    let mut n = gc.lock().unwrap();
349                    *n += 1;
350                    if *n > 1 { Err("aborted") } else { Ok(()) }
351                }
352            },
353            || {
354                let oc = Arc::clone(&oc);
355                async move {
356                    *oc.lock().unwrap() += 1;
357                    Err((Retry::Transient, "t".to_string()))
358                }
359            },
360        )
361        .await;
362
363        assert!(matches!(result, GatedOutcome::Aborted("aborted")));
364        // Gate: before try 1 (Ok) and before try 2 (Err).
365        assert_eq!(*gate_calls.lock().unwrap(), 2);
366        // Op: try 1 only; aborted before try 2.
367        assert_eq!(*op_calls.lock().unwrap(), 1);
368    }
369
370    #[tokio::test(start_paused = true)]
371    async fn gated_gate_abort_on_first_try_skips_op() {
372        let op_calls = Arc::new(Mutex::new(0u32));
373        let cfg = RetryCfg::default();
374
375        let oc = Arc::clone(&op_calls);
376        let result = with_retries_gated::<_, _, _, _, (), String, &'static str>(
377            &cfg,
378            || async { Err("immediate abort") },
379            || {
380                let oc = Arc::clone(&oc);
381                async move {
382                    *oc.lock().unwrap() += 1;
383                    Ok(())
384                }
385            },
386        )
387        .await;
388
389        assert!(matches!(result, GatedOutcome::Aborted("immediate abort")));
390        assert_eq!(
391            *op_calls.lock().unwrap(),
392            0,
393            "op never called when gate aborts first"
394        );
395    }
396
397    #[tokio::test(start_paused = true)]
398    async fn gated_fatal_short_circuits() {
399        let gate_calls = Arc::new(Mutex::new(0u32));
400        let cfg = RetryCfg {
401            max_retry_attempts: 5,
402            retry_base_delay: Duration::from_millis(10),
403        };
404
405        let gc = Arc::clone(&gate_calls);
406        let result = with_retries_gated::<_, _, _, _, (), String, ()>(
407            &cfg,
408            || {
409                let gc = Arc::clone(&gc);
410                async move {
411                    *gc.lock().unwrap() += 1;
412                    Ok(())
413                }
414            },
415            || async { Err((Retry::Permanent, "fatal".to_string())) },
416        )
417        .await;
418
419        assert!(matches!(result, GatedOutcome::Failed(_)));
420        assert_eq!(
421            *gate_calls.lock().unwrap(),
422            1,
423            "only one gate call for fatal"
424        );
425    }
426
427    #[tokio::test(start_paused = true)]
428    async fn gated_ok_on_first_try() {
429        let cfg = RetryCfg::default();
430        let result = with_retries_gated::<_, _, _, _, i32, String, ()>(
431            &cfg,
432            || async { Ok(()) },
433            || async { Ok(42) },
434        )
435        .await;
436        assert!(matches!(result, GatedOutcome::Ok(42)));
437    }
438}
bitwarden_access_connector/executor/retry.rs

bitwarden_access_connector/executor/
retry.rs
