use std::future::Future;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, PoisonError};
use tokio::sync::OwnedMutexGuard;
use super::{transaction_pool, ExecError, PoolTx};
use crate::sql::{Pool, PoolId};
type Callback = Box<dyn FnOnce() + Send>;
struct Slot {
state: Arc<tokio::sync::Mutex<TxState>>,
cancelled: AtomicUsize,
guards: AtomicUsize,
claims: AtomicUsize,
fatal: Arc<AtomicBool>,
tx_key: usize,
}
struct Claim<'a>(&'a Slot);
impl<'a> Claim<'a> {
fn new(slot: &'a Slot) -> Self {
slot.claims.fetch_add(1, Ordering::SeqCst);
Self(slot)
}
}
impl Drop for Claim<'_> {
fn drop(&mut self) {
self.0.claims.fetch_sub(1, Ordering::SeqCst);
}
}
struct TxState {
tx: Option<PoolTx<'static>>,
depth: usize,
poisoned: bool,
}
impl TxState {
fn tx(&mut self) -> &mut PoolTx<'static> {
self.tx.as_mut().expect("atomic transaction is open")
}
async fn savepoint(&mut self, verb: &str, depth: usize) -> Result<(), ExecError> {
let tx = self.tx();
let name = tx.dialect().quote_ident(&format!("rustango_sp_{depth}"));
let r = tx.execute_unprepared(&format!("{verb} {name}")).await;
if r.is_err() {
self.poisoned = true;
}
Ok(r?)
}
async fn rollback_to(&mut self, depth: usize) -> Result<(), ExecError> {
self.savepoint("ROLLBACK TO SAVEPOINT", depth).await?;
self.savepoint("RELEASE SAVEPOINT", depth).await?;
self.depth = depth - 1;
Ok(())
}
async fn check_open(&mut self, slot: &Slot) -> Result<(), ExecError> {
if self.tx().still_open().await {
return Ok(());
}
self.poisoned = true;
let mysql = self.tx().dialect().name() == "mysql";
Err(if mysql && !slot.fatal.load(Ordering::SeqCst) {
ExecError::AtomicEndedEarly
} else {
ExecError::AtomicAborted
})
}
async fn settle(&mut self, slot: &Slot) -> Result<(), ExecError> {
if self.poisoned || slot.fatal.load(Ordering::SeqCst) {
return Err(ExecError::AtomicAborted);
}
let d = slot.cancelled.swap(0, Ordering::SeqCst);
if d != 0 && d <= self.depth {
self.rollback_to(d).await?;
}
Ok(())
}
}
async fn in_background<T: Send + 'static>(
fut: impl Future<Output = Result<T, ExecError>> + Send + 'static,
) -> Result<T, ExecError> {
match tokio::spawn(fut).await {
Ok(r) => r,
Err(e) if e.is_panic() => std::panic::resume_unwind(e.into_panic()),
Err(_) => Err(ExecError::AtomicAborted),
}
}
async fn acquire(slot: &Arc<Slot>) -> Result<OwnedMutexGuard<TxState>, ExecError> {
let busy =
|s: &Slot| s.guards.load(Ordering::SeqCst) > 0 || s.claims.load(Ordering::SeqCst) > 0;
let guard = loop {
match Arc::clone(&slot.state).try_lock_owned() {
Ok(g) => break g,
Err(_) if busy(slot) => return Err(ExecError::NestedAtomic),
Err(_) => tokio::time::sleep(std::time::Duration::from_millis(1)).await,
}
};
if !guard.poisoned
&& !slot.fatal.load(Ordering::SeqCst)
&& slot.cancelled.load(Ordering::SeqCst) == 0
{
return Ok(guard);
}
let _claim = Claim::new(slot);
let task_slot = Arc::clone(slot);
in_background(async move {
let mut g = guard;
g.settle(&task_slot).await?;
Ok(g)
})
.await
}
pub(crate) struct Gate(Option<Arc<Slot>>);
impl Gate {
pub(crate) const OPEN: Self = Self(None);
}
pub(crate) fn gate(tx: &PoolTx<'_>) -> Result<Gate, ExecError> {
let key = std::ptr::from_ref::<PoolTx<'_>>(tx) as usize;
let slot = BLOCK
.try_with(|b| {
std::iter::successors(Some(Arc::clone(b)), |b| b.enclosing.clone())
.map(|b| Arc::clone(&b.slot))
.find(|s| s.tx_key == key)
})
.ok()
.flatten();
if slot
.as_ref()
.is_some_and(|s| s.fatal.load(Ordering::SeqCst))
{
return Err(ExecError::AtomicAborted);
}
Ok(Gate(slot))
}
impl Gate {
pub(crate) fn check<T>(&self, r: Result<T, ExecError>) -> Result<T, ExecError> {
if let (Some(slot), Err(ExecError::Driver(e))) = (&self.0, &r) {
if ends_transaction(e) {
slot.fatal.store(true, Ordering::SeqCst);
}
}
r
}
}
fn ends_transaction(e: &sqlx::Error) -> bool {
match e {
sqlx::Error::Io(_) | sqlx::Error::PoolClosed | sqlx::Error::WorkerCrashed => true,
#[allow(unused_variables)]
sqlx::Error::Database(db) => {
#[cfg(feature = "mysql")]
if db
.try_downcast_ref::<sqlx::mysql::MySqlDatabaseError>()
.is_some_and(|m| matches!(m.number(), 1213 | 1205))
{
return true;
}
#[cfg(feature = "postgres")]
if db
.try_downcast_ref::<sqlx::postgres::PgDatabaseError>()
.is_some_and(|p| {
use sqlx::postgres::PgSeverity;
matches!(p.severity(), PgSeverity::Fatal | PgSeverity::Panic)
})
{
return true;
}
false
}
_ => false,
}
}
struct Block {
pool: PoolId,
slot: Arc<Slot>,
depth: usize,
enclosing: Option<Arc<Block>>,
callbacks: Mutex<Vec<Callback>>,
committed: AtomicBool,
}
impl Block {
fn take_callbacks(&self) -> Vec<Callback> {
std::mem::take(
&mut self
.callbacks
.lock()
.unwrap_or_else(PoisonError::into_inner),
)
}
}
tokio::task_local! {
static BLOCK: Arc<Block>;
}
pub struct AtomicTx {
block: Arc<Block>,
}
impl AtomicTx {
pub async fn lock(&self) -> Result<TxGuard<'_>, ExecError> {
let slot = &self.block.slot;
let guard = acquire(slot).await?;
if guard.depth != self.block.depth {
return Err(ExecError::NestedAtomic);
}
slot.guards.fetch_add(1, Ordering::SeqCst);
Ok(TxGuard {
guard,
slot: Arc::clone(slot),
_handle: PhantomData,
})
}
}
pub struct TxGuard<'a> {
guard: OwnedMutexGuard<TxState>,
slot: Arc<Slot>,
_handle: PhantomData<&'a AtomicTx>,
}
impl Drop for TxGuard<'_> {
fn drop(&mut self) {
self.slot.guards.fetch_sub(1, Ordering::SeqCst);
}
}
impl std::ops::Deref for TxGuard<'_> {
type Target = PoolTx<'static>;
fn deref(&self) -> &PoolTx<'static> {
self.guard.tx.as_ref().expect("atomic transaction is open")
}
}
impl std::ops::DerefMut for TxGuard<'_> {
fn deref_mut(&mut self) -> &mut PoolTx<'static> {
self.guard.tx()
}
}
#[derive(Default)]
struct Opening {
opened: AtomicBool,
abandoned: AtomicBool,
}
struct OpenSavepoint {
slot: Arc<Slot>,
depth: usize,
opening: Arc<Opening>,
ending: bool,
}
fn mark_cancelled(slot: &Slot, depth: usize) {
let _ = slot
.cancelled
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |c| {
Some(if c == 0 { depth } else { c.min(depth) })
});
}
impl Drop for OpenSavepoint {
fn drop(&mut self) {
if !self.ending {
self.opening.abandoned.store(true, Ordering::SeqCst);
if self.opening.opened.load(Ordering::SeqCst) {
mark_cancelled(&self.slot, self.depth);
}
}
}
}
pub(crate) fn in_block(pool: &Pool) -> bool {
let id = PoolId::of(pool);
BLOCK
.try_with(|b| {
std::iter::successors(Some(Arc::clone(b)), |b| b.enclosing.clone())
.any(|b| b.pool.same(&id))
})
.unwrap_or(false)
}
#[cfg(feature = "jobs")]
pub(crate) fn in_any_block() -> bool {
BLOCK
.try_with(|b| !b.committed.load(Ordering::SeqCst))
.unwrap_or(false)
}
pub async fn atomic<F, T>(pool: &Pool, f: F) -> Result<T, ExecError>
where
F: for<'tx> FnOnce(
&'tx AtomicTx,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<T, ExecError>> + Send + 'tx>,
>,
{
run(pool, End::Commit, f).await
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum End {
Commit,
Rollback,
}
pub(crate) async fn rolled_back<F, T>(pool: &Pool, f: F) -> Result<T, ExecError>
where
F: for<'tx> FnOnce(
&'tx AtomicTx,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<T, ExecError>> + Send + 'tx>,
>,
{
run(pool, End::Rollback, f).await
}
async fn run<F, T>(pool: &Pool, end: End, f: F) -> Result<T, ExecError>
where
F: for<'tx> FnOnce(
&'tx AtomicTx,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<T, ExecError>> + Send + 'tx>,
>,
{
let id = PoolId::of(pool);
let current = BLOCK.try_with(Arc::clone).ok();
let same_pool =
std::iter::successors(current.clone(), |b| b.enclosing.clone()).find(|b| b.pool.same(&id));
match same_pool {
Some(parent) => nested(parent, current, end, f).await,
None => outermost(pool, id, current, end, f).await,
}
}
async fn outermost<F, T>(
pool: &Pool,
id: PoolId,
enclosing: Option<Arc<Block>>,
end: End,
f: F,
) -> Result<T, ExecError>
where
F: for<'tx> FnOnce(
&'tx AtomicTx,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<T, ExecError>> + Send + 'tx>,
>,
{
let fatal = Arc::new(AtomicBool::new(false));
let mut tx = transaction_pool(pool).await?;
let hook = Arc::clone(&fatal);
tx.on_sqlite_rollback(move || hook.store(true, Ordering::SeqCst))
.await?;
let state = Arc::new(tokio::sync::Mutex::new(TxState {
tx: Some(tx),
depth: 0,
poisoned: false,
}));
let tx_key = state
.try_lock()
.ok()
.and_then(|g| g.tx.as_ref().map(|t| std::ptr::from_ref(t) as usize))
.unwrap_or(0);
let slot = Arc::new(Slot {
state,
cancelled: AtomicUsize::new(0),
guards: AtomicUsize::new(0),
claims: AtomicUsize::new(0),
fatal,
tx_key,
});
let block = Arc::new(Block {
pool: id,
slot: Arc::clone(&slot),
depth: 0,
enclosing,
callbacks: Mutex::new(Vec::new()),
committed: AtomicBool::new(false),
});
let handle = AtomicTx {
block: Arc::clone(&block),
};
let res = BLOCK
.scope(Arc::clone(&block), async { f(&handle).await })
.await;
let mut st = slot.state.lock().await;
let settled = match st.settle(&slot).await {
Ok(()) if st.depth != 0 => Err(ExecError::AtomicAborted),
Ok(()) if res.is_ok() && end == End::Commit => st.check_open(&slot).await,
other => other,
};
let mut tx = st.tx.take().expect("atomic transaction is open");
drop(st);
tx.clear_sqlite_rollback().await;
match res.and_then(|v| settled.map(|()| v)) {
Ok(v) if end == End::Rollback => {
tx.rollback().await?;
Ok(v)
}
Ok(v) => {
tx.commit().await?;
block.committed.store(true, Ordering::SeqCst);
let callbacks = block.take_callbacks();
BLOCK.sync_scope(block, || {
for cb in callbacks {
let run = std::panic::catch_unwind(std::panic::AssertUnwindSafe(cb));
if let Err(panic) = run {
tracing::error!(
target: "rustango::error",
panic = crate::panic_guard::panic_message(&*panic),
"on_commit callback panicked"
);
}
}
});
Ok(v)
}
Err(e) => {
let _ = tx.rollback().await;
Err(e)
}
}
}
async fn nested<F, T>(
parent: Arc<Block>,
enclosing: Option<Arc<Block>>,
end: End,
f: F,
) -> Result<T, ExecError>
where
F: for<'tx> FnOnce(
&'tx AtomicTx,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<T, ExecError>> + Send + 'tx>,
>,
{
let slot = Arc::clone(&parent.slot);
let guard = acquire(&slot).await?;
let parent_depth = parent.depth;
let depth = parent_depth + 1;
let opening = Arc::new(Opening::default());
let mut open = OpenSavepoint {
slot: Arc::clone(&slot),
depth,
opening: Arc::clone(&opening),
ending: false,
};
let (task_slot, task_opening) = (Arc::clone(&slot), Arc::clone(&opening));
let claim = Claim::new(&slot);
let opened = in_background(async move {
let mut g = guard;
if g.depth != parent_depth {
return Err(ExecError::NestedAtomic);
}
g.savepoint("SAVEPOINT", depth).await?;
g.depth = depth;
task_opening.opened.store(true, Ordering::SeqCst);
if task_opening.abandoned.load(Ordering::SeqCst) {
mark_cancelled(&task_slot, depth);
}
Ok(())
})
.await;
drop(claim);
if let Err(e) = opened {
open.ending = true;
return Err(e);
}
let block = Arc::new(Block {
pool: parent.pool.clone(),
slot: Arc::clone(&slot),
depth,
enclosing,
callbacks: Mutex::new(Vec::new()),
committed: AtomicBool::new(false),
});
let handle = AtomicTx {
block: Arc::clone(&block),
};
let res = BLOCK
.scope(Arc::clone(&block), async { f(&handle).await })
.await;
let guard = match acquire(&slot).await {
Ok(g) => g,
Err(e) => return Err(res.err().unwrap_or(e)),
};
open.ending = true;
let release = res.is_ok() && end == End::Commit;
let claim = Claim::new(&slot);
let ended = in_background(async move {
let mut g = guard;
if g.depth != depth {
g.poisoned = true;
return Err(ExecError::AtomicAborted);
}
if release {
g.savepoint("RELEASE SAVEPOINT", depth).await?;
g.depth = depth - 1;
let callbacks = block.take_callbacks();
parent
.callbacks
.lock()
.unwrap_or_else(PoisonError::into_inner)
.extend(callbacks);
Ok(())
} else {
g.rollback_to(depth).await
}
})
.await;
drop(claim);
let v = res?;
ended?;
Ok(v)
}
#[macro_export]
macro_rules! atomic {
($pool:expr, |$tx:ident| $body:block) => {
async {
let __rustango_atomic_pool = ::core::clone::Clone::clone($pool);
$crate::sql::atomic(&__rustango_atomic_pool, |$tx| {
::std::boxed::Box::pin(async move { $body })
})
.await
}
};
}
pub fn on_commit<F>(f: F)
where
F: FnOnce() + Send + 'static,
{
let block = BLOCK.try_with(Arc::clone).unwrap_or_else(|_| {
panic!(
"rustango::sql::on_commit called outside an `atomic` block — \
the callback would never fire. Wrap the caller in \
`atomic(&pool, |tx| async move {{ ... on_commit(...) ... }})`."
);
});
if block.committed.load(Ordering::SeqCst) {
f();
} else {
block
.callbacks
.lock()
.unwrap_or_else(PoisonError::into_inner)
.push(Box::new(f));
}
}
#[must_use]
pub fn on_commit_pending() -> usize {
std::iter::successors(BLOCK.try_with(Arc::clone).ok(), |b| b.enclosing.clone())
.map(|b| {
b.callbacks
.lock()
.unwrap_or_else(PoisonError::into_inner)
.len()
})
.sum()
}