cratestack_sqlx/
isolation.rs

1//! Helpers for running banking-grade multi-row mutations under explicit
2//! transaction isolation, with retry on serialization failure.
3//!
4//! These are the hand-rolled form, for code that owns a pool. A
5//! *procedure* declaring `@isolation("...")` in the schema does not need
6//! them: its generated dispatch runs the procedure's authorization and
7//! body inside a transaction at the declared level, with its own retry
8//! loop, on every transport (REST, RPC, RPC batch, MCP and
9//! `invoke_with_db`) — see `SqlxRuntime::run_isolated` and
10//! docs/design/procedure-isolation.md. Before GHSA-r67q-4qqq-g9gm was
11//! fixed, this comment claimed a `ProcedureMetadata` constant recorded
12//! the level for dispatch to use; no such constant ever existed and the
13//! attribute had no effect.
14//!
15//! Like `db.transaction(...)` (see `crate::transaction`'s module doc,
16//! cratestack#534), composing write-builder `run_in_tx` calls inside
17//! `body` here does not get you automatic `AuditSink` fan-out or `@@emit`
18//! delivery either, for the identical reason: `body` returns to this
19//! helper's own retry loop, not to a commit hook this crate owns, so
20//! dispatch remains the caller's responsibility after the outer call
21//! returns `Ok`.
22use crate::sqlx;
23
24use std::future::Future;
25
26use cratestack_core::{CratestackError, TransactionIsolation};
27
28use crate::error::cratestack_error_from_sqlx;
29
30pub(crate) const MAX_RETRIES_DEFAULT: u32 = 3;
31
32/// Begin a transaction at the requested isolation level, run `body` against
33/// the live transaction, and commit. On `40001` (serialization_failure) or
34/// `40P01` (deadlock_detected) the transaction is rolled back and the body
35/// runs again, up to `MAX_RETRIES_DEFAULT` times. Other errors propagate
36/// immediately. Only a database error is retried: an error `body` builds
37/// itself (`Validation`, `Internal`, …) is returned as is even when its
38/// text contains `40001` or `deadlock detected`, and a typed database error
39/// is retried by its SQLSTATE alone, never by its text.
40///
41/// `body` receives a mutable transaction reference; it should run all of
42/// its SQL through that reference so the writes participate in the same
43/// transaction.
44pub async fn run_in_isolated_tx<F, Fut, T>(
45    pool: &sqlx::PgPool,
46    isolation: TransactionIsolation,
47    body: F,
48) -> Result<T, CratestackError>
49where
50    F: FnMut(sqlx::Transaction<'static, sqlx::Postgres>) -> Fut,
51    Fut: Future<Output = Result<(T, sqlx::Transaction<'static, sqlx::Postgres>), CratestackError>>,
52{
53    run_in_isolated_tx_with_retries(pool, isolation, MAX_RETRIES_DEFAULT, body).await
54}
55
56/// Same as [`run_in_isolated_tx`] but with a caller-chosen retry budget.
57/// Banks running long-tail contended writes sometimes want a higher cap
58/// (5–10); single-row CAS workflows can drop to 1 to fail fast.
59pub async fn run_in_isolated_tx_with_retries<F, Fut, T>(
60    pool: &sqlx::PgPool,
61    isolation: TransactionIsolation,
62    max_retries: u32,
63    mut body: F,
64) -> Result<T, CratestackError>
65where
66    F: FnMut(sqlx::Transaction<'static, sqlx::Postgres>) -> Fut,
67    Fut: Future<Output = Result<(T, sqlx::Transaction<'static, sqlx::Postgres>), CratestackError>>,
68{
69    let mut attempts = 0u32;
70    loop {
71        attempts += 1;
72        let mut tx = pool.begin().await.map_err(cratestack_error_from_sqlx)?;
73        let set_stmt = format!("SET TRANSACTION ISOLATION LEVEL {}", isolation.as_sql());
74        // `AssertSqlSafe`: `isolation.as_sql()` returns one of four
75        // `&'static str` literals, so nothing request-derived reaches this
76        // statement (sqlx 0.9's `SqlSafeStr` bound).
77        sqlx::query(sqlx::AssertSqlSafe(set_stmt))
78            .execute(&mut *tx)
79            .await
80            .map_err(cratestack_error_from_sqlx)?;
81
82        match body(tx).await {
83            Ok((value, tx)) => match tx.commit().await {
84                Ok(()) => return Ok(value),
85                Err(commit_error) => {
86                    // PG can defer a serialization anomaly all the way to
87                    // COMMIT: the body's SQL runs cleanly, then the engine
88                    // detects the conflict during the predicate-lock check
89                    // at commit and rolls the transaction back with
90                    // SQLSTATE 40001 (the docs are explicit that the
91                    // *entire* transaction must be retried). Without this
92                    // branch we'd advertise automatic retries but still
93                    // leak a transient 40001 to callers when the conflict
94                    // is detected at the commit boundary.
95                    let promoted = cratestack_error_from_sqlx(commit_error);
96                    if attempts <= max_retries && is_retriable(&promoted) {
97                        tokio::task::yield_now().await;
98                        continue;
99                    }
100                    return Err(promoted);
101                }
102            },
103            Err(error) => {
104                if attempts <= max_retries && is_retriable(&error) {
105                    // Backoff is intentionally trivial — banks running this
106                    // under heavy contention should swap to a more thoughtful
107                    // jittered backoff. Sub-millisecond pause yields the
108                    // current task without keeping a tx open.
109                    tokio::task::yield_now().await;
110                    continue;
111                }
112                return Err(error);
113            }
114        }
115    }
116}
117
118/// The same classifier as an `@isolation` procedure's dispatch
119/// ([`crate::retriable::retriable_sqlstate`]): the SQLSTATE of a typed
120/// database error — authoritative, its text never consulted — else the
121/// driver's text of the untyped `Database(String)` variant. Only database
122/// errors: an application or validation error whose message happens to
123/// contain `40001` (a body echoing request data) is not a serialization
124/// failure, and neither is a typed `P0001`/`22P02` whose message echoes it;
125/// retrying either re-ran `body` with its side effects and returned the
126/// wrong error (docs/design/procedure-isolation.md §5).
127fn is_retriable(error: &CratestackError) -> bool {
128    crate::retriable::retriable_sqlstate(error).is_some()
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    #[test]
136    fn parses_all_isolation_levels() {
137        assert_eq!(
138            TransactionIsolation::parse("serializable").unwrap(),
139            TransactionIsolation::Serializable,
140        );
141        assert_eq!(
142            TransactionIsolation::parse("Repeatable_Read").unwrap(),
143            TransactionIsolation::RepeatableRead,
144        );
145        assert_eq!(
146            TransactionIsolation::parse("read committed").unwrap(),
147            TransactionIsolation::ReadCommitted,
148        );
149        assert!(TransactionIsolation::parse("snapshot").is_err());
150    }
151
152    #[test]
153    fn sql_strings_match_pg_grammar() {
154        assert_eq!(TransactionIsolation::Serializable.as_sql(), "SERIALIZABLE");
155        assert_eq!(
156            TransactionIsolation::RepeatableRead.as_sql(),
157            "REPEATABLE READ",
158        );
159        assert_eq!(
160            TransactionIsolation::ReadCommitted.as_sql(),
161            "READ COMMITTED",
162        );
163    }
164
165    #[test]
166    fn retriable_on_serialization_failure_sqlstate() {
167        let err = CratestackError::Database(
168            "Database(PgDatabaseError { severity: ERROR, code: \"40001\", \
169             message: \"could not serialize access due to concurrent update\" })"
170                .to_owned(),
171        );
172        assert!(is_retriable(&err));
173    }
174
175    #[test]
176    fn retriable_on_deadlock_sqlstate() {
177        let err = CratestackError::Database(
178            "Database(PgDatabaseError { code: \"40P01\", \
179             message: \"deadlock detected\" })"
180                .to_owned(),
181        );
182        assert!(is_retriable(&err));
183    }
184
185    #[test]
186    fn not_retriable_on_an_application_error_that_mentions_40001() {
187        for err in [
188            CratestackError::Validation("field 'memo' length 40001 exceeds maximum 100".to_owned()),
189            CratestackError::Internal("deadlock detected in my own lock manager".to_owned()),
190        ] {
191            assert!(!is_retriable(&err), "{err:?}");
192        }
193    }
194
195    #[test]
196    fn not_retriable_on_unique_violation() {
197        let err = CratestackError::Database(
198            "duplicate key value violates unique constraint \"accounts_pkey\"".to_owned(),
199        );
200        assert!(!is_retriable(&err));
201    }
202
203    #[test]
204    fn retriable_when_serialization_failure_is_raised_at_commit_time() {
205        // PG SSI can defer the 40001 to COMMIT. The sqlx error surfaced
206        // by `tx.commit()` carries the same SQLSTATE; the loop now
207        // promotes that into `CratestackError::Database` and feeds it through
208        // `is_retriable` so the commit-time path is no longer leaked to
209        // callers despite the API advertising automatic retries.
210        let err = CratestackError::Database(
211            "Database(PgDatabaseError { severity: ERROR, code: \"40001\", \
212             message: \"could not serialize access due to read/write dependencies among transactions\" })"
213                .to_owned(),
214        );
215        assert!(is_retriable(&err));
216    }
217
218    // --- typed-variant paths ---
219
220    #[test]
221    fn retriable_typed_serialization_failure() {
222        use cratestack_core::DbErrorInfo;
223        let err = CratestackError::DatabaseTyped(DbErrorInfo {
224            detail: "could not serialize access due to concurrent update".to_owned(),
225            sqlstate: Some("40001".to_owned()),
226            constraint: None,
227        });
228        assert!(
229            is_retriable(&err),
230            "DatabaseTyped with 40001 sqlstate must be retriable via the fast path",
231        );
232    }
233
234    #[test]
235    fn retriable_typed_deadlock() {
236        use cratestack_core::DbErrorInfo;
237        let err = CratestackError::DatabaseTyped(DbErrorInfo {
238            detail: "deadlock detected".to_owned(),
239            sqlstate: Some("40P01".to_owned()),
240            constraint: None,
241        });
242        assert!(
243            is_retriable(&err),
244            "DatabaseTyped with 40P01 sqlstate must be retriable via the fast path",
245        );
246    }
247
248    #[test]
249    fn not_retriable_typed_unique_violation() {
250        use cratestack_core::DbErrorInfo;
251        let err = CratestackError::DatabaseTyped(DbErrorInfo {
252            detail: "duplicate key value violates unique constraint \"accounts_pkey\"".to_owned(),
253            sqlstate: Some("23505".to_owned()),
254            constraint: Some("accounts_pkey".to_owned()),
255        });
256        assert!(
257            !is_retriable(&err),
258            "unique_violation (23505) must not be retried",
259        );
260    }
261
262    #[test]
263    fn typed_variant_sqlstate_is_authoritative_over_its_detail() {
264        // A typed error's SQLSTATE decides; its text is never consulted.
265        // The text of other SQLSTATEs echoes request data: a body's
266        // `RAISE EXCEPTION 'insufficient funds: requested %'` (P0001), a
267        // cast of `'40001x'` (22P02). Matching it re-ran the body (and, in
268        // `@isolation` dispatch, answered 409 and released the idempotency
269        // key) — request-triggerable. An unknown SQLSTATE is not retried
270        // either, whatever its detail says.
271        use cratestack_core::DbErrorInfo;
272        for (sqlstate, detail) in [
273            (
274                "XX999",
275                "could not serialize access due to read/write dependencies",
276            ),
277            ("P0001", "insufficient funds: requested 40001"),
278            ("22P02", "invalid input syntax for type bigint: \"40001x\""),
279            ("55P03", "deadlock detected while waiting for 40P01"),
280        ] {
281            let err = CratestackError::DatabaseTyped(DbErrorInfo {
282                detail: detail.to_owned(),
283                sqlstate: Some(sqlstate.to_owned()),
284                constraint: None,
285            });
286            assert!(
287                !is_retriable(&err),
288                "typed {sqlstate} must not be retried because of its text",
289            );
290        }
291    }
292
293    #[test]
294    fn typed_variant_exposes_constraint_for_unique_violation() {
295        use cratestack_core::DbErrorInfo;
296        let err = CratestackError::DatabaseTyped(DbErrorInfo {
297            detail: "duplicate key value violates unique constraint \"wallets_owner_key\""
298                .to_owned(),
299            sqlstate: Some("23505".to_owned()),
300            constraint: Some("wallets_owner_key".to_owned()),
301        });
302        assert_eq!(err.db_sqlstate(), Some("23505"));
303        assert_eq!(err.db_constraint(), Some("wallets_owner_key"));
304        // Public message must remain canned — no detail leak.
305        assert_eq!(err.public_message(), "internal error");
306    }
307}
cratestack_sqlx/isolation.rs

cratestack_sqlx/
isolation.rs
