use core::sync::atomic::{AtomicU32, Ordering};
use drizzle_core::error::{DrizzleError, Result};
fn cleanup_error(
scope: &str,
original: DrizzleError,
action: &str,
cleanup: DrizzleError,
) -> DrizzleError {
DrizzleError::TransactionError(
format!("{scope} callback failed: {original}; {action} failed: {cleanup}").into(),
)
}
fn trace_panic_cleanup_error(scope: &str, name: &str, action: &str, err: &DrizzleError) {
#[cfg(feature = "tracing")]
tracing::error!(
scope,
name,
action,
error = %err,
"transaction cleanup failed after panic"
);
#[cfg(not(feature = "tracing"))]
let _ = (scope, name, action, err);
}
pub fn sync_transaction<Tx, R>(
transaction: Tx,
trace_name: &'static str,
trace_commit: impl Fn(),
trace_rollback: impl Fn(),
body: impl FnOnce(&Tx) -> Result<R>,
commit: impl FnOnce(Tx) -> Result<()>,
rollback: impl FnOnce(Tx) -> Result<()>,
) -> Result<R> {
let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| body(&transaction)));
match outcome {
Ok(Ok(value)) => {
trace_commit();
commit(transaction)?;
Ok(value)
}
Ok(Err(e)) => {
trace_rollback();
match rollback(transaction) {
Ok(()) => Err(e),
Err(rollback_err) => Err(cleanup_error("transaction", e, "rollback", rollback_err)),
}
}
Err(panic_payload) => {
trace_rollback();
if let Err(err) = rollback(transaction) {
trace_panic_cleanup_error("transaction", trace_name, "rollback", &err);
}
std::panic::resume_unwind(panic_payload);
}
}
}
pub fn sync_savepoint<R>(
depth: &AtomicU32,
mut execute_raw: impl FnMut(&str) -> Result<()>,
body: impl FnOnce() -> Result<R>,
) -> Result<R> {
let level = depth.load(Ordering::Relaxed);
let sp = format!("drizzle_sp_{level}");
depth.store(level + 1, Ordering::Relaxed);
execute_raw(&format!("SAVEPOINT {sp}"))?;
let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(body));
depth.store(level, Ordering::Relaxed);
match outcome {
Ok(Ok(value)) => {
execute_raw(&format!("RELEASE SAVEPOINT {sp}"))?;
Ok(value)
}
Ok(Err(e)) => {
if let Err(rollback_err) = execute_raw(&format!("ROLLBACK TO SAVEPOINT {sp}")) {
return Err(cleanup_error(
"savepoint",
e,
"rollback to savepoint",
rollback_err,
));
}
if let Err(release_err) = execute_raw(&format!("RELEASE SAVEPOINT {sp}")) {
return Err(cleanup_error(
"savepoint",
e,
"release savepoint after rollback",
release_err,
));
}
Err(e)
}
Err(panic_payload) => {
if let Err(err) = execute_raw(&format!("ROLLBACK TO SAVEPOINT {sp}")) {
trace_panic_cleanup_error("savepoint", &sp, "rollback to savepoint", &err);
}
if let Err(err) = execute_raw(&format!("RELEASE SAVEPOINT {sp}")) {
trace_panic_cleanup_error(
"savepoint",
&sp,
"release savepoint after rollback",
&err,
);
}
std::panic::resume_unwind(panic_payload);
}
}
}
pub async fn async_savepoint<R, Exec, ExecFut, BodyFut>(
depth: &AtomicU32,
mut execute_raw: Exec,
body: BodyFut,
) -> Result<R>
where
Exec: FnMut(String) -> ExecFut,
ExecFut: core::future::Future<Output = Result<()>>,
BodyFut: core::future::Future<Output = Result<R>>,
{
let level = depth.load(Ordering::Relaxed);
let sp = format!("drizzle_sp_{level}");
depth.store(level + 1, Ordering::Relaxed);
execute_raw(format!("SAVEPOINT {sp}")).await?;
let outcome = body.await;
depth.store(level, Ordering::Relaxed);
match outcome {
Ok(value) => {
execute_raw(format!("RELEASE SAVEPOINT {sp}")).await?;
Ok(value)
}
Err(e) => {
if let Err(rollback_err) = execute_raw(format!("ROLLBACK TO SAVEPOINT {sp}")).await {
return Err(cleanup_error(
"savepoint",
e,
"rollback to savepoint",
rollback_err,
));
}
if let Err(release_err) = execute_raw(format!("RELEASE SAVEPOINT {sp}")).await {
return Err(cleanup_error(
"savepoint",
e,
"release savepoint after rollback",
release_err,
));
}
Err(e)
}
}
}