rskit_messaging/middleware/
retry.rs

1//! Retry middleware backed by [`rskit_resilience::RetryPolicy`].
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use rskit_errors::AppResult;
7use rskit_resilience::{ConstantBackoff, LinearBackoff, RetryError, RetryPolicy};
8
9use crate::handler::{HandlerMiddleware, MessageHandler};
10use crate::message::Message;
11
12/// Canonical retry configuration for messaging middleware.
13///
14/// Messaging retries historically retried every handler failure by default.
15/// `RetryPolicy` defaults to retrying only errors marked retryable,
16/// so this wrapper encodes the messaging-specific always-retry default explicitly.
17#[derive(Debug, Clone)]
18pub struct RetryConfig {
19    policy: RetryPolicy,
20}
21
22impl Default for RetryConfig {
23    fn default() -> Self {
24        Self {
25            policy: RetryPolicy::new().with_retry_if(|_| true),
26        }
27    }
28}
29
30impl RetryConfig {
31    /// Create a retry config with messaging defaults.
32    #[must_use]
33    pub fn new() -> Self {
34        Self::default()
35    }
36
37    /// Wrap an explicit resilience retry policy.
38    #[must_use]
39    pub const fn from_policy(policy: RetryPolicy) -> Self {
40        Self { policy }
41    }
42
43    /// Set the maximum number of attempts (including the first call).
44    #[must_use]
45    pub fn with_max_attempts(mut self, n: usize) -> Self {
46        self.policy = self.policy.with_max_attempts(n);
47        self
48    }
49
50    /// Set the initial backoff delay before the first retry.
51    #[must_use]
52    pub fn with_initial_backoff(mut self, d: std::time::Duration) -> Self {
53        self.policy = self.policy.with_initial_backoff(d);
54        self
55    }
56
57    /// Set the upper bound on any single backoff delay.
58    #[must_use]
59    pub fn with_max_backoff(mut self, d: std::time::Duration) -> Self {
60        self.policy = self.policy.with_max_backoff(d);
61        self
62    }
63
64    /// Set the exponential backoff multiplication factor.
65    #[must_use]
66    pub fn with_backoff_factor(mut self, factor: f64) -> Self {
67        self.policy = self.policy.with_backoff_factor(factor);
68        self
69    }
70
71    /// Use a fixed delay for every retry attempt.
72    #[must_use]
73    pub fn with_constant_backoff(mut self, backoff: ConstantBackoff) -> Self {
74        self.policy = self.policy.with_constant_backoff(backoff);
75        self
76    }
77
78    /// Use a linearly increasing retry delay.
79    #[must_use]
80    pub fn with_linear_backoff(mut self, backoff: LinearBackoff) -> Self {
81        self.policy = self.policy.with_linear_backoff(backoff);
82        self
83    }
84
85    /// Enable or disable retry jitter.
86    #[must_use]
87    pub fn with_jitter(mut self, enabled: bool) -> Self {
88        self.policy = self.policy.with_jitter(enabled);
89        self
90    }
91
92    /// Override the predicate used to decide whether an error is retried.
93    #[must_use]
94    pub fn with_retry_if(
95        mut self,
96        f: impl Fn(&rskit_errors::AppError) -> bool + Send + Sync + 'static,
97    ) -> Self {
98        self.policy = self.policy.with_retry_if(f);
99        self
100    }
101
102    /// Register a callback called after each failed attempt before the next backoff sleep.
103    #[must_use]
104    pub fn with_on_retry(
105        mut self,
106        f: impl Fn(u32, &rskit_errors::AppError) + Send + Sync + 'static,
107    ) -> Self {
108        self.policy = self.policy.with_on_retry(f);
109        self
110    }
111
112    /// Borrow the underlying resilience policy.
113    #[must_use]
114    pub const fn policy(&self) -> &RetryPolicy {
115        &self.policy
116    }
117
118    /// Convert into the underlying resilience policy.
119    #[must_use]
120    pub fn into_policy(self) -> RetryPolicy {
121        self.policy
122    }
123
124    async fn execute<F, Fut, T>(&self, f: F) -> Result<T, RetryError>
125    where
126        F: FnMut() -> Fut,
127        Fut: std::future::Future<Output = AppResult<T>>,
128    {
129        self.policy.execute(f).await
130    }
131}
132
133impl From<RetryPolicy> for RetryConfig {
134    fn from(policy: RetryPolicy) -> Self {
135        Self::from_policy(policy)
136    }
137}
138
139/// Create a retry middleware with the given canonical retry policy.
140pub fn retry<T: Send + Sync + Clone + 'static>(config: RetryConfig) -> impl HandlerMiddleware<T> {
141    RetryMiddleware { config }
142}
143
144struct RetryMiddleware {
145    config: RetryConfig,
146}
147
148impl<T: Send + Sync + Clone + 'static> HandlerMiddleware<T> for RetryMiddleware {
149    fn wrap(&self, next: Arc<dyn MessageHandler<T>>) -> Arc<dyn MessageHandler<T>> {
150        Arc::new(RetryHandler {
151            config: self.config.clone(),
152            next,
153        })
154    }
155}
156
157struct RetryHandler<T: Send + Sync + 'static> {
158    config: RetryConfig,
159    next: Arc<dyn MessageHandler<T>>,
160}
161
162#[async_trait]
163impl<T: Send + Sync + Clone + 'static> MessageHandler<T> for RetryHandler<T> {
164    async fn handle(&self, msg: Message<T>) -> AppResult<()> {
165        let next = self.next.clone();
166        self.config
167            .execute(|| {
168                let next = next.clone();
169                let msg = msg.clone();
170                async move { next.handle(msg).await }
171            })
172            .await
173            .map_err(|err| err.last_error)
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use std::sync::atomic::{AtomicU32, Ordering};
180    use std::time::Duration;
181
182    use super::*;
183    use crate::handler::{FnHandler, chain_handlers};
184    use rskit_errors::{AppError, ErrorCode};
185
186    fn test_policy() -> RetryConfig {
187        RetryConfig::new()
188            .with_max_attempts(3)
189            .with_initial_backoff(Duration::from_millis(1))
190            .with_jitter(false)
191    }
192
193    #[tokio::test]
194    async fn retry_success_on_first_attempt() {
195        let counter = Arc::new(AtomicU32::new(0));
196        let c = counter.clone();
197        let base: Arc<dyn MessageHandler<String>> =
198            Arc::new(FnHandler::new(move |_msg: Message<String>| {
199                let c = c.clone();
200                async move {
201                    c.fetch_add(1, Ordering::SeqCst);
202                    Ok(())
203                }
204            }));
205
206        let handler = chain_handlers(
207            base,
208            &[Arc::new(RetryMiddleware {
209                config: test_policy(),
210            })],
211        );
212
213        handler
214            .handle(Message::new("t", "data".to_string()))
215            .await
216            .unwrap();
217        assert_eq!(counter.load(Ordering::SeqCst), 1);
218    }
219
220    #[tokio::test]
221    async fn retry_succeeds_on_second_attempt() {
222        let counter = Arc::new(AtomicU32::new(0));
223        let c = counter.clone();
224        let base: Arc<dyn MessageHandler<String>> =
225            Arc::new(FnHandler::new(move |_msg: Message<String>| {
226                let c = c.clone();
227                async move {
228                    let n = c.fetch_add(1, Ordering::SeqCst);
229                    if n == 0 {
230                        Err(AppError::new(ErrorCode::Internal, "transient"))
231                    } else {
232                        Ok(())
233                    }
234                }
235            }));
236
237        let handler = chain_handlers(
238            base,
239            &[Arc::new(RetryMiddleware {
240                config: test_policy(),
241            })],
242        );
243
244        handler
245            .handle(Message::new("t", "data".to_string()))
246            .await
247            .unwrap();
248        assert_eq!(counter.load(Ordering::SeqCst), 2);
249    }
250
251    #[tokio::test]
252    async fn retry_exhausts_all_attempts() {
253        let counter = Arc::new(AtomicU32::new(0));
254        let c = counter.clone();
255        let base: Arc<dyn MessageHandler<String>> =
256            Arc::new(FnHandler::new(move |_msg: Message<String>| {
257                let c = c.clone();
258                async move {
259                    c.fetch_add(1, Ordering::SeqCst);
260                    Err(AppError::new(ErrorCode::Internal, "always fails"))
261                }
262            }));
263
264        let handler = chain_handlers(
265            base,
266            &[Arc::new(RetryMiddleware {
267                config: test_policy(),
268            })],
269        );
270
271        let result = handler.handle(Message::new("t", "data".to_string())).await;
272        assert!(result.is_err());
273        assert_eq!(counter.load(Ordering::SeqCst), 3);
274    }
275    #[tokio::test]
276    async fn default_retry_config_retries_non_retryable_handler_errors() {
277        let counter = Arc::new(AtomicU32::new(0));
278        let c = counter.clone();
279        let base: Arc<dyn MessageHandler<String>> =
280            Arc::new(FnHandler::new(move |_msg: Message<String>| {
281                let c = c.clone();
282                async move {
283                    c.fetch_add(1, Ordering::SeqCst);
284                    Err(AppError::new(
285                        ErrorCode::InvalidInput,
286                        "invalid but retried",
287                    ))
288                }
289            }));
290
291        let handler = chain_handlers(
292            base,
293            &[Arc::new(RetryMiddleware {
294                config: RetryConfig::new()
295                    .with_max_attempts(2)
296                    .with_initial_backoff(Duration::from_millis(1))
297                    .with_jitter(false),
298            })],
299        );
300
301        let result = handler.handle(Message::new("t", "data".to_string())).await;
302
303        assert!(result.is_err());
304        assert_eq!(counter.load(Ordering::SeqCst), 2);
305    }
306
307    #[tokio::test]
308    async fn public_retry_constructor_and_config_builders_wrap_handler() {
309        let callback_count = Arc::new(AtomicU32::new(0));
310        let callback_seen = callback_count.clone();
311        let config = RetryConfig::from_policy(RetryPolicy::new())
312            .with_max_attempts(1)
313            .with_max_backoff(Duration::from_millis(5))
314            .with_backoff_factor(1.0)
315            .with_constant_backoff(ConstantBackoff::new(Duration::from_millis(1)))
316            .with_linear_backoff(LinearBackoff::new(
317                Duration::from_millis(1),
318                Duration::from_millis(1),
319                Duration::from_millis(5),
320            ))
321            .with_retry_if(|_| false)
322            .with_on_retry(move |_, _| {
323                callback_seen.fetch_add(1, Ordering::SeqCst);
324            });
325        let _policy = config.clone().into_policy();
326        let _borrowed = config.policy();
327
328        let base: Arc<dyn MessageHandler<String>> =
329            Arc::new(FnHandler::new(|_msg: Message<String>| async { Ok(()) }));
330        let middleware: Arc<dyn HandlerMiddleware<String>> = Arc::new(retry(config));
331        let handler = chain_handlers(base, &[middleware]);
332
333        handler
334            .handle(Message::new("t", "ok".to_string()))
335            .await
336            .unwrap();
337        assert_eq!(callback_count.load(Ordering::SeqCst), 0);
338    }
339}
rskit_messaging/middleware/retry.rs

rskit_messaging/middleware/
retry.rs
