Skip to main content

aster_forge_db/
transaction.rs

1//! Transaction helpers with consistent error mapping and rollback tracing.
2//!
3//! Service code can either manually begin/commit/rollback with uniform error conversion or run a
4//! callback inside `with_transaction`. The rollback guard logs dropped transactions, which helps
5//! catch early-return paths that rely on rollback-on-drop instead of making rollback explicit.
6//!
7//! ## Pattern
8//!
9//! Standard service-layer transaction usage:
10//! ```ignore
11//! transaction::with_transaction(db, async |txn| {
12//!     repo::operation(txn, ...).await?;
13//!     repo::another_operation(txn, ...).await?;
14//!     Ok(())
15//! })
16//! .await?;
17//! ```
18
19use crate::{DbError, Result, database_error_kind};
20use sea_orm::TransactionSession;
21use std::panic::Location;
22use std::{future::Future, ops::AsyncFnOnce, pin::Pin, time::Duration};
23
24struct RollbackGuard {
25    file: &'static str,
26    line: u32,
27    armed: bool,
28}
29
30impl RollbackGuard {
31    fn new(location: &'static Location<'static>) -> Self {
32        Self {
33            file: location.file(),
34            line: location.line(),
35            armed: true,
36        }
37    }
38
39    fn disarm(&mut self) {
40        self.armed = false;
41    }
42}
43
44impl Drop for RollbackGuard {
45    fn drop(&mut self) {
46        if self.armed {
47            tracing::warn!(
48                file = self.file,
49                line = self.line,
50                "transaction dropped before explicit commit/rollback; relying on rollback-on-drop"
51            );
52        }
53    }
54}
55
56/// Begins and returns a transaction so the caller can commit or roll it back.
57///
58/// This centralizes `begin` error mapping.
59pub async fn begin<C: sea_orm::TransactionTrait>(db: &C) -> Result<C::Transaction> {
60    db.begin()
61        .await
62        .map_err(|error| database_operation_with_context(error, "begin transaction"))
63}
64
65/// Commits a transaction and maps errors consistently.
66pub async fn commit<T: sea_orm::TransactionSession>(txn: T) -> Result<()> {
67    txn.commit()
68        .await
69        .map_err(|error| database_operation_with_context(error, "commit transaction"))
70}
71
72/// Rolls back a transaction and maps errors consistently.
73pub async fn rollback<T: sea_orm::TransactionSession>(txn: T) -> Result<()> {
74    txn.rollback()
75        .await
76        .map_err(|error| database_operation_with_context(error, "rollback transaction"))
77}
78
79fn database_operation_with_context(error: sea_orm::DbErr, context: &str) -> DbError {
80    let kind = database_error_kind(&error);
81    let message = format!("{context}: {error}");
82    match kind {
83        Some(kind) => DbError::database_operation_classified(message, kind),
84        None => DbError::database_operation(message),
85    }
86}
87
88fn transaction_delay(config: &crate::retry::RetryConfig, attempt: u32) -> Duration {
89    aster_forge_utils::backoff::cap_delay(
90        aster_forge_utils::backoff::exponential_delay(
91            Duration::from_millis(config.base_delay_ms),
92            attempt,
93        ),
94        Duration::from_millis(config.max_delay_ms),
95    )
96}
97
98#[derive(Clone, Copy, Debug, Eq, PartialEq)]
99enum CommitFailureAction {
100    Retry,
101    ReturnKnownFailure,
102    ReturnOutcomeUnknown,
103}
104
105fn commit_failure_action(
106    retryable: bool,
107    outcome_known_rolled_back: bool,
108    attempt: u32,
109    max_retries: u32,
110) -> CommitFailureAction {
111    if !outcome_known_rolled_back {
112        CommitFailureAction::ReturnOutcomeUnknown
113    } else if retryable {
114        if attempt < max_retries {
115            CommitFailureAction::Retry
116        } else {
117            CommitFailureAction::ReturnKnownFailure
118        }
119    } else {
120        CommitFailureAction::ReturnKnownFailure
121    }
122}
123
124fn commit_outcome_known_rolled_back(kind: Option<crate::DatabaseErrorKind>) -> bool {
125    matches!(
126        kind,
127        Some(crate::DatabaseErrorKind::Deadlock | crate::DatabaseErrorKind::SerializationFailure)
128    )
129}
130
131/// Runs a complete transaction boundary with bounded retries selected by the product.
132///
133/// The callback is rerun only after rollback or a commit failure known to have rolled back the
134/// transaction. Any commit error with an uncertain server-side outcome is returned as
135/// `DbError::CommitOutcomeUnknown`, because the server may have committed the transaction even
136/// though the client did not receive a success response.
137///
138/// Use [`crate::retry::RetryConfig::deadlock`] as the config profile: deadlock and
139/// serialization conflicts resolve inside short lock-wait windows, so the fast profile fits.
140pub async fn with_transaction_retry<C, F, T, E, P>(
141    db: &C,
142    config: &crate::retry::RetryConfig,
143    mut operation: F,
144    should_retry: P,
145) -> std::result::Result<T, E>
146where
147    C: sea_orm::TransactionTrait,
148    F: for<'txn> FnMut(
149        &'txn C::Transaction,
150    )
151        -> Pin<Box<dyn Future<Output = std::result::Result<T, E>> + Send + 'txn>>,
152    E: From<DbError> + std::fmt::Display,
153    P: Fn(&E) -> bool,
154{
155    let mut attempt = 0;
156    loop {
157        let txn = match db.begin().await {
158            Ok(txn) => txn,
159            Err(error) => {
160                let error = E::from(database_operation_with_context(error, "begin transaction"));
161                if attempt < config.max_retries && should_retry(&error) {
162                    tokio::time::sleep(transaction_delay(config, attempt)).await;
163                    attempt += 1;
164                    continue;
165                }
166                return Err(error);
167            }
168        };
169
170        match operation(&txn).await {
171            Ok(value) => match txn.commit().await {
172                Ok(()) => return Ok(value),
173                Err(error) => {
174                    let kind = database_error_kind(&error);
175                    let classified_error = E::from(match kind {
176                        Some(kind) => DbError::database_operation_classified(
177                            format!("commit transaction: {error}"),
178                            kind,
179                        ),
180                        None => DbError::database_operation(format!("commit transaction: {error}")),
181                    });
182                    match commit_failure_action(
183                        should_retry(&classified_error),
184                        commit_outcome_known_rolled_back(kind),
185                        attempt,
186                        config.max_retries,
187                    ) {
188                        CommitFailureAction::Retry => {
189                            tokio::time::sleep(transaction_delay(config, attempt)).await;
190                            attempt += 1;
191                            continue;
192                        }
193                        CommitFailureAction::ReturnKnownFailure => return Err(classified_error),
194                        CommitFailureAction::ReturnOutcomeUnknown => {
195                            return Err(E::from(DbError::commit_outcome_unknown(
196                                format!("commit transaction: {error}"),
197                                kind,
198                            )));
199                        }
200                    }
201                }
202            },
203            Err(error) => {
204                if let Err(rollback_error) = txn.rollback().await {
205                    tracing::warn!(
206                        callback_error = %error,
207                        rollback_error = %rollback_error,
208                        "transaction rollback failed after callback error"
209                    );
210                }
211                if attempt < config.max_retries && should_retry(&error) {
212                    tokio::time::sleep(transaction_delay(config, attempt)).await;
213                    attempt += 1;
214                    continue;
215                }
216                return Err(error);
217            }
218        }
219    }
220}
221
222/// Runs a transaction callback with consistent tracing and rollback guarding.
223///
224/// The callback may return a product-specific error type. Forge-created transaction boundary
225/// errors are converted through `E: From<DbError>`, while callback errors are preserved unchanged.
226pub async fn with_transaction<C, F, T, E>(db: &C, operation: F) -> std::result::Result<T, E>
227where
228    C: sea_orm::TransactionTrait,
229    F: for<'txn> AsyncFnOnce(&'txn C::Transaction) -> std::result::Result<T, E>,
230    E: From<DbError> + std::fmt::Display,
231{
232    let location = Location::caller();
233    tracing::debug!(
234        file = location.file(),
235        line = location.line(),
236        "beginning transaction"
237    );
238    let txn = begin(db).await.map_err(E::from)?;
239    let mut rollback_guard = RollbackGuard::new(location);
240
241    match operation(&txn).await {
242        Ok(value) => {
243            rollback_guard.disarm();
244            commit(txn).await.map_err(E::from)?;
245            tracing::debug!(
246                file = location.file(),
247                line = location.line(),
248                "committed transaction"
249            );
250            Ok(value)
251        }
252        Err(error) => {
253            tracing::debug!(
254                file = location.file(),
255                line = location.line(),
256                error = %error,
257                "rolling back transaction after callback error"
258            );
259            rollback_guard.disarm();
260            if let Err(rollback_error) = rollback(txn).await {
261                tracing::error!(
262                    file = location.file(),
263                    line = location.line(),
264                    callback_error = %error,
265                    rollback_error = %rollback_error,
266                    "transaction rollback failed after callback error"
267                );
268            }
269            Err(error)
270        }
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::{
277        CommitFailureAction, commit_failure_action, commit_outcome_known_rolled_back, rollback,
278        transaction_delay, with_transaction,
279    };
280    use crate::{DbError, connection::DatabaseConfig};
281    use sea_orm::{ConnectionTrait, DatabaseConnection, Statement, TransactionTrait};
282    use std::fmt;
283
284    async fn sqlite_db() -> DatabaseConnection {
285        crate::connection::connect(&DatabaseConfig {
286            url: "sqlite::memory:".to_string(),
287            pool_size: 1,
288            retry_count: 0,
289        })
290        .await
291        .expect("sqlite memory database should connect")
292    }
293
294    async fn count_rows(db: &DatabaseConnection) -> i64 {
295        let statement = Statement::from_string(
296            sea_orm::DbBackend::Sqlite,
297            "SELECT COUNT(*) FROM transaction_items",
298        );
299        let row = db
300            .query_one_raw(statement)
301            .await
302            .expect("count query should succeed")
303            .expect("count query should return one row");
304        row.try_get_by_index(0).expect("count should decode")
305    }
306
307    #[derive(Debug, PartialEq, Eq)]
308    enum ProductError {
309        Db(String),
310        Validation(&'static str),
311    }
312
313    impl From<DbError> for ProductError {
314        fn from(value: DbError) -> Self {
315            Self::Db(value.to_string())
316        }
317    }
318
319    impl fmt::Display for ProductError {
320        fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
321            match self {
322                Self::Db(message) => formatter.write_str(message),
323                Self::Validation(message) => formatter.write_str(message),
324            }
325        }
326    }
327
328    #[tokio::test]
329    async fn with_transaction_commits_successful_callback() {
330        let db = sqlite_db().await;
331        db.execute_unprepared("CREATE TABLE transaction_items (id INTEGER PRIMARY KEY);")
332            .await
333            .expect("table should be created");
334
335        let value = with_transaction(&db, async |txn| {
336            txn.execute_unprepared("INSERT INTO transaction_items (id) VALUES (1);")
337                .await
338                .map_err(DbError::from)?;
339            Ok::<_, DbError>("committed")
340        })
341        .await
342        .expect("transaction should commit");
343
344        assert_eq!(value, "committed");
345        assert_eq!(count_rows(&db).await, 1);
346    }
347
348    #[tokio::test]
349    async fn with_transaction_rolls_back_callback_error() {
350        let db = sqlite_db().await;
351        db.execute_unprepared("CREATE TABLE transaction_items (id INTEGER PRIMARY KEY);")
352            .await
353            .expect("table should be created");
354
355        let error = with_transaction(&db, async |txn| {
356            txn.execute_unprepared("INSERT INTO transaction_items (id) VALUES (1);")
357                .await
358                .map_err(DbError::from)?;
359            Err::<(), _>(DbError::database_operation("forced failure"))
360        })
361        .await
362        .expect_err("callback error should propagate");
363
364        assert!(matches!(error, DbError::DatabaseOperation(_)));
365        assert_eq!(count_rows(&db).await, 0);
366    }
367
368    #[tokio::test]
369    async fn with_transaction_preserves_product_callback_errors() {
370        let db = sqlite_db().await;
371        db.execute_unprepared("CREATE TABLE transaction_items (id INTEGER PRIMARY KEY);")
372            .await
373            .expect("table should be created");
374
375        let error = with_transaction(&db, async |txn| {
376            txn.execute_unprepared("INSERT INTO transaction_items (id) VALUES (1);")
377                .await
378                .map_err(DbError::from)
379                .map_err(ProductError::from)?;
380            Err::<(), _>(ProductError::Validation("business validation failed"))
381        })
382        .await
383        .expect_err("callback error should propagate");
384
385        assert_eq!(
386            error,
387            ProductError::Validation("business validation failed")
388        );
389        assert_eq!(count_rows(&db).await, 0);
390    }
391
392    #[tokio::test]
393    async fn with_transaction_retry_restarts_after_callback_failure() {
394        use std::sync::{
395            Arc,
396            atomic::{AtomicUsize, Ordering},
397        };
398
399        let db = sqlite_db().await;
400        let attempts = Arc::new(AtomicUsize::new(0));
401        let config = crate::retry::RetryConfig {
402            max_retries: 2,
403            base_delay_ms: 0,
404            max_delay_ms: 0,
405        };
406        let result = {
407            let attempts = Arc::clone(&attempts);
408            super::with_transaction_retry(
409                &db,
410                &config,
411                move |_txn| {
412                    let attempts = Arc::clone(&attempts);
413                    Box::pin(async move {
414                        let attempt = attempts.fetch_add(1, Ordering::SeqCst);
415                        if attempt < 2 {
416                            Err(DbError::database_operation("temporary failure"))
417                        } else {
418                            Ok::<_, DbError>("ok")
419                        }
420                    })
421                },
422                |error| matches!(error, DbError::DatabaseOperation(_)),
423            )
424            .await
425        };
426
427        assert_eq!(result.unwrap(), "ok");
428        assert_eq!(attempts.load(Ordering::SeqCst), 3);
429    }
430
431    #[test]
432    fn transaction_delay_doubles_caps_and_handles_boundaries() {
433        let config = crate::retry::RetryConfig {
434            max_retries: 3,
435            base_delay_ms: 100,
436            max_delay_ms: 250,
437        };
438        assert_eq!(
439            transaction_delay(&config, 0),
440            std::time::Duration::from_millis(100)
441        );
442        assert_eq!(
443            transaction_delay(&config, 1),
444            std::time::Duration::from_millis(200)
445        );
446        assert_eq!(
447            transaction_delay(&config, 2),
448            std::time::Duration::from_millis(250)
449        );
450        assert_eq!(
451            transaction_delay(&config, u32::MAX),
452            std::time::Duration::from_millis(250)
453        );
454
455        let zero_max = crate::retry::RetryConfig {
456            max_retries: 1,
457            base_delay_ms: 100,
458            max_delay_ms: 0,
459        };
460        assert_eq!(transaction_delay(&zero_max, 0), std::time::Duration::ZERO);
461
462        let initial_above_max = crate::retry::RetryConfig {
463            max_retries: 1,
464            base_delay_ms: 1_000,
465            max_delay_ms: 250,
466        };
467        assert_eq!(
468            transaction_delay(&initial_above_max, 0),
469            std::time::Duration::from_millis(250)
470        );
471    }
472
473    #[tokio::test]
474    async fn with_transaction_retry_zero_budget_runs_callback_once() {
475        use std::sync::{
476            Arc,
477            atomic::{AtomicUsize, Ordering},
478        };
479
480        let db = sqlite_db().await;
481        let attempts = Arc::new(AtomicUsize::new(0));
482        let config = crate::retry::RetryConfig {
483            max_retries: 0,
484            base_delay_ms: 0,
485            max_delay_ms: 0,
486        };
487        let error = {
488            let attempts = attempts.clone();
489            super::with_transaction_retry(
490                &db,
491                &config,
492                move |_txn| {
493                    let attempts = attempts.clone();
494                    Box::pin(async move {
495                        attempts.fetch_add(1, Ordering::SeqCst);
496                        Err::<(), _>(DbError::database_operation("temporary failure"))
497                    })
498                },
499                |error| matches!(error, DbError::DatabaseOperation(_)),
500            )
501            .await
502            .unwrap_err()
503        };
504
505        assert!(matches!(error, DbError::DatabaseOperation(_)));
506        assert_eq!(attempts.load(Ordering::SeqCst), 1);
507    }
508
509    #[test]
510    fn commit_failure_action_preserves_known_exhausted_deadlocks() {
511        assert_eq!(
512            commit_failure_action(true, true, 0, 3),
513            CommitFailureAction::Retry
514        );
515        assert_eq!(
516            commit_failure_action(true, true, 3, 3),
517            CommitFailureAction::ReturnKnownFailure
518        );
519        assert_eq!(
520            commit_failure_action(false, false, 0, 3),
521            CommitFailureAction::ReturnOutcomeUnknown
522        );
523        assert_eq!(
524            commit_failure_action(true, false, 0, 3),
525            CommitFailureAction::ReturnOutcomeUnknown
526        );
527        assert_eq!(
528            commit_failure_action(false, true, 0, 3),
529            CommitFailureAction::ReturnKnownFailure
530        );
531        assert!(commit_outcome_known_rolled_back(Some(
532            crate::DatabaseErrorKind::Deadlock
533        )));
534        assert!(commit_outcome_known_rolled_back(Some(
535            crate::DatabaseErrorKind::SerializationFailure
536        )));
537        assert!(!commit_outcome_known_rolled_back(Some(
538            crate::DatabaseErrorKind::LockTimeout
539        )));
540    }
541
542    #[tokio::test]
543    async fn rollback_helper_discards_pending_changes() {
544        let db = sqlite_db().await;
545        db.execute_unprepared("CREATE TABLE transaction_items (id INTEGER PRIMARY KEY);")
546            .await
547            .expect("table should be created");
548        let txn = db.begin().await.expect("transaction should begin");
549        txn.execute_unprepared("INSERT INTO transaction_items (id) VALUES (1);")
550            .await
551            .expect("insert should succeed");
552
553        rollback(txn).await.expect("rollback should succeed");
554
555        assert_eq!(count_rows(&db).await, 0);
556    }
557}