txn_db/

txn.rs

//! Transactions and snapshots — the read and write handles a [`Db`](crate::Db)
//! hands out.
//!
//! A [`Transaction`] is the read-write unit of work. It takes a snapshot of the
//! database when it begins, serves every read from that snapshot (plus its own
//! uncommitted writes), buffers writes locally, and applies them atomically at
//! [`commit`](Transaction::commit) — or discards them on
//! [`rollback`](Transaction::rollback) or drop. Because reads come from a fixed
//! snapshot, a transaction never blocks writers and is never blocked by them.
//!
//! Transactions run under snapshot isolation by default. A serializable
//! transaction (from [`Db::begin_serializable`](crate::Db::begin_serializable),
//! behind the `serializable` feature) additionally records every key it reads so
//! that the read set can be validated at commit.
//!
//! Both a transaction and a [`Snapshot`] hold an [`ActiveReader`], which
//! registers the read timestamp with the database while the handle is alive and
//! unregisters it on drop. That registry is what lets garbage collection know
//! the oldest snapshot still in use, so it never reclaims a version a live reader
//! can still observe.

use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use crate::db::Inner;
use crate::error::Result;
use crate::store::{MemoryStore, VersionStore};
use crate::timestamp::Timestamp;

/// A handle's registration in the database's live-reader set.
///
/// Constructing one registers a read timestamp; dropping it unregisters that
/// timestamp. Holding it keeps the database from garbage-collecting versions the
/// reader can still see. It is a field of both [`Transaction`] and [`Snapshot`]
/// rather than a `Drop` on those types directly, so a transaction can still move
/// its write buffer out at commit.
struct ActiveReader<S: VersionStore> {
    inner: Arc<Inner<S>>,
    read_ts: Timestamp,
}

impl<S: VersionStore> ActiveReader<S> {
    /// Register a reader against `inner`, reading at the current watermark.
    fn new(inner: Arc<Inner<S>>) -> Self {
        let read_ts = inner.begin_reader();
        ActiveReader { inner, read_ts }
    }
}

impl<S: VersionStore> Drop for ActiveReader<S> {
    fn drop(&mut self) {
        self.inner.end_reader(self.read_ts);
    }
}

/// A read-write transaction over a consistent snapshot of the database.
///
/// A transaction is created by [`Db::begin`](crate::Db::begin) (snapshot
/// isolation) or `Db::begin_serializable` (serializable, with the `serializable`
/// feature). It reads as of the snapshot timestamp captured at that moment, so
/// concurrent commits by other transactions are invisible to it. Writes are
/// buffered in the transaction and become visible to others only when
/// [`commit`](Transaction::commit) succeeds; within the transaction, a read of a
/// key it has written returns that pending write (read-your-own-writes).
///
/// At commit the database checks every written key for a write-write conflict:
/// if another transaction committed a change to any of those keys after this
/// transaction's snapshot, the commit is rejected with a retryable
/// [`TxnError::Conflict`](crate::TxnError::Conflict) and none of the writes are
/// applied. A serializable transaction also validates its read set, rejecting
/// commits whose reads are no longer current.
///
/// Dropping a transaction without committing discards its buffered writes; it
/// is equivalent to [`rollback`](Transaction::rollback).
///
/// # Examples
///
/// ```
/// use txn_db::Db;
///
/// let db = Db::new();
///
/// let mut tx = db.begin();
/// tx.put(b"account:1".to_vec(), 100u64.to_le_bytes().to_vec());
/// tx.put(b"account:2".to_vec(), 50u64.to_le_bytes().to_vec());
/// let commit_ts = tx.commit()?;
///
/// // A fresh transaction sees the committed state.
/// let tx = db.begin();
/// assert!(tx.get(b"account:1")?.is_some());
/// assert!(commit_ts > txn_db::Timestamp::ZERO);
/// # Ok::<(), txn_db::TxnError>(())
/// ```
#[must_use = "a transaction buffers writes that are discarded unless it is committed"]
pub struct Transaction<S: VersionStore = MemoryStore> {
    active: ActiveReader<S>,
    writes: HashMap<Arc<[u8]>, Option<Arc<[u8]>>>,
    /// The set of keys read from the snapshot, tracked only for serializable
    /// transactions. `None` under snapshot isolation, where reads are not
    /// validated. Interior mutability lets [`get`](Self::get) record reads
    /// through a shared reference.
    reads: Option<RefCell<HashSet<Arc<[u8]>>>>,
}

impl<S: VersionStore> Transaction<S> {
    /// Construct a transaction over `inner`. When `serializable` is set the
    /// transaction records its read set for validation at commit.
    pub(crate) fn new(inner: Arc<Inner<S>>, serializable: bool) -> Self {
        Transaction {
            active: ActiveReader::new(inner),
            writes: HashMap::new(),
            reads: serializable.then(|| RefCell::new(HashSet::new())),
        }
    }

    /// The snapshot timestamp this transaction reads at.
    ///
    /// Every read that is not served from the transaction's own write buffer
    /// observes the database as of this timestamp.
    ///
    /// # Examples
    ///
    /// ```
    /// use txn_db::Db;
    ///
    /// let db = Db::new();
    /// let tx = db.begin();
    /// // Nothing has committed yet, so the snapshot is the empty database.
    /// assert_eq!(tx.read_timestamp(), txn_db::Timestamp::ZERO);
    /// ```
    #[inline]
    #[must_use]
    pub fn read_timestamp(&self) -> Timestamp {
        self.active.read_ts
    }

    /// Read the value of `key` as this transaction sees it.
    ///
    /// If the transaction has written `key`, the pending write is returned
    /// (read-your-own-writes), including `None` if it has deleted the key.
    /// Otherwise the value is read from the transaction's snapshot: the newest
    /// version committed at or before the snapshot timestamp, or `None` if the
    /// key does not exist as of the snapshot. For a serializable transaction the
    /// key is recorded in the read set for validation at commit.
    ///
    /// # Errors
    ///
    /// Returns [`TxnError::Store`](crate::TxnError::Store) if the backing
    /// [`VersionStore`](crate::VersionStore) fails the read. The default
    /// in-memory store never fails.
    ///
    /// # Examples
    ///
    /// ```
    /// use txn_db::Db;
    ///
    /// let db = Db::new();
    /// let mut tx = db.begin();
    ///
    /// assert_eq!(tx.get(b"k")?, None);          // absent
    /// tx.put(b"k".to_vec(), b"v".to_vec());
    /// assert_eq!(tx.get(b"k")?.as_deref(), Some(&b"v"[..])); // its own write
    /// tx.delete(b"k".to_vec());
    /// assert_eq!(tx.get(b"k")?, None);          // its own delete
    /// # Ok::<(), txn_db::TxnError>(())
    /// ```
    pub fn get(&self, key: &[u8]) -> Result<Option<Arc<[u8]>>> {
        if let Some(pending) = self.writes.get(key) {
            return Ok(pending.clone());
        }
        let value = self.active.inner.store.get(key, self.active.read_ts)?;
        // A serializable transaction records the key — present or absent — so a
        // later writer to it is caught at commit.
        if let Some(reads) = &self.reads {
            let _ = reads.borrow_mut().insert(Arc::from(key));
        }
        Ok(value)
    }

    /// Buffer a write of `value` to `key`, to be applied at commit.
    ///
    /// The write is local to this transaction until [`commit`](Self::commit)
    /// succeeds; other transactions do not see it. Writing the same key twice
    /// keeps the last value. Both arguments accept anything convertible into an
    /// `Arc<[u8]>` — passing an owned `Vec<u8>` or `Arc<[u8]>` moves it in
    /// without copying the bytes.
    ///
    /// # Examples
    ///
    /// ```
    /// use txn_db::Db;
    ///
    /// let db = Db::new();
    /// let mut tx = db.begin();
    /// tx.put(b"city".to_vec(), b"oslo".to_vec());
    /// tx.put(b"city".to_vec(), b"bergen".to_vec()); // overwrites within the txn
    /// assert_eq!(tx.get(b"city")?.as_deref(), Some(&b"bergen"[..]));
    /// # Ok::<(), txn_db::TxnError>(())
    /// ```
    pub fn put(&mut self, key: impl Into<Arc<[u8]>>, value: impl Into<Arc<[u8]>>) {
        let _ = self.writes.insert(key.into(), Some(value.into()));
    }

    /// Buffer a delete of `key`, to be applied at commit.
    ///
    /// After this call the transaction reads `key` as absent. At commit a
    /// tombstone is written so that snapshots taken after the commit also see
    /// the key as absent. Deleting a key that does not exist is a no-op that
    /// still participates in conflict detection, so a delete races other
    /// writers the same way a `put` does.
    ///
    /// # Examples
    ///
    /// ```
    /// use txn_db::Db;
    ///
    /// let db = Db::new();
    /// let mut setup = db.begin();
    /// setup.put(b"k".to_vec(), b"v".to_vec());
    /// setup.commit()?;
    ///
    /// let mut tx = db.begin();
    /// tx.delete(b"k".to_vec());
    /// tx.commit()?;
    ///
    /// assert_eq!(db.begin().get(b"k")?, None);
    /// # Ok::<(), txn_db::TxnError>(())
    /// ```
    pub fn delete(&mut self, key: impl Into<Arc<[u8]>>) {
        let _ = self.writes.insert(key.into(), None);
    }

    /// Commit the transaction, applying all buffered writes atomically.
    ///
    /// On success every buffered write becomes visible to transactions that
    /// begin afterward, and the commit timestamp is returned. A transaction
    /// that buffered no writes commits trivially and returns its snapshot
    /// timestamp without allocating a new one — including a serializable
    /// read-only transaction, which has observed a consistent snapshot and needs
    /// no validation.
    ///
    /// # Errors
    ///
    /// Returns [`TxnError::Conflict`](crate::TxnError::Conflict) — which is
    /// retryable — if any written key was changed by another transaction that
    /// committed after this one's snapshot, or, for a serializable transaction,
    /// if any key it read has since changed. In either case no writes are
    /// applied. Returns [`TxnError::Store`](crate::TxnError::Store) if the
    /// backing store fails to apply the batch, or
    /// [`TxnError::Durability`](crate::TxnError::Durability) if a durable commit
    /// cannot be made durable.
    ///
    /// # Examples
    ///
    /// ```
    /// use txn_db::Db;
    ///
    /// let db = Db::new();
    /// let mut tx = db.begin();
    /// tx.put(b"k".to_vec(), b"v".to_vec());
    /// let ts = tx.commit()?;
    /// assert!(ts > txn_db::Timestamp::ZERO);
    /// # Ok::<(), txn_db::TxnError>(())
    /// ```
    pub fn commit(self) -> Result<Timestamp> {
        let read_ts = self.active.read_ts;
        if self.writes.is_empty() {
            return Ok(read_ts);
        }
        // The read set, minus keys also in the write set (those are covered by
        // the write-write check). Empty for snapshot-isolation transactions.
        let reads: Vec<Arc<[u8]>> = match self.reads {
            Some(set) => set
                .into_inner()
                .into_iter()
                .filter(|key| !self.writes.contains_key(key))
                .collect(),
            None => Vec::new(),
        };
        let batch = self.writes.into_iter().collect();
        self.active.inner.commit_writes(read_ts, batch, &reads)
    }

    /// Discard the transaction and all of its buffered writes.
    ///
    /// This is explicit; simply dropping the transaction has the same effect.
    /// Rolling back never fails and never touches the shared store.
    ///
    /// # Examples
    ///
    /// ```
    /// use txn_db::Db;
    ///
    /// let db = Db::new();
    /// let mut tx = db.begin();
    /// tx.put(b"k".to_vec(), b"v".to_vec());
    /// tx.rollback();
    ///
    /// // The write never reached the database.
    /// assert_eq!(db.begin().get(b"k")?, None);
    /// # Ok::<(), txn_db::TxnError>(())
    /// ```
    #[inline]
    pub fn rollback(self) {
        // Dropping `self` releases the buffered writes and unregisters the
        // reader; this method documents the intent and consumes the transaction
        // so it cannot be used again.
    }
}

/// A read-only, point-in-time view of the database.
///
/// A snapshot is created by [`Db::snapshot`](crate::Db::snapshot) and reads as
/// of the moment it was taken. It has no write buffer and nothing to commit, so
/// it is cheaper than a transaction when all you need is to read several keys at
/// one consistent instant. While it is alive it pins that instant: garbage
/// collection will not reclaim a version the snapshot can still observe.
///
/// # Examples
///
/// ```
/// use txn_db::Db;
///
/// let db = Db::new();
/// let mut tx = db.begin();
/// tx.put(b"k".to_vec(), b"v1".to_vec());
/// tx.commit()?;
///
/// // Capture a snapshot, then change the database.
/// let snap = db.snapshot();
/// let mut tx = db.begin();
/// tx.put(b"k".to_vec(), b"v2".to_vec());
/// tx.commit()?;
///
/// // The snapshot still sees the value as of when it was taken.
/// assert_eq!(snap.get(b"k")?.as_deref(), Some(&b"v1"[..]));
/// assert_eq!(db.snapshot().get(b"k")?.as_deref(), Some(&b"v2"[..]));
/// # Ok::<(), txn_db::TxnError>(())
/// ```
pub struct Snapshot<S: VersionStore = MemoryStore> {
    active: ActiveReader<S>,
}

impl<S: VersionStore> Snapshot<S> {
    /// Construct a snapshot over `inner`, reading at the current watermark.
    pub(crate) fn new(inner: Arc<Inner<S>>) -> Self {
        Snapshot {
            active: ActiveReader::new(inner),
        }
    }

    /// The timestamp this snapshot reads at.
    ///
    /// # Examples
    ///
    /// ```
    /// use txn_db::Db;
    ///
    /// let db = Db::new();
    /// assert_eq!(db.snapshot().read_timestamp(), txn_db::Timestamp::ZERO);
    /// ```
    #[inline]
    #[must_use]
    pub fn read_timestamp(&self) -> Timestamp {
        self.active.read_ts
    }

    /// Read the value of `key` as of this snapshot.
    ///
    /// Returns the newest version committed at or before the snapshot
    /// timestamp, or `None` if the key does not exist as of that instant.
    ///
    /// # Errors
    ///
    /// Returns [`TxnError::Store`](crate::TxnError::Store) if the backing store
    /// fails the read. The default in-memory store never fails.
    ///
    /// # Examples
    ///
    /// ```
    /// use txn_db::Db;
    ///
    /// let db = Db::new();
    /// assert_eq!(db.snapshot().get(b"missing")?, None);
    /// # Ok::<(), txn_db::TxnError>(())
    /// ```
    pub fn get(&self, key: &[u8]) -> Result<Option<Arc<[u8]>>> {
        self.active.inner.store.get(key, self.active.read_ts)
    }
}