aster_forge_db/
retry.rs

1//! Retry helpers for transient database operations.
2//!
3//! Three levels exist, from preferred to most specific:
4//!
5//! - Multi-statement transactions: [`crate::transaction::with_transaction_retry`] retries the
6//!   whole transaction callback and classifies commit outcomes. This is the only retry level
7//!   that is safe for transactions.
8//! - Idempotent single statements (reads, upserts, deletes by key): [`with_sea_orm_retry`] and
9//!   [`with_sea_orm_retry_timeout`].
10//! - Crate-internal [`DbError`] workflows such as connection setup: [`with_retry`].
11//!
12//! Every retryability decision derives from [`crate::database_error_kind`] and driver-native
13//! error codes; error message text is never inspected. Non-retryable errors return
14//! immediately so application bugs are not hidden behind sleep loops.
15
16use std::time::Duration;
17use tokio::time::{sleep, timeout};
18
19use crate::{DbError, Result};
20
21/// Retry configuration for async database operations.
22///
23/// One config type serves every retry level; construct it through a profile
24/// ([`RetryConfig::connection`], [`RetryConfig::deadlock`]) and override individual fields
25/// when the product needs a different budget.
26#[derive(Clone, Copy, Debug)]
27pub struct RetryConfig {
28    /// Maximum number of retry attempts after the initial attempt.
29    pub max_retries: u32,
30    /// Base exponential-backoff delay in milliseconds.
31    pub base_delay_ms: u64,
32    /// Maximum backoff delay in milliseconds.
33    pub max_delay_ms: u64,
34}
35
36impl RetryConfig {
37    /// Profile for connection setup and acquisition: the peer may need time to recover,
38    /// so back off slowly (3 retries, 100ms base, 5s cap).
39    #[must_use]
40    pub fn connection() -> Self {
41        Self {
42            max_retries: 3,
43            base_delay_ms: 100,
44            max_delay_ms: 5000,
45        }
46    }
47
48    /// Profile for deadlock/serialization-failure retries: lock-wait windows are short,
49    /// so retry quickly (3 retries, 5ms base, 50ms cap).
50    #[must_use]
51    pub fn deadlock() -> Self {
52        Self {
53            max_retries: 3,
54            base_delay_ms: 5,
55            max_delay_ms: 50,
56        }
57    }
58}
59
60impl Default for RetryConfig {
61    /// Defaults to the [`RetryConfig::connection`] profile.
62    fn default() -> Self {
63        Self::connection()
64    }
65}
66
67/// Returns whether a `SeaORM` error represents a transient database failure.
68///
69/// Connection acquisition and connection failures are always retryable: they happen before
70/// the statement ran, so retrying cannot duplicate work. Query and execution failures are
71/// retried only when [`crate::database_error_kind`] proves a transient locking conflict
72/// (deadlock, serialization failure, lock timeout) from driver-native error codes. Error
73/// message text is never inspected.
74pub fn is_retryable_sea_orm_error(error: &sea_orm::DbErr) -> bool {
75    use sea_orm::DbErr;
76
77    match error {
78        DbErr::ConnectionAcquire(_) | DbErr::Conn(_) => true,
79        _ => crate::database_error_kind(error)
80            .is_some_and(crate::DatabaseErrorKind::is_transient_locking),
81    }
82}
83
84/// Executes a `SeaORM` operation with shared transient-error classification and backoff.
85///
86/// This is a statement-level escape hatch for idempotent single statements only (reads,
87/// upserts, deletes by key). **Never use it around a multi-statement transaction**: a
88/// deadlock rolls the whole transaction back, and re-running individual statements
89/// afterwards executes them in autocommit mode, producing partial writes. Transactions
90/// belong in [`crate::transaction::with_transaction_retry`], which retries the entire
91/// callback and classifies commit outcomes.
92///
93/// # Errors
94///
95/// Returns an error when retry orchestration or the database operation fails.
96pub async fn with_sea_orm_retry<F, Fut, T>(
97    operation_name: &str,
98    config: RetryConfig,
99    mut operation: F,
100) -> std::result::Result<T, sea_orm::DbErr>
101where
102    F: FnMut() -> Fut,
103    Fut: std::future::Future<Output = std::result::Result<T, sea_orm::DbErr>>,
104{
105    let mut attempt = 0_u32;
106    loop {
107        match operation().await {
108            Ok(value) => return Ok(value),
109            Err(error) if attempt < config.max_retries && is_retryable_sea_orm_error(&error) => {
110                let delay = calculate_delay(&config, attempt);
111                tracing::warn!(
112                    operation = operation_name,
113                    attempt = attempt + 1,
114                    max_attempts = config.max_retries + 1,
115                    delay_ms = duration_millis_u64(delay),
116                    error = %error,
117                    "retrying SeaORM operation"
118                );
119                sleep(delay).await;
120                attempt += 1;
121            }
122            Err(error) => return Err(error),
123        }
124    }
125}
126
127/// Executes a `SeaORM` operation with retry and a timeout applied to every attempt.
128///
129/// Same classification and boundary rules as [`with_sea_orm_retry`]; a timed-out attempt
130/// is retried regardless of error classification because the operation produced no outcome.
131/// Only wrap work that stays safe when a timed-out attempt keeps running in the background.
132///
133/// # Errors
134///
135/// Returns an error when retry orchestration or the database operation fails.
136pub async fn with_sea_orm_retry_timeout<F, Fut, T>(
137    operation_name: &str,
138    config: RetryConfig,
139    attempt_timeout: Duration,
140    mut operation: F,
141) -> std::result::Result<T, sea_orm::DbErr>
142where
143    F: FnMut() -> Fut,
144    Fut: std::future::Future<Output = std::result::Result<T, sea_orm::DbErr>>,
145{
146    let mut attempt = 0_u32;
147    loop {
148        match timeout(attempt_timeout, operation()).await {
149            Ok(Ok(value)) => return Ok(value),
150            Ok(Err(error))
151                if attempt < config.max_retries && is_retryable_sea_orm_error(&error) =>
152            {
153                let delay = calculate_delay(&config, attempt);
154                tracing::warn!(
155                    operation = operation_name,
156                    attempt = attempt + 1,
157                    max_attempts = config.max_retries + 1,
158                    delay_ms = duration_millis_u64(delay),
159                    error = %error,
160                    "retrying SeaORM operation"
161                );
162                sleep(delay).await;
163                attempt += 1;
164            }
165            Ok(Err(error)) => return Err(error),
166            Err(_) if attempt < config.max_retries => {
167                let delay = calculate_delay(&config, attempt);
168                tracing::warn!(
169                    operation = operation_name,
170                    attempt = attempt + 1,
171                    max_attempts = config.max_retries + 1,
172                    timeout_ms = duration_millis_u64(attempt_timeout),
173                    delay_ms = duration_millis_u64(delay),
174                    "SeaORM operation attempt timed out; retrying"
175                );
176                sleep(delay).await;
177                attempt += 1;
178            }
179            Err(_) => {
180                return Err(sea_orm::DbErr::Custom(format!(
181                    "operation '{operation_name}' timed out after {}ms",
182                    duration_millis_u64(attempt_timeout)
183                )));
184            }
185        }
186    }
187}
188
189/// Execute an async operation with exponential backoff retry
190///
191/// # Errors
192///
193/// Returns an error when retry orchestration or the database operation fails.
194pub async fn with_retry<F, Fut, T>(config: &RetryConfig, operation: F) -> Result<T>
195where
196    F: Fn() -> Fut,
197    Fut: std::future::Future<Output = Result<T>>,
198{
199    let mut last_err = None;
200    for attempt in 0..=config.max_retries {
201        match operation().await {
202            Ok(val) => return Ok(val),
203            Err(e) => {
204                if attempt == config.max_retries || !is_retryable(&e) {
205                    return Err(e);
206                }
207                let delay = calculate_delay(config, attempt);
208                tracing::warn!(
209                    attempt = attempt + 1,
210                    max = config.max_retries,
211                    delay_ms = duration_millis_u64(delay),
212                    error = %e,
213                    "retrying operation"
214                );
215                last_err = Some(e);
216                sleep(delay).await;
217            }
218        }
219    }
220    Err(last_err.unwrap_or(DbError::RetryExhausted))
221}
222
223fn is_retryable(err: &DbError) -> bool {
224    err.is_retryable()
225}
226
227fn calculate_delay(config: &RetryConfig, attempt: u32) -> Duration {
228    use aster_forge_utils::backoff::{cap_delay, exponential_delay, randomized_jitter};
229
230    let raw = exponential_delay(Duration::from_millis(config.base_delay_ms), attempt);
231    cap_delay(
232        randomized_jitter(raw, 50, 150),
233        Duration::from_millis(config.max_delay_ms),
234    )
235}
236
237fn duration_millis_u64(duration: Duration) -> u64 {
238    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use std::sync::{
245        Arc,
246        atomic::{AtomicUsize, Ordering},
247    };
248
249    fn zero_delay_config(max_retries: u32) -> RetryConfig {
250        RetryConfig {
251            max_retries,
252            base_delay_ms: 0,
253            max_delay_ms: 0,
254        }
255    }
256
257    #[test]
258    fn retryable_error_classification_only_allows_database_errors() {
259        assert!(is_retryable(&DbError::database_connection(
260            "connection lost"
261        )));
262        assert!(is_retryable(&DbError::database_operation_classified(
263            "deadlock",
264            crate::DatabaseErrorKind::Deadlock
265        )));
266        // Unclassified operation errors carry no driver evidence of retry safety.
267        assert!(!is_retryable(&DbError::database_operation("deadlock")));
268        assert!(!is_retryable(&DbError::non_retryable("invalid input")));
269        assert!(!is_retryable(&DbError::non_retryable("forbidden")));
270    }
271
272    #[test]
273    fn sea_orm_retryability_allows_connection_failures_before_any_statement() {
274        let conn_error = sea_orm::DbErr::Conn(sea_orm::error::RuntimeErr::Internal(
275            "connection reset".to_string(),
276        ));
277
278        assert!(is_retryable_sea_orm_error(&conn_error));
279    }
280
281    #[test]
282    fn sea_orm_retryability_rejects_non_driver_errors_without_reading_messages() {
283        // Message text must not influence classification: these mention deadlock but are
284        // not driver database errors, so they are not retryable.
285        assert!(!is_retryable_sea_orm_error(&sea_orm::DbErr::Custom(
286            "deadlock detected".to_string()
287        )));
288        assert!(!is_retryable_sea_orm_error(
289            &sea_orm::DbErr::RecordNotFound("deadlock".to_string())
290        ));
291    }
292
293    #[tokio::test]
294    async fn sea_orm_retryability_classifies_real_sqlite_busy_as_retryable() {
295        use aster_forge_test::temp::SqliteTestDatabase;
296        use sea_orm::{ConnectOptions, ConnectionTrait, SqlxSqliteConnector};
297
298        let database = SqliteTestDatabase::new("retry-busy");
299        let locker = SqlxSqliteConnector::connect(ConnectOptions::new(database.url()))
300            .await
301            .unwrap();
302        let contender = SqlxSqliteConnector::connect(ConnectOptions::new(database.url()))
303            .await
304            .unwrap();
305        contender
306            .execute_unprepared("PRAGMA busy_timeout=0;")
307            .await
308            .unwrap();
309        locker
310            .execute_unprepared("CREATE TABLE items (id INTEGER PRIMARY KEY);")
311            .await
312            .unwrap();
313        locker.execute_unprepared("BEGIN IMMEDIATE;").await.unwrap();
314        locker
315            .execute_unprepared("INSERT INTO items (id) VALUES (1);")
316            .await
317            .unwrap();
318
319        let error = contender
320            .execute_unprepared("INSERT INTO items (id) VALUES (2);")
321            .await
322            .unwrap_err();
323
324        locker.execute_unprepared("ROLLBACK;").await.unwrap();
325        contender.close().await.unwrap();
326        locker.close().await.unwrap();
327
328        assert!(
329            is_retryable_sea_orm_error(&error),
330            "SQLITE_BUSY from a locked database should be retryable, got: {error}"
331        );
332        assert_eq!(
333            crate::database_error_kind(&error),
334            Some(crate::DatabaseErrorKind::LockTimeout)
335        );
336    }
337
338    #[tokio::test]
339    async fn sea_orm_retryability_rejects_sqlite_unique_violation() {
340        use sea_orm::{ConnectionTrait, Database};
341
342        let db = Database::connect("sqlite::memory:").await.unwrap();
343        db.execute_unprepared("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT UNIQUE);")
344            .await
345            .unwrap();
346        db.execute_unprepared("INSERT INTO items (name) VALUES ('a');")
347            .await
348            .unwrap();
349
350        let error = db
351            .execute_unprepared("INSERT INTO items (name) VALUES ('a');")
352            .await
353            .unwrap_err();
354
355        assert!(!is_retryable_sea_orm_error(&error));
356        assert_eq!(
357            crate::database_error_kind(&error),
358            Some(crate::DatabaseErrorKind::UniqueConstraint)
359        );
360    }
361
362    #[test]
363    fn calculate_delay_applies_jitter_and_hard_caps_max_delay() {
364        let config = RetryConfig {
365            max_retries: 3,
366            base_delay_ms: 100,
367            max_delay_ms: 250,
368        };
369
370        let expected_bounds = [(0, 50, 150), (1, 100, 250), (2, 200, 250), (8, 250, 250)];
371
372        for (attempt, min_ms, max_ms) in expected_bounds {
373            for _ in 0..64 {
374                let delay_ms = duration_millis_u64(calculate_delay(&config, attempt));
375                assert!(
376                    (min_ms..=max_ms).contains(&delay_ms),
377                    "attempt {attempt} produced {delay_ms}ms outside [{min_ms}, {max_ms}]"
378                );
379            }
380        }
381    }
382
383    #[test]
384    fn calculate_delay_handles_zero_and_initial_above_max_boundaries() {
385        let zero_base = RetryConfig {
386            max_retries: 1,
387            base_delay_ms: 0,
388            max_delay_ms: 100,
389        };
390        let zero_max = RetryConfig {
391            max_retries: 1,
392            base_delay_ms: 100,
393            max_delay_ms: 0,
394        };
395        let initial_above_max = RetryConfig {
396            max_retries: 1,
397            base_delay_ms: 1_000,
398            max_delay_ms: 250,
399        };
400
401        for attempt in [0, 1, u32::MAX] {
402            assert_eq!(calculate_delay(&zero_base, attempt), Duration::ZERO);
403            assert_eq!(calculate_delay(&zero_max, attempt), Duration::ZERO);
404            assert_eq!(
405                calculate_delay(&initial_above_max, attempt),
406                Duration::from_millis(250)
407            );
408        }
409    }
410
411    #[tokio::test]
412    async fn with_retry_retries_retryable_errors_until_success() {
413        let attempts = Arc::new(AtomicUsize::new(0));
414        let result = {
415            let attempts = Arc::clone(&attempts);
416            with_retry(&zero_delay_config(3), move || {
417                let attempts = Arc::clone(&attempts);
418                async move {
419                    let attempt = attempts.fetch_add(1, Ordering::SeqCst);
420                    if attempt < 2 {
421                        Err(DbError::database_operation_classified(
422                            "deadlock",
423                            crate::DatabaseErrorKind::Deadlock,
424                        ))
425                    } else {
426                        Ok("ok")
427                    }
428                }
429            })
430            .await
431        };
432
433        assert_eq!(result.unwrap(), "ok");
434        assert_eq!(attempts.load(Ordering::SeqCst), 3);
435    }
436
437    #[tokio::test]
438    async fn with_retry_stops_immediately_for_non_retryable_errors() {
439        let attempts = Arc::new(AtomicUsize::new(0));
440        let result = {
441            let attempts = Arc::clone(&attempts);
442            with_retry(&zero_delay_config(3), move || {
443                let attempts = Arc::clone(&attempts);
444                async move {
445                    attempts.fetch_add(1, Ordering::SeqCst);
446                    Err::<(), _>(DbError::non_retryable("bad request"))
447                }
448            })
449            .await
450        };
451
452        assert!(matches!(result.unwrap_err(), DbError::NonRetryable(_)));
453        assert_eq!(attempts.load(Ordering::SeqCst), 1);
454    }
455
456    #[tokio::test]
457    async fn with_retry_stops_after_exhausting_retry_budget() {
458        let attempts = Arc::new(AtomicUsize::new(0));
459        let result = {
460            let attempts = Arc::clone(&attempts);
461            with_retry(&zero_delay_config(2), move || {
462                let attempts = Arc::clone(&attempts);
463                async move {
464                    attempts.fetch_add(1, Ordering::SeqCst);
465                    Err::<(), _>(DbError::database_connection("still failing"))
466                }
467            })
468            .await
469        };
470
471        assert!(matches!(
472            result.unwrap_err(),
473            DbError::DatabaseConnection(_)
474        ));
475        assert_eq!(attempts.load(Ordering::SeqCst), 3);
476    }
477
478    #[tokio::test]
479    async fn with_retry_zero_budget_runs_exactly_one_attempt() {
480        let attempts = Arc::new(AtomicUsize::new(0));
481        let result = {
482            let attempts = Arc::clone(&attempts);
483            with_retry(&zero_delay_config(0), move || {
484                let attempts = Arc::clone(&attempts);
485                async move {
486                    attempts.fetch_add(1, Ordering::SeqCst);
487                    Err::<(), _>(DbError::database_connection("still failing"))
488                }
489            })
490            .await
491        };
492
493        assert!(matches!(
494            result.unwrap_err(),
495            DbError::DatabaseConnection(_)
496        ));
497        assert_eq!(attempts.load(Ordering::SeqCst), 1);
498    }
499}
aster_forge_db/retry.rs

aster_forge_db/
retry.rs
