use std::error::Error;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use hydracache_core::CacheCodec;
use hydracache_db::{
CollectedInvalidationReport, CommitPosition, DbCache, InvalidationCollector,
SqlxInvalidationOutbox,
};
use sqlx::{Sqlite, SqlitePool, Transaction};
use crate::{SqlxTransactionError, TransactionResult};
pub type SqlxTransactionFuture<'tx, E> =
Pin<Box<dyn Future<Output = std::result::Result<(), E>> + Send + 'tx>>;
pub trait SqlxTransactionExt<C>
where
C: CacheCodec,
{
fn sqlx_transactions(&self) -> SqlxTransactionCompanion<C>;
}
impl<C> SqlxTransactionExt<C> for DbCache<C>
where
C: CacheCodec,
{
fn sqlx_transactions(&self) -> SqlxTransactionCompanion<C> {
SqlxTransactionCompanion::new(self.clone())
}
}
#[derive(Debug)]
pub struct SqlxTransactionCompanion<C = hydracache::PostcardCodec>
where
C: CacheCodec,
{
queries: DbCache<C>,
outbox: Option<SqlxInvalidationOutbox>,
counters: Arc<SqlxTransactionCounters>,
}
impl<C> Clone for SqlxTransactionCompanion<C>
where
C: CacheCodec,
{
fn clone(&self) -> Self {
Self {
queries: self.queries.clone(),
outbox: self.outbox.clone(),
counters: self.counters.clone(),
}
}
}
impl<C> SqlxTransactionCompanion<C>
where
C: CacheCodec,
{
pub fn new(queries: DbCache<C>) -> Self {
Self {
queries,
outbox: None,
counters: Arc::new(SqlxTransactionCounters::default()),
}
}
pub fn with_outbox(mut self, outbox: SqlxInvalidationOutbox) -> Self {
self.outbox = Some(outbox);
self
}
pub fn diagnostics(&self) -> SqlxTransactionDiagnostics {
self.counters.snapshot()
}
pub async fn sqlite_durable<F, E>(
&self,
pool: &SqlitePool,
reason: impl Into<String>,
body: F,
) -> TransactionResult<SqlxTransactionReport, E>
where
F: for<'tx> FnOnce(
&'tx mut Transaction<'_, Sqlite>,
&'tx mut InvalidationCollector,
) -> SqlxTransactionFuture<'tx, E>
+ Send,
E: Error + Send + Sync + 'static,
{
let Some(outbox) = self.outbox.as_ref() else {
self.counters
.enqueue_failures
.fetch_add(1, Ordering::Relaxed);
return Err(SqlxTransactionError::MissingOutbox);
};
let mut tx = pool.begin().await.map_err(SqlxTransactionError::Sqlx)?;
let mut collector = InvalidationCollector::new(self.queries.namespace(), reason);
if let Err(error) = body(&mut tx, &mut collector).await {
self.counters.body_errors.fetch_add(1, Ordering::Relaxed);
rollback_and_count(tx, &self.counters).await;
return Err(SqlxTransactionError::Body(error));
}
let collected = collector.into_collected();
let intent_count = collected.len();
let inserted = if collected.is_empty() {
0
} else {
match outbox
.enqueue_in_sqlite_tx(
&mut tx,
collected.namespace(),
&sqlite_commit_position(),
collected.batch(),
)
.await
{
Ok(inserted) => inserted,
Err(error) => {
self.counters
.enqueue_failures
.fetch_add(1, Ordering::Relaxed);
rollback_and_count(tx, &self.counters).await;
return Err(SqlxTransactionError::Outbox(error));
}
}
};
match tx.commit().await {
Ok(()) => {
self.counters.commits.fetch_add(1, Ordering::Relaxed);
Ok(SqlxTransactionReport {
intent_count,
durable_rows: inserted,
local_report: None,
})
}
Err(error) => {
self.counters
.commit_failures
.fetch_add(1, Ordering::Relaxed);
Err(SqlxTransactionError::Sqlx(error))
}
}
}
pub async fn sqlite_local<F, E>(
&self,
pool: &SqlitePool,
reason: impl Into<String>,
body: F,
) -> TransactionResult<SqlxTransactionReport, E>
where
F: for<'tx> FnOnce(
&'tx mut Transaction<'_, Sqlite>,
&'tx mut InvalidationCollector,
) -> SqlxTransactionFuture<'tx, E>
+ Send,
E: Error + Send + Sync + 'static,
{
let mut tx = pool.begin().await.map_err(SqlxTransactionError::Sqlx)?;
let mut collector = InvalidationCollector::new(self.queries.namespace(), reason);
if let Err(error) = body(&mut tx, &mut collector).await {
self.counters.body_errors.fetch_add(1, Ordering::Relaxed);
rollback_and_count(tx, &self.counters).await;
return Err(SqlxTransactionError::Body(error));
}
let collected = collector.into_collected();
let intent_count = collected.len();
match tx.commit().await {
Ok(()) => {
self.counters.commits.fetch_add(1, Ordering::Relaxed);
}
Err(error) => {
self.counters
.commit_failures
.fetch_add(1, Ordering::Relaxed);
return Err(SqlxTransactionError::Sqlx(error));
}
}
let local_report = collected
.execute_local(self.queries.cache())
.await
.map_err(SqlxTransactionError::LocalInvalidation)?;
self.counters
.local_invalidations
.fetch_add(1, Ordering::Relaxed);
Ok(SqlxTransactionReport {
intent_count,
durable_rows: 0,
local_report: Some(local_report),
})
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct SqlxTransactionReport {
pub intent_count: usize,
pub durable_rows: usize,
pub local_report: Option<CollectedInvalidationReport>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct SqlxTransactionDiagnostics {
pub commits: u64,
pub rollbacks: u64,
pub body_errors: u64,
pub enqueue_failures: u64,
pub commit_failures: u64,
pub local_invalidations: u64,
}
#[derive(Debug, Default)]
struct SqlxTransactionCounters {
commits: AtomicU64,
rollbacks: AtomicU64,
body_errors: AtomicU64,
enqueue_failures: AtomicU64,
commit_failures: AtomicU64,
local_invalidations: AtomicU64,
}
impl SqlxTransactionCounters {
fn snapshot(&self) -> SqlxTransactionDiagnostics {
SqlxTransactionDiagnostics {
commits: self.commits.load(Ordering::Relaxed),
rollbacks: self.rollbacks.load(Ordering::Relaxed),
body_errors: self.body_errors.load(Ordering::Relaxed),
enqueue_failures: self.enqueue_failures.load(Ordering::Relaxed),
commit_failures: self.commit_failures.load(Ordering::Relaxed),
local_invalidations: self.local_invalidations.load(Ordering::Relaxed),
}
}
}
async fn rollback_and_count(tx: Transaction<'_, Sqlite>, counters: &SqlxTransactionCounters) {
if tx.rollback().await.is_ok() {
counters.rollbacks.fetch_add(1, Ordering::Relaxed);
}
}
fn sqlite_commit_position() -> CommitPosition {
static NEXT: AtomicU64 = AtomicU64::new(1);
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_nanos();
let sequence = NEXT.fetch_add(1, Ordering::Relaxed);
CommitPosition::new(format!("sqlite:{now}:{sequence}"))
}