1use 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
56pub 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
65pub 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
72pub 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
131pub 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
222pub 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}