reddb_engine/transaction/

snapshot.rs

//! Live optimistic MVCC snapshot, outcome and SSI tracking.
//!
//! This in-process state allocates transaction IDs, captures transaction and
//! statement snapshots, retains vacuum horizons, tracks aborted writers and
//! detects dangerous serializable read/write structures. The host runtime owns
//! write conflict validation, WAL publication and recovery of transaction outcomes.
//! Moving this state into the shared engine does not itself provide durability.

use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use super::IsolationLevel;
use crate::storage::entity_id::EntityId;

/// Default autocommit-xid pool batch size. Each refill reserves this
/// many xids in a single `fetch_add` so back-to-back autocommit inserts
/// share one atomic op instead of paying it per row. Sized small to
/// keep a pristine `peek_next_xid()` close to the truth — VACUUM and
/// diagnostics observe the allocator high-water mark. Snapshots exclude
/// reserved-but-unused slots from visibility until a later snapshot.
pub(crate) const AUTOCOMMIT_POOL_BATCH: u64 = 16;

/// A transaction identifier. Monotonic across the lifetime of the process.
pub type Xid = u64;

/// Reserved xid meaning "not inside a transaction" — pre-MVCC rows stamp
/// this value so they stay visible to every snapshot.
pub const XID_NONE: Xid = 0;

/// Immutable snapshot taken at transaction start or statement start.
///
/// Callers evaluate `UnifiedEntity::is_visible(snapshot.xid)` on every
/// row returned from storage to filter out rows created by concurrent
/// transactions that hadn't committed when the snapshot was taken.
#[derive(Debug, Clone)]
pub struct Snapshot {
    /// The snapshot's xid — every row with `xmin <= xid` created before
    /// the snapshot is visible (assuming `xmax` hasn't passed).
    pub xid: Xid,
    /// Active transactions and reserved-but-unissued pool xids at capture.
    /// Their writes must be *hidden* even when `xmin <= xid`, because
    /// the writer hadn't committed yet from this snapshot's point of view.
    pub in_progress: HashSet<Xid>,
}

impl Snapshot {
    /// Is a row with this xmin/xmax visible under this snapshot?
    ///
    /// Delegates to [`super::visibility::is_visible`] — the deep
    /// module that owns the full MVCC visibility predicate. The
    /// `aborted` argument is empty here because `Snapshot` does not
    /// carry the manager-level aborted set; callers that need the
    /// rolled-back-writer rule should consult [`SnapshotManager`]
    /// directly, or evolve `Snapshot` to embed an aborted view.
    pub fn sees(&self, xmin: Xid, xmax: Xid) -> bool {
        super::visibility::is_visible(xmin, xmax, self.xid, &self.in_progress, &HashSet::new())
    }
}

/// Per-transaction state tracked on the runtime while BEGIN/COMMIT/ROLLBACK
/// is active. Attached to a connection by the runtime Transaction State module.
#[derive(Debug, Clone)]
pub struct TxnContext {
    pub xid: Xid,
    pub isolation: IsolationLevel,
    /// Snapshot captured at BEGIN (SnapshotIsolation / Serializable) or
    /// refreshed per-statement (ReadCommitted).
    pub snapshot: Snapshot,
    /// Ordered list of `(savepoint_name, sub_xid)` entries (Phase
    /// 2.3.2e savepoints). Each SAVEPOINT pushes a freshly-allocated
    /// xid onto this stack; writes stamp xmin/xmax with the top entry
    /// so ROLLBACK TO SAVEPOINT can mark only those writes as aborted.
    /// RELEASE SAVEPOINT pops the named level plus everything above it
    /// without aborting — the sub-xids keep their effects and commit
    /// together with the parent. Empty stack means "writes use `xid`
    /// directly", matching pre-savepoint behaviour.
    pub savepoints: Vec<(String, Xid)>,
    /// Sub-xids popped by `RELEASE SAVEPOINT` that should still commit
    /// alongside the parent. PG semantics: released subtxns keep their
    /// writes — they're promoted to parent-visible at COMMIT. Stored
    /// separately from `savepoints` so their names are gone (cannot be
    /// rolled back or released again) while their xids remain trackable.
    pub released_sub_xids: Vec<Xid>,
}

impl TxnContext {
    /// Xid new writes in this connection should stamp onto tuples — the
    /// innermost open savepoint, or the parent xid when no savepoint is
    /// active.
    pub fn writer_xid(&self) -> Xid {
        self.savepoints.last().map(|(_, x)| *x).unwrap_or(self.xid)
    }
}

/// Central allocator and liveness tracker.
///
/// Uses an atomic counter for xid allocation and a parking_lot-guarded
/// HashSet for in-progress/aborted bookkeeping. The sets stay small —
/// only unfinished transactions plus a finite rollback history — so a
/// plain HashSet outperforms more complex data structures here.
///
/// # Autocommit xid pool
///
/// Single-row autocommit writes (`MutationEngine::append_one`) need an
/// xid that's "born committed" — they call `begin()` then `commit()`
/// back-to-back before the row is even durable. The pre-commit pool
/// (`autocommit_pool_*`) batches the reservation: one
/// `next_xid.fetch_add(BATCH)` reserves a contiguous range of xids,
/// each handed out via a single atomic without touching the
/// `RwLock<ManagerState>`. Pool xids are never inserted into `active`
/// or `aborted` so they look like already-committed transactions to
/// every snapshot — identical visibility semantics to the legacy
/// `begin()/commit()` pair (which also leaves the xid in neither set).
pub struct SnapshotManager {
    next_xid: AtomicU64,
    visibility_revision: AtomicU64,
    state: parking_lot::RwLock<ManagerState>,
    /// Reservation window for the autocommit pool. A single
    /// `parking_lot::Mutex` protects two `u64`s — `next` (next xid to
    /// hand out) and `end` (exclusive upper bound). When `next == end`
    /// the next caller refills by reserving `AUTOCOMMIT_POOL_BATCH`
    /// xids in a single `next_xid.fetch_add`, dropping the lock cost
    /// from one acquire-per-xid (the legacy `begin()`+`commit()` pair)
    /// to one acquire-per-`AUTOCOMMIT_POOL_BATCH` xids. A plain Mutex
    /// is enough here — the critical section is two stores and an
    /// atomic add, and contention is bounded by the writer count.
    autocommit_pool: parking_lot::Mutex<AutocommitPool>,
}

pub type LogicalRowKey = (String, EntityId);

#[derive(Default)]
struct SerializableTxnState {
    read_set: HashSet<LogicalRowKey>,
    vault_reads: HashSet<(String, Option<String>, Option<String>)>,
    incoming_rw: HashSet<Xid>,
    outgoing_rw: HashSet<Xid>,
    committed: bool,
}

#[derive(Default)]
struct AutocommitPool {
    next: Xid,
    end: Xid,
}

#[derive(Default)]
struct ManagerState {
    /// xids that have started but not yet committed/rolled back.
    active: HashSet<Xid>,
    /// Oldest version required by each active transaction's captured snapshot.
    /// A newer reader can exclude an older writer even after that writer commits.
    snapshot_horizons: HashMap<Xid, Xid>,
    /// Captured statement frames also need history without an active SQL
    /// transaction. Counts let nested frames retain the same horizon safely.
    statement_horizons: HashMap<Xid, usize>,
    /// xids that rolled back. `is_visible` MUST treat these as invisible
    /// (the writer never committed). The set is pruned lazily by VACUUM.
    aborted: HashSet<Xid>,
    /// A low-level rollback of a non-active xid may invalidate already stored
    /// rows without a new row mutation. Checkpoint caches must validate fully.
    untracked_visibility_revision: u64,
    /// xids that must NOT be reclaimed by VACUUM because some higher-level
    /// object (a VCS commit, a long-lived replica snapshot) still points
    /// at them. Reference-counted so multiple pins coexist; decrementing
    /// to zero removes the entry. `prune_aborted` skips any xid present
    /// here so its row versions stay readable.
    pinned: HashMap<Xid, u32>,
    serializable: HashMap<Xid, SerializableTxnState>,
}

/// Releases a statement's history requirement on every frame exit path.
pub struct SnapshotRetentionGuard {
    manager: Arc<SnapshotManager>,
    horizon: Xid,
}

impl Drop for SnapshotRetentionGuard {
    fn drop(&mut self) {
        let mut state = self.manager.state.write();
        let count = state
            .statement_horizons
            .get_mut(&self.horizon)
            .expect("a live snapshot retention guard has a registered horizon");
        if *count == 1 {
            state.statement_horizons.remove(&self.horizon);
        } else {
            *count -= 1;
        }
    }
}

impl SnapshotManager {
    pub fn new() -> Self {
        Self {
            // Start at 1 so xid=0 keeps its pre-MVCC "everyone sees it" meaning.
            next_xid: AtomicU64::new(1),
            visibility_revision: AtomicU64::new(0),
            state: parking_lot::RwLock::new(ManagerState::default()),
            // Pool starts empty — first caller triggers a refill.
            autocommit_pool: parking_lot::Mutex::new(AutocommitPool::default()),
        }
    }

    /// Allocate a new xid and mark it active. Returns the xid for
    /// stamping onto `UnifiedEntity::xmin/xmax`.
    ///
    /// Allocation and the `active` insert happen together under the state
    /// write lock. This is load-bearing for first-committer-wins: a
    /// concurrent snapshot (state read lock) must never observe a
    /// *lower-numbered* still-running xid as absent from `active`. Were the
    /// `fetch_add` done outside the lock, a writer that grabbed the smaller
    /// xid could still be mid-`begin()` while a later transaction snapshots
    /// a larger xid — the smaller xid would be missing from `in_progress`,
    /// yet `xmin <= snapshot.xid`, so the snapshot would treat that
    /// concurrent (later possibly-committed) writer as visible. That both
    /// leaks an uncommitted row into a peer's read and lets two writers on
    /// the same row both commit. Serialising allocation under the lock
    /// guarantees any xid smaller than a taken snapshot's xid is already in
    /// `active` (or already finished), so `in_progress` is never short a
    /// concurrent writer.
    pub fn begin(&self) -> Xid {
        let mut state = self.state.write();
        // A real writer advances beyond the reserved autocommit range. Retire
        // its unused slots so later inserts cannot receive older identities
        // and pin fresh statements' vacuum horizons behind committed deletes.
        let mut pool = self.autocommit_pool.lock();
        pool.next = pool.end;
        let xid = self.next_xid.fetch_add(1, Ordering::Relaxed);
        state.active.insert(xid);
        self.visibility_revision.fetch_add(1, Ordering::Release);
        xid
    }

    /// Capture a fresh *read-committed* snapshot at this instant: sees every
    /// transaction that has committed so far and hides every still-active
    /// one. Unlike [`snapshot`], the caller supplies no xid — the snapshot
    /// id is the current allocator high-water mark, read *under the same
    /// state lock* that [`begin`] holds while it bumps the counter and
    /// inserts into `active`. That mutual exclusion is what makes the read
    /// atomic: a concurrent `begin()` is either wholly visible (counter
    /// bumped AND present in `active`, hence hidden as in-progress) or
    /// wholly invisible (neither) — never the torn state where its xid sits
    /// below the high-water mark yet is missing from `active`, which would
    /// leak an uncommitted writer as visible.
    ///
    /// Used for RMW pre-image reads that must always observe the latest
    /// committed row version rather than a pinned statement snapshot.
    pub fn fresh_read_snapshot(&self) -> Snapshot {
        let state = self.state.read();
        self.fresh_read_snapshot_locked(&state)
    }

    fn fresh_read_snapshot_locked(&self, state: &ManagerState) -> Snapshot {
        let pool = self.autocommit_pool.lock();
        // Visibility uses an inclusive bound. The allocator points at the
        // next, not the last, xid; including it admits the next future writer.
        let xid = self.next_xid.load(Ordering::Relaxed).saturating_sub(1);
        let mut in_progress: HashSet<Xid> = state.active.iter().copied().collect();
        // Reserved-but-unissued pool slots are future writers even though
        // their numbers can be below this snapshot's high-water mark.
        in_progress.extend(pool.next..pool.end);
        Snapshot { xid, in_progress }
    }

    /// Capture and register under the same lock used by GC cutoff selection.
    /// Registering only when the snapshot is installed leaves a capture-to-read
    /// interval in which a committed writer's old versions can be reclaimed.
    pub fn fresh_retained_snapshot(self: &Arc<Self>) -> (Snapshot, SnapshotRetentionGuard) {
        let mut state = self.state.write();
        let snapshot = self.fresh_read_snapshot_locked(&state);
        let guard = self.retain_snapshot_locked(&snapshot, &mut state);
        (snapshot, guard)
    }

    /// Retain a snapshot already protected by its parent transaction, outer
    /// statement frame, or historical reference while transferring ownership.
    pub fn retain_snapshot(
        self: &Arc<Self>,
        snapshot: Snapshot,
    ) -> (Snapshot, SnapshotRetentionGuard) {
        let mut state = self.state.write();
        let guard = self.retain_snapshot_locked(&snapshot, &mut state);
        (snapshot, guard)
    }

    fn retain_snapshot_locked(
        self: &Arc<Self>,
        snapshot: &Snapshot,
        state: &mut ManagerState,
    ) -> SnapshotRetentionGuard {
        // Deletes already visible to the snapshot need no old version. An
        // excluded writer can be older than its inclusive high-water mark.
        let horizon = snapshot
            .in_progress
            .iter()
            .copied()
            .fold(snapshot.xid.saturating_add(1), Xid::min);
        *state.statement_horizons.entry(horizon).or_default() += 1;
        SnapshotRetentionGuard {
            manager: Arc::clone(self),
            horizon,
        }
    }

    /// Capture a point-in-time snapshot. Must be called after `begin()`
    /// when using SnapshotIsolation/Serializable. ReadCommitted refreshes
    /// this per statement via the same call.
    pub fn snapshot(&self, xid: Xid) -> Snapshot {
        let mut state = self.state.write();
        // Active xids other than our own appear as "in-progress" to us.
        let pool = self.autocommit_pool.lock();
        let mut in_progress: HashSet<Xid> =
            state.active.iter().copied().filter(|&x| x != xid).collect();
        in_progress.extend(pool.next..pool.end);
        if state.active.contains(&xid) {
            let horizon = in_progress.iter().copied().fold(xid, Xid::min);
            state
                .snapshot_horizons
                .entry(xid)
                .and_modify(|previous| *previous = (*previous).min(horizon))
                .or_insert(horizon);
        }
        Snapshot { xid, in_progress }
    }

    /// Mark a transaction as committed. Its writes become visible to
    /// future snapshots; earlier snapshots keep their own view.
    pub fn commit(&self, xid: Xid) {
        let mut state = self.state.write();
        let was_active = state.active.remove(&xid);
        state.snapshot_horizons.remove(&xid);
        // Also clear from aborted set in case of prior rollback_to call
        // that touched this xid (defensive; normally a no-op).
        let was_aborted = state.aborted.remove(&xid);
        if was_active || was_aborted {
            self.visibility_revision.fetch_add(1, Ordering::Release);
        }
        if let Some(ssi) = state.serializable.get_mut(&xid) {
            ssi.committed = true;
        }
        Self::prune_serializable_state(&mut state);
    }

    /// Allocate an xid that is *born committed* — for autocommit
    /// callers (`MutationEngine::append_one`) that previously paid two
    /// `state.write()` lock acquisitions per row to insert-then-remove
    /// from the active set.
    ///
    /// The returned xid is never inserted into `active` and never into
    /// `aborted`, which matches the steady state of the legacy
    /// `begin()/commit()` pair when called back-to-back: the xid leaves
    /// the manager's tracking sets unobservably. Concurrent readers
    /// therefore see it as an already-committed transaction once
    /// `xmin <= snapshot.xid`, which is exactly the semantics the
    /// autocommit path needs.
    ///
    /// Implementation: a small reservation pool (`AUTOCOMMIT_POOL_BATCH`
    /// xids) is reserved with one `fetch_add`. Each caller hands itself
    /// the next xid via a single atomic. When the pool drains, the
    /// next caller serialises briefly through `autocommit_pool_refill`
    /// to bump the window, then falls back into the lock-free hot path.
    ///
    /// Durability note: this method does NOT make the row durable —
    /// it only allocates the identifier. The caller must complete the
    /// usual WAL-append + fsync cycle before acknowledging the write.
    /// Pre-allocating the xid is safe because the xid carries no
    /// promise that any row exists; it's just a number for `xmin`.
    pub fn allocate_committed_xid(&self) -> Xid {
        let mut pool = self.autocommit_pool.lock();
        if pool.next >= pool.end {
            // Reserve the next contiguous range. A single
            // `fetch_add(BATCH)` on the global counter — equivalent to
            // BATCH back-to-back `begin()` calls in terms of xid
            // numbering, but with zero `state.write()` traffic.
            let start = self
                .next_xid
                .fetch_add(AUTOCOMMIT_POOL_BATCH, Ordering::Relaxed);
            pool.next = start;
            pool.end = start + AUTOCOMMIT_POOL_BATCH;
        }
        let xid = pool.next;
        pool.next += 1;
        xid
    }

    /// Mark a transaction as rolled back. Its writes MUST stay hidden
    /// from every future read — `is_visible` consults the aborted set
    /// before honouring a row's `xmin`.
    pub fn rollback(&self, xid: Xid) {
        let mut state = self.state.write();
        state.snapshot_horizons.remove(&xid);
        let was_active = state.active.remove(&xid);
        let newly_aborted = state.aborted.insert(xid);
        if newly_aborted && !was_active {
            state.untracked_visibility_revision = state
                .untracked_visibility_revision
                .checked_add(1)
                .expect("visibility invalidation counter must not overflow");
        }
        if was_active || newly_aborted {
            self.visibility_revision.fetch_add(1, Ordering::Release);
        }
        state.serializable.remove(&xid);
        for ssi in state.serializable.values_mut() {
            ssi.incoming_rw.remove(&xid);
            ssi.outgoing_rw.remove(&xid);
        }
        Self::prune_serializable_state(&mut state);
    }

    pub fn begin_serializable(&self, xid: Xid) {
        self.state.write().serializable.entry(xid).or_default();
    }

    pub fn record_serializable_read(&self, reader: Xid, collection: &str, logical_id: EntityId) {
        let mut state = self.state.write();
        if let Some(ssi) = state.serializable.get_mut(&reader) {
            ssi.read_set.insert((collection.to_string(), logical_id));
        }
    }

    pub fn record_serializable_vault_read(
        &self,
        reader: Xid,
        collection: &str,
        tenant: Option<&str>,
        key: Option<&str>,
    ) {
        if let Some(ssi) = self.state.write().serializable.get_mut(&reader) {
            ssi.vault_reads.insert((
                collection.into(),
                tenant.map(str::to_owned),
                key.map(str::to_owned),
            ));
        }
    }

    pub fn serializable_vault_commit_would_be_dangerous(
        &self,
        writer: Xid,
        write_set: &HashSet<(String, Option<String>, String)>,
    ) -> bool {
        let mut state = self.state.write();
        if !state.serializable.contains_key(&writer) {
            return false;
        }
        let readers = state
            .serializable
            .iter()
            .filter_map(|(&reader, ssi)| {
                (reader != writer
                    && ssi.vault_reads.iter().any(|(collection, tenant, key)| {
                        write_set
                            .iter()
                            .any(|(written_collection, written_tenant, written_key)| {
                                collection == written_collection
                                    && tenant == written_tenant
                                    && key.as_ref().is_none_or(|key| key == written_key)
                            })
                    }))
                .then_some(reader)
            })
            .collect();
        Self::serializable_readers_would_be_dangerous(&mut state, writer, readers)
    }

    pub fn serializable_commit_would_be_dangerous(
        &self,
        writer: Xid,
        write_set: &HashSet<LogicalRowKey>,
    ) -> bool {
        if write_set.is_empty() {
            return false;
        }

        let mut state = self.state.write();
        if !state.serializable.contains_key(&writer) {
            return false;
        }

        let readers: Vec<Xid> = state
            .serializable
            .iter()
            .filter_map(|(&reader, ssi)| {
                if reader != writer && !ssi.read_set.is_disjoint(write_set) {
                    Some(reader)
                } else {
                    None
                }
            })
            .collect();

        Self::serializable_readers_would_be_dangerous(&mut state, writer, readers)
    }

    fn serializable_readers_would_be_dangerous(
        state: &mut ManagerState,
        writer: Xid,
        readers: Vec<Xid>,
    ) -> bool {
        for reader in readers {
            if let Some(reader_state) = state.serializable.get_mut(&reader) {
                reader_state.outgoing_rw.insert(writer);
            }
            if let Some(writer_state) = state.serializable.get_mut(&writer) {
                writer_state.incoming_rw.insert(reader);
            }
        }

        let Some(writer_state) = state.serializable.get(&writer) else {
            return false;
        };
        let has_incoming = !writer_state.incoming_rw.is_empty();
        let has_outgoing_to_committed = writer_state
            .outgoing_rw
            .iter()
            .any(|xid| state.serializable.get(xid).is_some_and(|ssi| ssi.committed));
        has_incoming && has_outgoing_to_committed
    }

    fn prune_serializable_state(state: &mut ManagerState) {
        if state.active.is_empty() {
            state.serializable.clear();
        }
    }

    /// Is this xid known to have rolled back? Called by the read path to
    /// skip tuples whose creator never committed.
    pub fn is_aborted(&self, xid: Xid) -> bool {
        self.state.read().aborted.contains(&xid)
    }

    /// Unfinished writers must be rolled back in a crash-recoverable checkpoint,
    /// without changing their live transaction state or their readers.
    pub fn checkpoint_excluded_state(&self) -> (u64, HashSet<Xid>) {
        let (revision, excluded, _) = self.checkpoint_projection_state();
        (revision, excluded)
    }

    pub fn checkpoint_projection_state(&self) -> (u64, HashSet<Xid>, u64) {
        let state = self.state.read();
        let mut excluded = state.aborted.clone();
        excluded.extend(state.active.iter().copied());
        (
            self.visibility_revision.load(Ordering::Acquire),
            excluded,
            state.untracked_visibility_revision,
        )
    }

    /// Is this xid still active?
    pub fn is_active(&self, xid: Xid) -> bool {
        self.state.read().active.contains(&xid)
    }

    /// Unique-key admission uses current transaction outcomes, not the
    /// reader's snapshot. An invisible active/newly committed writer still
    /// reserves its key; an aborted creator or completed replacement does not.
    pub fn row_reserves_unique_key(&self, xmin: Xid, xmax: Xid, own_xids: &HashSet<Xid>) -> bool {
        let state = self.state.read();
        if state.aborted.contains(&xmin) {
            return false;
        }
        if state.active.contains(&xmin) && !own_xids.contains(&xmin) {
            return true;
        }
        if xmax == 0 || state.aborted.contains(&xmax) {
            return true;
        }
        !own_xids.contains(&xmax) && state.active.contains(&xmax)
    }

    /// Snapshot of every still-active xid (for VACUUM oldest-active-xid
    /// calculation — any row with `xmax < min(active)` is reclaimable).
    pub fn oldest_active_xid(&self) -> Option<Xid> {
        self.state.read().active.iter().copied().min()
    }

    /// GC must preserve writers excluded by an active reader's snapshot, even
    /// when those writers have since committed and left the active set.
    pub fn oldest_required_xid(&self) -> Option<Xid> {
        let state = self.state.read();
        state
            .active
            .iter()
            .chain(state.snapshot_horizons.values())
            .chain(state.statement_horizons.keys())
            .chain(state.pinned.keys())
            .copied()
            .min()
    }

    /// Oldest externally pinned xid. Pinned snapshots behave like active
    /// snapshots for VACUUM: any tuple visible to that xid must survive even
    /// when no SQL transaction is currently active.
    pub fn oldest_pinned_xid(&self) -> Option<Xid> {
        self.state.read().pinned.keys().copied().min()
    }

    /// Return the next xid that would be allocated. Useful for diagnostics
    /// and for VACUUM to know the upper bound of aborted-xid retention.
    pub fn peek_next_xid(&self) -> Xid {
        self.next_xid.load(Ordering::Relaxed)
    }

    /// Advance the allocator so future snapshots consider an xid
    /// recovered from storage/WAL to be in the committed past.
    pub fn observe_committed_xid(&self, xid: Xid) {
        if xid == XID_NONE {
            return;
        }
        // Replay can import an xid from our still-unused reservation range.
        // Retire those slots before publishing the floor: they must neither
        // hide the recovered row in fresh snapshots nor be allocated again.
        // Captured snapshots retain their original exclusions.
        let mut pool = self.autocommit_pool.lock();
        let target = xid.saturating_add(1);
        pool.next = pool.next.max(target.min(pool.end));
        let mut current = self.next_xid.load(Ordering::Relaxed);
        while current < target {
            match self.next_xid.compare_exchange(
                current,
                target,
                Ordering::Relaxed,
                Ordering::Relaxed,
            ) {
                Ok(_) => return,
                Err(actual) => current = actual,
            }
        }
    }

    /// Prune outcomes only after the caller has removed or normalized every
    /// physical version that references them across all collections and models.
    /// A snapshot cutoff alone is insufficient: an aborted insert with xmin=x
    /// becomes visible again if x is forgotten while that version still exists.
    /// Pinned xids remain retained for higher-level references.
    pub fn prune_aborted(&self, below: Xid) {
        let mut state = self.state.write();
        let ManagerState {
            aborted, pinned, ..
        } = &mut *state;
        let previous = aborted.len();
        aborted.retain(|&x| x >= below || pinned.contains_key(&x));
        if previous != aborted.len() {
            self.visibility_revision.fetch_add(1, Ordering::Release);
        }
    }

    /// Pin an xid so its row versions stay reclaim-safe across VACUUM.
    /// Reference-counted — call `unpin` once per `pin` to release.
    pub fn pin(&self, xid: Xid) {
        if xid == XID_NONE {
            return;
        }
        let mut state = self.state.write();
        *state.pinned.entry(xid).or_insert(0) += 1;
    }

    /// Decrement an xid's pin count. At zero it is removed and becomes
    /// VACUUM-eligible again. No-op if the xid was never pinned.
    pub fn unpin(&self, xid: Xid) {
        if xid == XID_NONE {
            return;
        }
        let mut state = self.state.write();
        if let Some(count) = state.pinned.get_mut(&xid) {
            if *count <= 1 {
                state.pinned.remove(&xid);
            } else {
                *count -= 1;
            }
        }
    }

    /// Is this xid currently pinned?
    pub fn is_pinned(&self, xid: Xid) -> bool {
        self.state.read().pinned.contains_key(&xid)
    }

    /// Current pin count for an xid (0 if not pinned). Diagnostic only.
    pub fn pin_count(&self, xid: Xid) -> u32 {
        self.state.read().pinned.get(&xid).copied().unwrap_or(0)
    }
}

impl Default for SnapshotManager {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn xids_are_monotonic() {
        let m = SnapshotManager::new();
        let a = m.begin();
        let b = m.begin();
        let c = m.begin();
        assert!(a < b && b < c);
    }

    #[test]
    fn snapshot_excludes_concurrent_writers() {
        let m = SnapshotManager::new();
        let writer = m.begin();
        let reader = m.begin();
        let snap = m.snapshot(reader);
        // Writer is active from reader's perspective → in_progress set.
        assert!(snap.in_progress.contains(&writer));
        // A row written by `writer` with xmin=writer must be invisible.
        assert!(!snap.sees(writer, XID_NONE));
    }

    #[test]
    fn fresh_read_snapshot_sees_latest_committed_hides_active() {
        let m = SnapshotManager::new();
        // A committed writer must be visible to a fresh read snapshot; an
        // xid still active must be hidden. This is the RMW pre-image
        // contract: latest committed, never an uncommitted sibling.
        let committed = m.begin();
        m.commit(committed);
        let active = m.begin();

        let snap = m.fresh_read_snapshot();
        assert!(snap.sees(committed, XID_NONE));
        assert!(!snap.sees(active, XID_NONE));
        assert!(snap.in_progress.contains(&active));
    }

    #[test]
    fn fresh_read_snapshot_refreshes_between_calls() {
        let m = SnapshotManager::new();
        // Value committed *after* a first read must be visible to a second
        // fresh read — proving the snapshot is re-evaluated, not pinned.
        let later = m.begin();
        let first = m.fresh_read_snapshot();
        assert!(!first.sees(later, XID_NONE));
        m.commit(later);
        let second = m.fresh_read_snapshot();
        assert!(second.sees(later, XID_NONE));
    }

    #[test]
    fn committed_rows_become_visible() {
        let m = SnapshotManager::new();
        let writer = m.begin();
        m.commit(writer);
        let reader = m.begin();
        let snap = m.snapshot(reader);
        // Row stamped with writer's xid is now visible (writer < reader & committed).
        assert!(snap.sees(writer, XID_NONE));
    }

    #[test]
    fn rolled_back_writers_stay_hidden() {
        let m = SnapshotManager::new();
        let writer = m.begin();
        m.rollback(writer);
        assert!(m.is_aborted(writer));
        // Future callers skip tuples with xmin == writer by also consulting is_aborted.
    }

    #[test]
    fn pre_mvcc_rows_always_visible() {
        let m = SnapshotManager::new();
        let reader = m.begin();
        let snap = m.snapshot(reader);
        assert!(snap.sees(XID_NONE, XID_NONE));
    }

    #[test]
    fn deletion_xmax_respected() {
        let m = SnapshotManager::new();
        let creator = m.begin();
        m.commit(creator);
        let deleter = m.begin();
        m.commit(deleter);
        let reader = m.begin();
        let snap = m.snapshot(reader);
        // Reader opens *after* delete → row must be hidden.
        assert!(!snap.sees(creator, deleter));
    }

    #[test]
    fn pin_blocks_prune_of_aborted_xid() {
        let m = SnapshotManager::new();
        let writer = m.begin();
        m.rollback(writer);
        assert!(m.is_aborted(writer));
        m.pin(writer);
        // Even with a high watermark, pinned xid survives prune.
        m.prune_aborted(writer + 1);
        assert!(m.is_aborted(writer));
        m.unpin(writer);
        m.prune_aborted(writer + 1);
        assert!(!m.is_aborted(writer));
    }

    #[test]
    fn pin_is_reference_counted() {
        let m = SnapshotManager::new();
        let x = m.begin();
        m.pin(x);
        m.pin(x);
        assert_eq!(m.pin_count(x), 2);
        m.unpin(x);
        assert_eq!(m.pin_count(x), 1);
        assert!(m.is_pinned(x));
        m.unpin(x);
        assert_eq!(m.pin_count(x), 0);
        assert!(!m.is_pinned(x));
        // Extra unpin is a no-op.
        m.unpin(x);
        assert_eq!(m.pin_count(x), 0);
    }

    #[test]
    fn pin_xid_none_is_noop() {
        let m = SnapshotManager::new();
        m.pin(XID_NONE);
        assert!(!m.is_pinned(XID_NONE));
        assert_eq!(m.pin_count(XID_NONE), 0);
    }

    #[test]
    fn allocate_committed_xid_is_monotonic_and_unique() {
        let m = SnapshotManager::new();
        let mut seen = HashSet::new();
        let mut last = 0u64;
        // Drive at least three pool refills (BATCH=16 → 50 covers it).
        for _ in 0..50 {
            let x = m.allocate_committed_xid();
            assert!(x > last, "xids must be strictly increasing: {x} > {last}");
            assert!(seen.insert(x), "duplicate xid handed out: {x}");
            last = x;
        }
    }

    #[test]
    fn mixed_allocations_preserve_transaction_order_and_reclaim_horizon() {
        let manager = Arc::new(SnapshotManager::new());
        let old = manager.allocate_committed_xid();
        let deletion = manager.begin();
        manager.commit(deletion);
        let next_insert = manager.allocate_committed_xid();
        assert!(
            next_insert > deletion,
            "a later insert cannot reuse an older reserved xid"
        );
        let (snapshot, _retention) = manager.fresh_retained_snapshot();
        assert!(!snapshot.sees(old, deletion));
        assert!(snapshot.sees(next_insert, XID_NONE));
        assert!(
            manager.oldest_required_xid().expect("retained statement") > deletion,
            "unused allocator slots cannot keep already visible deletes resident"
        );
    }

    #[test]
    fn allocate_committed_xid_skips_active_set() {
        let m = SnapshotManager::new();
        let _x = m.allocate_committed_xid();
        // Pool xids must never appear in the active set — they are
        // born committed. `oldest_active_xid` reflects only `begin()`
        // callers (real BEGIN-wrapped transactions).
        assert_eq!(m.oldest_active_xid(), None);
    }

    #[test]
    fn allocate_committed_xid_visible_to_subsequent_snapshots() {
        let m = SnapshotManager::new();
        let writer = m.allocate_committed_xid();
        let reader = m.begin();
        let snap = m.snapshot(reader);
        // Pool xid must be invisible to in_progress/aborted (it's in
        // neither) and visible because writer < reader. This matches
        // the legacy begin()+commit() pair's visibility exactly.
        assert!(!snap.in_progress.contains(&writer));
        assert!(!m.is_aborted(writer));
        assert!(snap.sees(writer, XID_NONE));
    }

    #[test]
    fn allocate_committed_xid_does_not_block_concurrent_begin() {
        // Smoke test: an open BEGIN-wrapped tx coexists with pool
        // allocation; pool xids end up between the begin and commit
        // without being added to `active`.
        let m = SnapshotManager::new();
        let tx = m.begin();
        let auto1 = m.allocate_committed_xid();
        let auto2 = m.allocate_committed_xid();
        m.commit(tx);
        assert!(tx < auto1 && auto1 < auto2);
        // `active` should be empty after commit.
        assert_eq!(m.oldest_active_xid(), None);
    }

    #[test]
    fn oldest_active_is_min_live_xid() {
        let m = SnapshotManager::new();
        let a = m.begin();
        let b = m.begin();
        assert_eq!(m.oldest_active_xid(), Some(a));
        m.commit(a);
        assert_eq!(m.oldest_active_xid(), Some(b));
        m.commit(b);
        assert_eq!(m.oldest_active_xid(), None);
    }
    #[test]
    fn read_snapshot_excludes_first_future_writer_and_deleter() {
        let manager = SnapshotManager::new();
        let old = manager.begin();
        manager.commit(old);
        let snapshot = manager.fresh_read_snapshot();
        let future = manager.begin();
        manager.commit(future);
        assert!(!snapshot.sees(future, XID_NONE));
        assert!(snapshot.sees(old, future));
        let fresh = manager.fresh_read_snapshot();
        assert!(fresh.sees(future, XID_NONE));
        assert!(!fresh.sees(old, future));
    }

    #[test]
    fn snapshots_hide_autocommit_writers_allocated_after_begin() {
        let manager = SnapshotManager::new();
        let old = manager.allocate_committed_xid();
        let reader = manager.begin();
        let transaction_snapshot = manager.snapshot(reader);
        let statement_snapshot = manager.fresh_read_snapshot();
        let future = manager.allocate_committed_xid();
        assert!(future > reader, "begin retires the older reserved range");
        for snapshot in [transaction_snapshot, statement_snapshot] {
            assert!(snapshot.sees(old, XID_NONE));
            assert!(!snapshot.sees(future, XID_NONE));
            assert!(snapshot.sees(old, future));
        }
        assert!(manager.fresh_read_snapshot().sees(future, XID_NONE));
        manager.rollback(reader);
    }

    #[test]
    fn retiring_unused_slots_preserves_an_older_readers_history() {
        let manager = Arc::new(SnapshotManager::new());
        let old = manager.allocate_committed_xid();
        let (snapshot, retention) = manager.fresh_retained_snapshot();
        let deletion = manager.begin();
        manager.commit(deletion);
        let next_insert = manager.allocate_committed_xid();
        let (_current, _current_retention) = manager.fresh_retained_snapshot();
        assert!(snapshot.sees(old, deletion));
        assert!(!snapshot.sees(next_insert, XID_NONE));
        assert!(manager.oldest_required_xid().expect("old reader retained") <= deletion);
        drop(retention);
        assert!(
            manager
                .oldest_required_xid()
                .expect("current reader retained")
                > deletion
        );
    }

    #[test]
    fn replay_retires_reserved_xids_without_changing_captured_snapshots() {
        let manager = SnapshotManager::new();
        let first = manager.allocate_committed_xid();
        let captured = manager.fresh_read_snapshot();
        let replayed = first + 2;
        manager.observe_committed_xid(replayed);
        assert!(!captured.sees(replayed, XID_NONE));
        assert!(manager.fresh_read_snapshot().sees(replayed, XID_NONE));
        assert!(manager.allocate_committed_xid() > replayed);

        let beyond_pool = manager.peek_next_xid() + 10;
        let captured = manager.fresh_read_snapshot();
        manager.observe_committed_xid(beyond_pool);
        assert!(!captured.sees(beyond_pool, XID_NONE));
        assert!(manager.fresh_read_snapshot().sees(beyond_pool, XID_NONE));
        assert!(manager.allocate_committed_xid() > beyond_pool);
    }
}