use crate::{Database, Transaction};
use anyhow::{Result, anyhow, bail};
use serde::{Serialize, de::DeserializeOwned};
use std::{
any::type_name,
cell::UnsafeCell,
collections::HashMap,
marker::PhantomData,
ops::{Deref, DerefMut},
sync::{
Arc, Mutex, Weak,
atomic::{AtomicBool, Ordering},
},
};
use tracing::Instrument;
#[repr(transparent)]
pub(crate) struct SyncUnsafeCell<T>(UnsafeCell<T>);
unsafe impl<T: Send> Send for SyncUnsafeCell<T> {}
unsafe impl<T: Send> Sync for SyncUnsafeCell<T> {}
impl<T> SyncUnsafeCell<T> {
fn new(value: T) -> Self {
Self(UnsafeCell::new(value))
}
fn get(&self) -> *mut T {
self.0.get()
}
}
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct DocKey {
pub pk: String,
pub sk: String,
}
impl DocKey {
pub fn new(pk: impl Into<String>, sk: impl Into<String>) -> Self {
Self {
pk: pk.into(),
sk: sk.into(),
}
}
}
pub trait Document: Serialize + DeserializeOwned + Send + Sync + 'static {
fn key(&self) -> DocKey;
}
pub trait DocGet {
type Doc: Document;
fn key(&self) -> DocKey;
}
#[allow(async_fn_in_trait)]
pub trait TrxRead: Sized {
type Output;
fn collect_keys(&self, keys: &mut Vec<DocKey>);
async fn finalize(
self,
tx: &Trx,
results: &mut std::vec::IntoIter<Option<crate::turso::StoredDoc>>,
) -> Result<Self::Output>;
}
impl<R> TrxRead for R
where
R: DocGet,
{
type Output = Option<DocHandle<R::Doc>>;
fn collect_keys(&self, keys: &mut Vec<DocKey>) {
keys.push(self.key());
}
async fn finalize(
self,
tx: &Trx,
results: &mut std::vec::IntoIter<Option<crate::turso::StoredDoc>>,
) -> Result<Self::Output> {
let stored = results
.next()
.ok_or_else(|| anyhow!("trx batch result missing for read"))?;
let key = self.key();
tx.inner
.lock()
.unwrap()
.register_loaded::<R::Doc>(key, stored)
}
}
macro_rules! impl_trx_read_tuple {
($($T:ident),+) => {
#[allow(non_snake_case)]
impl<$($T: TrxRead),+> TrxRead for ($($T,)+) {
type Output = ($($T::Output,)+);
fn collect_keys(&self, keys: &mut Vec<DocKey>) {
let ($($T,)+) = self;
$($T.collect_keys(keys);)+
}
async fn finalize(
self,
tx: &Trx,
results: &mut std::vec::IntoIter<Option<crate::turso::StoredDoc>>,
) -> Result<Self::Output> {
let ($($T,)+) = self;
Ok(($($T.finalize(tx, results).await?,)+))
}
}
};
}
impl_trx_read_tuple!(A);
impl_trx_read_tuple!(A, B);
impl_trx_read_tuple!(A, B, C);
impl_trx_read_tuple!(A, B, C, D);
impl_trx_read_tuple!(A, B, C, D, E);
impl_trx_read_tuple!(A, B, C, D, E, F);
impl_trx_read_tuple!(A, B, C, D, E, F, G);
impl_trx_read_tuple!(A, B, C, D, E, F, G, H);
impl_trx_read_tuple!(A, B, C, D, E, F, G, H, I);
impl_trx_read_tuple!(A, B, C, D, E, F, G, H, I, J);
impl_trx_read_tuple!(A, B, C, D, E, F, G, H, I, J, K);
impl_trx_read_tuple!(A, B, C, D, E, F, G, H, I, J, K, L);
pub struct DocHandle<T> {
data: Arc<SyncUnsafeCell<T>>,
dirty: Arc<AtomicBool>,
deleted: Arc<AtomicBool>,
_alive: Arc<()>,
_marker: PhantomData<Arc<T>>,
}
impl<T> DocHandle<T> {
pub fn delete(&self) {
self.deleted.store(true, Ordering::Release);
}
}
impl<T> Deref for DocHandle<T> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { &*self.data.get() }
}
}
impl<T> DerefMut for DocHandle<T> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.dirty.store(true, Ordering::Release);
unsafe { &mut *self.data.get() }
}
}
pub struct Trx {
inner: Arc<Mutex<TrxState>>,
}
impl Trx {
#[tracing::instrument(skip_all)]
pub async fn get<R>(&self, request: R) -> Result<R::Output>
where
R: TrxRead,
{
let mut keys = Vec::new();
request.collect_keys(&mut keys);
{
let state = self.inner.lock().unwrap();
for key in &keys {
if state.index.contains_key(key) {
bail!("duplicate trx key access: {}/{}", key.pk, key.sk);
}
}
}
let key_pairs: Vec<(String, String)> =
keys.iter().map(|k| (k.pk.clone(), k.sk.clone())).collect();
let stored = self.batch_load(&key_pairs).await?;
let mut iter = stored.into_iter();
request.finalize(self, &mut iter).await
}
pub fn create<T>(&self, doc: T) -> Result<DocHandle<T>>
where
T: Document,
{
self.inner.lock().unwrap().create(doc)
}
pub fn commit<Out, Cancel>(self, out: Out) -> Result<TrxControl<Out, Cancel>> {
Ok(TrxControl {
inner: TrxControlInner::Commit(out),
})
}
pub fn cancel<Out, Cancel>(self, reason: Cancel) -> Result<TrxControl<Out, Cancel>> {
Ok(TrxControl {
inner: TrxControlInner::Cancel(reason),
})
}
async fn batch_load(
&self,
keys: &[(String, String)],
) -> Result<Vec<Option<crate::turso::StoredDoc>>> {
if keys.is_empty() {
return Ok(vec![]);
}
let mut tx_opt = self.inner.lock().unwrap().tx.take();
let result = match &mut tx_opt {
Some(tx) => tx.batch_get_with_version(keys).await,
None => {
let db = self.inner.lock().unwrap().db.clone();
match db.begin_immediate_with_reads(keys).await {
Ok((tx, docs)) => {
tx_opt = Some(tx);
Ok(docs)
}
Err(e) => Err(e),
}
}
};
self.inner.lock().unwrap().tx = tx_opt;
result
}
}
pub struct TrxControl<Out, Cancel> {
inner: TrxControlInner<Out, Cancel>,
}
enum TrxControlInner<Out, Cancel> {
Commit(Out),
Cancel(Cancel),
}
#[derive(Debug)]
pub enum TrxResult<Out, Cancel, Err> {
Committed(Out),
Cancelled(Cancel),
Conflict(ConflictDetails),
Err(Err),
}
#[derive(Clone, Debug, Default)]
pub struct ConflictDetails {
pub keys: Vec<ConflictKey>,
}
#[derive(Clone, Debug)]
pub struct ConflictKey {
pub key: DocKey,
pub expected_version: Option<i64>,
pub actual_version: Option<i64>,
}
const MAX_ATTEMPTS: u32 = 5;
const BACKOFF_BASE_MS: u64 = 50;
const BACKOFF_CAP_MS: u64 = 1000;
pub(crate) async fn run<F, Fut, Out, Cancel, E>(db: Database, mut f: F) -> TrxResult<Out, Cancel, E>
where
F: FnMut(Trx) -> Fut,
Fut: std::future::Future<Output = Result<TrxControl<Out, Cancel>, E>>,
E: From<anyhow::Error>,
{
let mut attempt: u32 = 0;
loop {
let attempt_span = tracing::info_span!("trx_attempt", attempt = attempt);
let result = run_attempt(&db, &mut f).instrument(attempt_span).await;
match result {
AttemptOutcome::Done(r) => return r,
AttemptOutcome::Conflict(details) => {
if attempt + 1 >= MAX_ATTEMPTS {
return TrxResult::Conflict(details);
}
let backoff_span = tracing::info_span!("trx_backoff", attempt = attempt);
async {
let backoff = conflict_backoff(attempt).await;
crate::runtime::sleep(backoff).await;
}
.instrument(backoff_span)
.await;
attempt += 1;
}
}
}
}
enum AttemptOutcome<Out, Cancel, E> {
Done(TrxResult<Out, Cancel, E>),
Conflict(ConflictDetails),
}
async fn run_attempt<F, Fut, Out, Cancel, E>(
db: &Database,
f: &mut F,
) -> AttemptOutcome<Out, Cancel, E>
where
F: FnMut(Trx) -> Fut,
Fut: std::future::Future<Output = Result<TrxControl<Out, Cancel>, E>>,
E: From<anyhow::Error>,
{
let state = Arc::new(Mutex::new(TrxState::new(db.clone())));
let tx = Trx {
inner: state.clone(),
};
let user_span = tracing::info_span!("trx_user_closure");
let control = match f(tx).instrument(user_span).await {
Ok(control) => control,
Err(err) => return AttemptOutcome::Done(TrxResult::Err(err)),
};
match control.inner {
TrxControlInner::Commit(out) => {
let (commit_db, tx, entries_result) = take_entries_and_tx(state);
let entries = match entries_result {
Ok(e) => e,
Err(err) => {
if let Some(tx) = tx {
let _ = tx.rollback().await;
}
return AttemptOutcome::Done(TrxResult::Err(E::from(err)));
}
};
match commit_entries(commit_db, tx, entries).await {
Ok(()) => AttemptOutcome::Done(TrxResult::Committed(out)),
Err(CommitFailure::Conflict(details)) => AttemptOutcome::Conflict(details),
Err(CommitFailure::Err(err)) => AttemptOutcome::Done(TrxResult::Err(E::from(err))),
}
}
TrxControlInner::Cancel(reason) => {
let tx_to_rollback = state.lock().unwrap().tx.take();
if let Some(tx) = tx_to_rollback {
let _ = tx.rollback().await;
}
AttemptOutcome::Done(TrxResult::Cancelled(reason))
}
}
}
async fn conflict_backoff(attempt: u32) -> std::time::Duration {
let ceiling = BACKOFF_BASE_MS
.checked_shl(attempt)
.unwrap_or(BACKOFF_CAP_MS)
.min(BACKOFF_CAP_MS);
let mut buf = [0u8; 8];
crate::runtime::random_bytes(&mut buf).await;
let raw = u64::from_le_bytes(buf);
let delay_ms = raw % (ceiling + 1);
std::time::Duration::from_millis(delay_ms)
}
struct TrxState {
db: Database,
entries: Vec<TrackedEntry>,
index: HashMap<DocKey, usize>,
tx: Option<Transaction>,
}
impl TrxState {
fn new(db: Database) -> Self {
Self {
db,
entries: Vec::new(),
index: HashMap::new(),
tx: None,
}
}
fn register_loaded<T>(
&mut self,
key: DocKey,
stored: Option<crate::turso::StoredDoc>,
) -> Result<Option<DocHandle<T>>>
where
T: Document,
{
if self.index.contains_key(&key) {
bail!("duplicate trx key access: {}/{}", key.pk, key.sk);
}
let idx = self.entries.len();
self.index.insert(key.clone(), idx);
match stored {
Some(stored) => {
let doc = serde_json::from_slice::<T>(&stored.data).map_err(|err| {
anyhow!(
"failed to deserialize {} at {}/{}: {}",
type_name::<T>(),
key.pk,
key.sk,
err
)
})?;
let (shared, handle) = new_shared_doc(doc);
self.entries.push(TrackedEntry {
key,
expected_version: Some(stored.version),
state: TrackedState::Managed {
shared,
created: false,
},
});
Ok(Some(handle))
}
None => {
self.entries.push(TrackedEntry {
key,
expected_version: None,
state: TrackedState::Missing,
});
Ok(None)
}
}
}
fn create<T>(&mut self, doc: T) -> Result<DocHandle<T>>
where
T: Document,
{
let key = doc.key();
let (shared, handle) = new_shared_doc(doc);
match self.index.get(&key).copied() {
None => {
let idx = self.entries.len();
self.index.insert(key.clone(), idx);
self.entries.push(TrackedEntry {
key,
expected_version: None,
state: TrackedState::Managed {
shared,
created: true,
},
});
Ok(handle)
}
Some(idx) => match self.entries.get_mut(idx) {
Some(TrackedEntry {
expected_version: None,
state: TrackedState::Missing,
..
}) => {
self.entries[idx].state = TrackedState::Managed {
shared,
created: true,
};
Ok(handle)
}
_ => bail!("duplicate trx key access: {}/{}", key.pk, key.sk),
},
}
}
fn take_entries_and_tx(
&mut self,
) -> (Database, Option<Transaction>, Result<Vec<TrackedEntry>>) {
let tx = self.tx.take();
let db = self.db.clone();
for entry in &self.entries {
if let TrackedState::Managed { shared, .. } = &entry.state
&& shared.handle_alive.upgrade().is_some()
{
let err = anyhow!(
"live doc handle escaped trx for key {}/{}; commit outputs must not contain DocHandle values",
entry.key.pk,
entry.key.sk
);
self.entries.clear();
return (db, tx, Err(err));
}
}
(db, tx, Ok(std::mem::take(&mut self.entries)))
}
}
struct TrackedEntry {
key: DocKey,
expected_version: Option<i64>,
state: TrackedState,
}
impl TrackedEntry {
fn write(&self) -> Result<PendingWrite> {
match &self.state {
TrackedState::Missing => Ok(PendingWrite::None),
TrackedState::Managed { shared, created } => {
if shared.deleted.load(Ordering::Acquire) {
if *created {
return Ok(PendingWrite::None);
}
let expected_version = self.expected_version.ok_or_else(|| {
anyhow!("existing tracked doc missing expected version for delete")
})?;
return Ok(PendingWrite::Delete(expected_version));
}
if *created {
return Ok(PendingWrite::Insert((shared.serialize)()?));
}
if shared.dirty.load(Ordering::Acquire) {
let expected_version = self.expected_version.ok_or_else(|| {
anyhow!("existing tracked doc missing expected version for update")
})?;
return Ok(PendingWrite::Update {
expected_version,
data: (shared.serialize)()?,
});
}
Ok(PendingWrite::None)
}
}
}
}
enum TrackedState {
Missing,
Managed { shared: SharedDoc, created: bool },
}
struct SharedDoc {
dirty: Arc<AtomicBool>,
deleted: Arc<AtomicBool>,
handle_alive: Weak<()>,
serialize: Box<dyn Fn() -> Result<Vec<u8>> + Send + Sync>,
}
fn new_shared_doc<T>(doc: T) -> (SharedDoc, DocHandle<T>)
where
T: Document + Send + Sync,
{
let data = Arc::new(SyncUnsafeCell::new(doc));
let dirty = Arc::new(AtomicBool::new(false));
let deleted = Arc::new(AtomicBool::new(false));
let alive = Arc::new(());
let serialize_data = data.clone();
let shared = SharedDoc {
dirty: dirty.clone(),
deleted: deleted.clone(),
handle_alive: Arc::downgrade(&alive),
serialize: Box::new(move || {
let doc_ref = unsafe { &*serialize_data.get() };
serde_json::to_vec(doc_ref).map_err(Into::into)
}),
};
let handle = DocHandle {
data,
dirty,
deleted,
_alive: alive,
_marker: PhantomData,
};
(shared, handle)
}
fn take_entries_and_tx(
state: Arc<Mutex<TrxState>>,
) -> (Database, Option<Transaction>, Result<Vec<TrackedEntry>>) {
state.lock().unwrap().take_entries_and_tx()
}
enum PendingWrite {
None,
Insert(Vec<u8>),
Update {
expected_version: i64,
data: Vec<u8>,
},
Delete(i64),
}
enum CommitFailure {
Conflict(ConflictDetails),
Err(anyhow::Error),
}
#[tracing::instrument(skip_all, fields(entries = entries.len(), reused_tx = tx.is_some()))]
async fn commit_entries(
db: Database,
tx: Option<Transaction>,
entries: Vec<TrackedEntry>,
) -> std::result::Result<(), CommitFailure> {
let mut writes: Vec<crate::WriteOp> = Vec::new();
for entry in &entries {
match entry.write().map_err(CommitFailure::Err)? {
PendingWrite::None => {}
PendingWrite::Insert(data) => writes.push(crate::WriteOp::Insert {
pk: entry.key.pk.clone(),
sk: entry.key.sk.clone(),
data,
}),
PendingWrite::Update {
expected_version,
data,
} => writes.push(crate::WriteOp::Update {
pk: entry.key.pk.clone(),
sk: entry.key.sk.clone(),
expected_version,
data,
}),
PendingWrite::Delete(expected_version) => writes.push(crate::WriteOp::Delete {
pk: entry.key.pk.clone(),
sk: entry.key.sk.clone(),
expected_version,
}),
}
}
if writes.is_empty() && tx.is_none() {
return Ok(());
}
let mut tx = match tx {
Some(t) => t,
None => {
let begin_span = tracing::info_span!("commit_begin_tx");
async { db.transaction().await.map_err(CommitFailure::Err) }
.instrument(begin_span)
.await?
}
};
let outcome = tx
.apply_writes_and_commit(&writes)
.await
.map_err(CommitFailure::Err)?;
let mut conflicts = Vec::new();
if let Some(info) = outcome.conflict
&& let Some(op) = writes.get(info.step_index)
{
let (pk, sk, expected) = write_key_and_expected(op);
conflicts.push(ConflictKey {
key: DocKey { pk, sk },
expected_version: expected,
actual_version: None,
});
}
for (i, count) in outcome.affected_counts.iter().enumerate() {
if *count == 1 {
continue;
}
let Some(op) = writes.get(i) else { continue };
let (pk, sk, expected) = write_key_and_expected(op);
let key = DocKey { pk, sk };
if conflicts.iter().any(|c| c.key == key) {
continue;
}
conflicts.push(ConflictKey {
key,
expected_version: expected,
actual_version: None,
});
}
if !conflicts.is_empty() {
let key_pairs: Vec<(String, String)> = conflicts
.iter()
.map(|c| (c.key.pk.clone(), c.key.sk.clone()))
.collect();
if let Ok(stored) = db.batch_get_with_version(&key_pairs).await {
for (c, slot) in conflicts.iter_mut().zip(stored.into_iter()) {
c.actual_version = slot.map(|d| d.version);
}
}
return Err(CommitFailure::Conflict(ConflictDetails { keys: conflicts }));
}
Ok(())
}
fn write_key_and_expected(op: &crate::WriteOp) -> (String, String, Option<i64>) {
match op {
crate::WriteOp::Insert { pk, sk, .. } => (pk.clone(), sk.clone(), None),
crate::WriteOp::Update {
pk,
sk,
expected_version,
..
}
| crate::WriteOp::Delete {
pk,
sk,
expected_version,
} => (pk.clone(), sk.clone(), Some(*expected_version)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::turso_with_config;
#[derive(Clone, serde::Serialize, serde::Deserialize)]
struct TestDoc {
id: String,
value: i32,
}
impl Document for TestDoc {
fn key(&self) -> DocKey {
DocKey::new("TestDoc", format!("id={}", self.id))
}
}
struct TestDocGet {
id: String,
}
impl DocGet for TestDocGet {
type Doc = TestDoc;
fn key(&self) -> DocKey {
DocKey::new("TestDoc", format!("id={}", self.id))
}
}
fn test_state() -> TrxState {
TrxState::new(turso_with_config(
"http://127.0.0.1:0".to_string(),
String::new(),
))
}
#[test]
fn create_produces_insert_even_without_mutation() {
let mut state = test_state();
let _handle = state
.create(TestDoc {
id: "a".into(),
value: 1,
})
.expect("create should succeed");
let write = state.entries[0].write().expect("write plan");
match write {
PendingWrite::Insert(data) => {
let doc: TestDoc = serde_json::from_slice(&data).expect("deserialize insert");
assert_eq!(doc.id, "a");
assert_eq!(doc.value, 1);
}
_ => panic!("expected insert"),
}
}
#[test]
fn loaded_doc_marks_dirty_on_deref_mut() {
let mut state = test_state();
let doc = TestDoc {
id: "a".into(),
value: 1,
};
let key = TestDocGet { id: "a".into() }.key();
let mut handle = state
.register_loaded::<TestDoc>(
key,
Some(crate::turso::StoredDoc {
data: serde_json::to_vec(&doc).expect("serialize").into(),
version: 7,
}),
)
.expect("load should succeed")
.expect("doc should exist");
handle.value = 5;
match state.entries[0].write().expect("write plan") {
PendingWrite::Update {
expected_version,
data,
} => {
assert_eq!(expected_version, 7);
let doc: TestDoc = serde_json::from_slice(&data).expect("deserialize update");
assert_eq!(doc.value, 5);
}
_ => panic!("expected update"),
}
}
#[test]
fn loaded_doc_delete_produces_delete_write() {
let mut state = test_state();
let doc = TestDoc {
id: "a".into(),
value: 1,
};
let key = TestDocGet { id: "a".into() }.key();
let handle = state
.register_loaded::<TestDoc>(
key,
Some(crate::turso::StoredDoc {
data: serde_json::to_vec(&doc).expect("serialize").into(),
version: 7,
}),
)
.expect("load should succeed")
.expect("doc should exist");
handle.delete();
match state.entries[0].write().expect("write plan") {
PendingWrite::Delete(expected_version) => assert_eq!(expected_version, 7),
_ => panic!("expected delete"),
}
}
#[test]
fn missing_read_can_be_promoted_to_create() {
let mut state = test_state();
let key = TestDocGet { id: "a".into() }.key();
let loaded = state
.register_loaded::<TestDoc>(key, None)
.expect("register missing should succeed");
assert!(loaded.is_none());
let handle = state
.create(TestDoc {
id: "a".into(),
value: 3,
})
.expect("create after missing get should succeed");
assert_eq!(handle.value, 3);
assert!(matches!(
state.entries[0].write().expect("write plan"),
PendingWrite::Insert(_)
));
}
#[test]
fn duplicate_key_access_is_rejected() {
let mut state = test_state();
let first = state.register_loaded::<TestDoc>(TestDocGet { id: "a".into() }.key(), None);
assert!(first.is_ok());
let second = state.register_loaded::<TestDoc>(TestDocGet { id: "a".into() }.key(), None);
assert!(second.is_err());
}
#[test]
fn live_handle_cannot_escape_commit_boundary() {
let mut state = test_state();
let _handle = state
.create(TestDoc {
id: "a".into(),
value: 1,
})
.expect("create should succeed");
let (_, _, result) = state.take_entries_and_tx();
match result {
Ok(_) => panic!("live handle should fail"),
Err(err) => assert!(err.to_string().contains("live doc handle escaped trx")),
}
}
}