backon/
blocking_retry_with_context.rs

1use core::time::Duration;
2
3use crate::backoff::BackoffBuilder;
4use crate::blocking_sleep::MaybeBlockingSleeper;
5use crate::Backoff;
6use crate::BlockingSleeper;
7use crate::DefaultBlockingSleeper;
8
9/// BlockingRetryableWithContext adds retry support for blocking functions.
10pub trait BlockingRetryableWithContext<
11    B: BackoffBuilder,
12    T,
13    E,
14    Ctx,
15    F: FnMut(Ctx) -> (Ctx, Result<T, E>),
16>
17{
18    /// Generate a new retry
19    fn retry(self, builder: B) -> BlockingRetryWithContext<B::Backoff, T, E, Ctx, F>;
20}
21
22impl<B, T, E, Ctx, F> BlockingRetryableWithContext<B, T, E, Ctx, F> for F
23where
24    B: BackoffBuilder,
25    F: FnMut(Ctx) -> (Ctx, Result<T, E>),
26{
27    fn retry(self, builder: B) -> BlockingRetryWithContext<B::Backoff, T, E, Ctx, F> {
28        BlockingRetryWithContext::new(self, builder.build())
29    }
30}
31
32/// Retry structure generated by [`BlockingRetryableWithContext`].
33pub struct BlockingRetryWithContext<
34    B: Backoff,
35    T,
36    E,
37    Ctx,
38    F: FnMut(Ctx) -> (Ctx, Result<T, E>),
39    SF: MaybeBlockingSleeper = DefaultBlockingSleeper,
40    RF = fn(&E) -> bool,
41    NF = fn(&E, Duration),
42> {
43    backoff: B,
44    retryable: RF,
45    notify: NF,
46    f: F,
47    sleep_fn: SF,
48    ctx: Option<Ctx>,
49}
50
51impl<B, T, E, Ctx, F> BlockingRetryWithContext<B, T, E, Ctx, F>
52where
53    B: Backoff,
54    F: FnMut(Ctx) -> (Ctx, Result<T, E>),
55{
56    /// Create a new retry.
57    fn new(f: F, backoff: B) -> Self {
58        BlockingRetryWithContext {
59            backoff,
60            retryable: |_: &E| true,
61            notify: |_: &E, _: Duration| {},
62            sleep_fn: DefaultBlockingSleeper::default(),
63            f,
64            ctx: None,
65        }
66    }
67}
68
69impl<B, T, E, Ctx, F, SF, RF, NF> BlockingRetryWithContext<B, T, E, Ctx, F, SF, RF, NF>
70where
71    B: Backoff,
72    F: FnMut(Ctx) -> (Ctx, Result<T, E>),
73    SF: MaybeBlockingSleeper,
74    RF: FnMut(&E) -> bool,
75    NF: FnMut(&E, Duration),
76{
77    /// Set the context for retrying.
78    ///
79    /// Context is used to capture ownership manually to prevent lifetime issues.
80    pub fn context(self, context: Ctx) -> BlockingRetryWithContext<B, T, E, Ctx, F, SF, RF, NF> {
81        BlockingRetryWithContext {
82            backoff: self.backoff,
83            retryable: self.retryable,
84            notify: self.notify,
85            f: self.f,
86            sleep_fn: self.sleep_fn,
87            ctx: Some(context),
88        }
89    }
90
91    /// Set the sleeper for retrying.
92    ///
93    /// The sleeper should implement the [`BlockingSleeper`] trait. The simplest way is to use a closure like  `Fn(Duration)`.
94    ///
95    /// If not specified, we use the [`DefaultBlockingSleeper`].
96    pub fn sleep<SN: BlockingSleeper>(
97        self,
98        sleep_fn: SN,
99    ) -> BlockingRetryWithContext<B, T, E, Ctx, F, SN, RF, NF> {
100        BlockingRetryWithContext {
101            backoff: self.backoff,
102            retryable: self.retryable,
103            notify: self.notify,
104            f: self.f,
105            sleep_fn,
106            ctx: self.ctx,
107        }
108    }
109
110    /// Set the conditions for retrying.
111    ///
112    /// If not specified, all errors are considered retryable.
113    pub fn when<RN: FnMut(&E) -> bool>(
114        self,
115        retryable: RN,
116    ) -> BlockingRetryWithContext<B, T, E, Ctx, F, SF, RN, NF> {
117        BlockingRetryWithContext {
118            backoff: self.backoff,
119            retryable,
120            notify: self.notify,
121            f: self.f,
122            sleep_fn: self.sleep_fn,
123            ctx: self.ctx,
124        }
125    }
126
127    /// Set to notify for all retry attempts.
128    ///
129    /// When a retry happens, the input function will be invoked with the error and the sleep duration before pausing.
130    ///
131    /// If not specified, this operation does nothing.
132    pub fn notify<NN: FnMut(&E, Duration)>(
133        self,
134        notify: NN,
135    ) -> BlockingRetryWithContext<B, T, E, Ctx, F, SF, RF, NN> {
136        BlockingRetryWithContext {
137            backoff: self.backoff,
138            retryable: self.retryable,
139            notify,
140            f: self.f,
141            sleep_fn: self.sleep_fn,
142            ctx: self.ctx,
143        }
144    }
145}
146
147impl<B, T, E, Ctx, F, SF, RF, NF> BlockingRetryWithContext<B, T, E, Ctx, F, SF, RF, NF>
148where
149    B: Backoff,
150    F: FnMut(Ctx) -> (Ctx, Result<T, E>),
151    SF: BlockingSleeper,
152    RF: FnMut(&E) -> bool,
153    NF: FnMut(&E, Duration),
154{
155    /// Call the retried function.
156    ///
157    /// TODO: implement [`FnOnce`] after it stable.
158    pub fn call(mut self) -> (Ctx, Result<T, E>) {
159        let mut ctx = self.ctx.take().expect("context must be valid");
160        loop {
161            let (xctx, result) = (self.f)(ctx);
162            // return ctx ownership back
163            ctx = xctx;
164
165            match result {
166                Ok(v) => return (ctx, Ok(v)),
167                Err(err) => {
168                    if !(self.retryable)(&err) {
169                        return (ctx, Err(err));
170                    }
171
172                    match self.backoff.next() {
173                        None => return (ctx, Err(err)),
174                        Some(dur) => {
175                            (self.notify)(&err, dur);
176                            self.sleep_fn.sleep(dur);
177                        }
178                    }
179                }
180            }
181        }
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    extern crate alloc;
188
189    use alloc::string::ToString;
190    use core::time::Duration;
191
192    use anyhow::anyhow;
193    use anyhow::Result;
194    use spin::Mutex;
195
196    use super::*;
197    use crate::ExponentialBuilder;
198
199    struct Test;
200
201    impl Test {
202        fn hello(&mut self) -> Result<usize> {
203            Err(anyhow!("not retryable"))
204        }
205    }
206
207    #[test]
208    fn test_retry_with_not_retryable_error() -> Result<()> {
209        let error_times = Mutex::new(0);
210
211        let test = Test;
212
213        let backoff = ExponentialBuilder::default().with_min_delay(Duration::from_millis(1));
214
215        let (_, result) = {
216            |mut v: Test| {
217                let mut x = error_times.lock();
218                *x += 1;
219
220                let res = v.hello();
221                (v, res)
222            }
223        }
224        .retry(backoff)
225        .context(test)
226        // Only retry If error message is `retryable`
227        .when(|e| e.to_string() == "retryable")
228        .call();
229
230        assert!(result.is_err());
231        assert_eq!("not retryable", result.unwrap_err().to_string());
232        // `f` always returns error "not retryable", so it should be executed
233        // only once.
234        assert_eq!(*error_times.lock(), 1);
235        Ok(())
236    }
237}
backon/blocking_retry_with_context.rs

backon/
blocking_retry_with_context.rs
