reddb-io-engine 1.27.12

//! Per-connection transaction state and lifecycle transitions.
//!
//! This module owns the connection-to-transaction and transaction-local
//! tenant maps. Runtime dispatch supplies commit-time validation and pending
//! write finalization, while transaction and savepoint state changes stay
//! behind this interface.

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

use parking_lot::RwLock;

use crate::transaction::snapshot::{
    Snapshot, SnapshotManager, SnapshotRetentionGuard, TxnContext, Xid,
};
use crate::transaction::IsolationLevel;

/// A requested savepoint is absent from the active transaction.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SavepointError {
    pub name: String,
}

impl std::fmt::Display for SavepointError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(formatter, "savepoint {} does not exist", self.name)
    }
}

impl std::error::Error for SavepointError {}

pub struct RollbackToSavepoint {
    pub savepoint_xid: Xid,
    pub aborted_xids: Vec<Xid>,
}

/// Volatile connection state for a host transaction dispatcher.
///
/// The host serializes lifecycle calls per connection and validates transaction
/// control before calling this layer. Commit preparation supplies write-conflict
/// validation and durable publication; this state machine performs neither I/O
/// nor recovery, and its visibility transitions alone do not acknowledge a write.
pub struct TransactionState {
    snapshot_manager: Arc<SnapshotManager>,
    contexts: RwLock<HashMap<u64, TxnContext>>,
    local_tenants: RwLock<HashMap<u64, Option<String>>>,
}

impl TransactionState {
    pub fn new() -> Self {
        Self {
            snapshot_manager: Arc::new(SnapshotManager::new()),
            contexts: RwLock::new(HashMap::new()),
            local_tenants: RwLock::new(HashMap::new()),
        }
    }

    pub fn begin(&self, connection_id: u64, isolation: IsolationLevel) -> Xid {
        let xid = self.snapshot_manager.begin();
        if isolation == IsolationLevel::Serializable {
            self.snapshot_manager.begin_serializable(xid);
        }
        let context = TxnContext {
            xid,
            isolation,
            snapshot: self.snapshot_manager.snapshot(xid),
            savepoints: Vec::new(),
            released_sub_xids: Vec::new(),
        };
        self.contexts.write().insert(connection_id, context);
        xid
    }

    pub fn commit<E, Publication>(
        &self,
        connection_id: u64,
        prepare: impl FnOnce(&TxnContext) -> Result<Publication, E>,
    ) -> Result<Option<TxnContext>, E> {
        self.local_tenants.write().remove(&connection_id);
        let Some(context) = self.contexts.write().remove(&connection_id) else {
            return Ok(None);
        };

        // A writer can return its publication guard after validation and WAL
        // append. Keep it through the transaction's visibility transition.
        let _publication = match prepare(&context) {
            Ok(publication) => publication,
            Err(error) => {
                self.rollback_context(&context);
                return Err(error);
            }
        };

        for (_, xid) in &context.savepoints {
            self.snapshot_manager.commit(*xid);
        }
        for xid in &context.released_sub_xids {
            self.snapshot_manager.commit(*xid);
        }
        self.snapshot_manager.commit(context.xid);
        Ok(Some(context))
    }

    pub fn rollback(&self, connection_id: u64) -> Option<TxnContext> {
        self.local_tenants.write().remove(&connection_id);
        let context = self.contexts.write().remove(&connection_id)?;
        self.rollback_context(&context);
        Some(context)
    }

    fn rollback_context(&self, context: &TxnContext) {
        for (_, xid) in &context.savepoints {
            self.snapshot_manager.rollback(*xid);
        }
        for xid in &context.released_sub_xids {
            self.snapshot_manager.rollback(*xid);
        }
        self.snapshot_manager.rollback(context.xid);
    }

    pub fn has_savepoint(&self, connection_id: u64, name: &str) -> bool {
        self.contexts
            .read()
            .get(&connection_id)
            .is_some_and(|context| {
                context
                    .savepoints
                    .iter()
                    .any(|(existing, _)| existing == name)
            })
    }

    pub fn savepoint(&self, connection_id: u64, name: &str) -> Option<Xid> {
        let mut contexts = self.contexts.write();
        let context = contexts.get_mut(&connection_id)?;
        let xid = self.snapshot_manager.begin();
        context.savepoints.push((name.to_string(), xid));
        Some(xid)
    }

    pub fn release_savepoint(
        &self,
        connection_id: u64,
        name: &str,
    ) -> Result<Option<usize>, SavepointError> {
        let mut contexts = self.contexts.write();
        let Some(context) = contexts.get_mut(&connection_id) else {
            return Ok(None);
        };
        let position = context
            .savepoints
            .iter()
            .position(|(savepoint_name, _)| savepoint_name == name)
            .ok_or_else(|| SavepointError {
                name: name.to_string(),
            })?;
        let released = context.savepoints.len() - position;
        context.released_sub_xids.extend(
            context
                .savepoints
                .split_off(position)
                .into_iter()
                .map(|(_, xid)| xid),
        );
        Ok(Some(released))
    }

    pub fn rollback_to_savepoint(
        &self,
        connection_id: u64,
        name: &str,
    ) -> Result<Option<RollbackToSavepoint>, SavepointError> {
        self.rollback_to_savepoint_with(connection_id, name, |_| Ok::<_, SavepointError>(()))
            .map(|rollback| rollback.map(|(rollback, ())| rollback))
    }

    /// Keep subtransactions active while the host restores rows and indexes.
    /// Failed undo still aborts their visibility before returning the error.
    pub fn rollback_to_savepoint_with<T, E: From<SavepointError>>(
        &self,
        connection_id: u64,
        name: &str,
        prepare: impl FnOnce(&RollbackToSavepoint) -> Result<T, E>,
    ) -> Result<Option<(RollbackToSavepoint, T)>, E> {
        let mut contexts = self.contexts.write();
        let Some(context) = contexts.get_mut(&connection_id) else {
            return Ok(None);
        };
        let position = context
            .savepoints
            .iter()
            .position(|(savepoint_name, _)| savepoint_name == name)
            .ok_or_else(|| SavepointError {
                name: name.to_string(),
            })?;
        let savepoint_xid = context.savepoints[position].1;
        let aborted_xids = context
            .savepoints
            .split_off(position)
            .into_iter()
            .map(|(_, xid)| xid)
            .collect::<Vec<_>>();
        drop(contexts);

        let rollback = RollbackToSavepoint {
            savepoint_xid,
            aborted_xids,
        };
        let prepared = prepare(&rollback);
        for xid in &rollback.aborted_xids {
            self.snapshot_manager.rollback(*xid);
        }
        prepared.map(|prepared| Some((rollback, prepared)))
    }

    pub fn context(&self, connection_id: u64) -> Option<TxnContext> {
        self.contexts.read().get(&connection_id).cloned()
    }

    pub fn in_transaction(&self, connection_id: u64) -> bool {
        self.contexts.read().contains_key(&connection_id)
    }

    pub fn current_snapshot(&self, connection_id: u64) -> Snapshot {
        if let Some(context) = self.context(connection_id) {
            if context.isolation == IsolationLevel::ReadCommitted {
                return self.snapshot_manager.fresh_read_snapshot();
            }
            return context.snapshot;
        }
        self.snapshot_manager.fresh_read_snapshot()
    }

    pub fn retained_current_snapshot(
        &self,
        connection_id: u64,
    ) -> (Snapshot, SnapshotRetentionGuard) {
        if let Some(context) = self.context(connection_id) {
            if context.isolation != IsolationLevel::ReadCommitted {
                return self.snapshot_manager.retain_snapshot(context.snapshot);
            }
        }
        self.snapshot_manager.fresh_retained_snapshot()
    }

    pub fn writer_xid(&self, connection_id: u64) -> Option<Xid> {
        self.context(connection_id)
            .map(|context| context.writer_xid())
    }

    pub fn own_xids(&self, connection_id: u64) -> HashSet<Xid> {
        let mut xids = HashSet::new();
        if let Some(context) = self.context(connection_id) {
            xids.insert(context.xid);
            xids.extend(context.savepoints.into_iter().map(|(_, xid)| xid));
            xids.extend(context.released_sub_xids);
        }
        xids
    }

    pub fn set_local_tenant(&self, connection_id: u64, tenant: Option<String>) {
        self.local_tenants.write().insert(connection_id, tenant);
    }

    pub fn local_tenant(&self, connection_id: u64) -> Option<Option<String>> {
        self.local_tenants.read().get(&connection_id).cloned()
    }

    pub fn snapshot_manager(&self) -> Arc<SnapshotManager> {
        Arc::clone(&self.snapshot_manager)
    }
}

#[cfg(test)]
mod tests {
    use super::SavepointError;
    use crate::transaction::IsolationLevel;

    use super::TransactionState;

    #[test]
    fn begin_commit_and_rollback_update_snapshot_state() {
        let state = TransactionState::new();

        let committed_xid = state.begin(11, IsolationLevel::SnapshotIsolation);
        assert!(state.in_transaction(11));
        assert_eq!(state.writer_xid(11), Some(committed_xid));
        assert!(state.snapshot_manager().is_active(committed_xid));

        let committed = state
            .commit(11, |_| Ok::<_, ()>(()))
            .expect("commit preparation should succeed")
            .expect("connection should have an active transaction");
        assert_eq!(committed.xid, committed_xid);
        assert!(!state.in_transaction(11));
        assert!(!state.snapshot_manager().is_active(committed_xid));
        assert!(!state.snapshot_manager().is_aborted(committed_xid));

        let aborted_xid = state.begin(11, IsolationLevel::ReadCommitted);
        let aborted = state
            .rollback(11)
            .expect("connection should have an active transaction");
        assert_eq!(aborted.xid, aborted_xid);
        assert!(!state.in_transaction(11));
        assert!(state.snapshot_manager().is_aborted(aborted_xid));
    }

    #[test]
    fn savepoint_undo_retains_active_xids_until_preparation_finishes_even_on_error() {
        for fail in [false, true] {
            let state = TransactionState::new();
            state.begin(21, IsolationLevel::SnapshotIsolation);
            let xid = state.savepoint(21, "undo").expect("savepoint");
            let result = state.rollback_to_savepoint_with(21, "undo", |rollback| {
                assert_eq!(rollback.aborted_xids, vec![xid]);
                assert!(state.snapshot_manager().is_active(xid));
                assert!(!state.snapshot_manager().is_aborted(xid));
                if fail {
                    Err(SavepointError {
                        name: "undo failed".into(),
                    })
                } else {
                    Ok(())
                }
            });
            assert_eq!(result.is_err(), fail);
            assert!(!state.snapshot_manager().is_active(xid));
            assert!(state.snapshot_manager().is_aborted(xid));
            assert!(state.in_transaction(21));
        }
    }

    #[test]
    fn savepoint_release_and_rollback_follow_nested_sequence() {
        let state = TransactionState::new();
        let parent_xid = state.begin(21, IsolationLevel::SnapshotIsolation);
        let outer_xid = state
            .savepoint(21, "outer")
            .expect("transaction should accept a savepoint");
        let released_xid = state
            .savepoint(21, "released")
            .expect("transaction should accept a nested savepoint");

        assert_eq!(
            state
                .release_savepoint(21, "released")
                .expect("savepoint should exist"),
            Some(1)
        );
        assert_eq!(state.writer_xid(21), Some(outer_xid));

        let nested_xid = state
            .savepoint(21, "nested")
            .expect("transaction should accept another nested savepoint");
        let rolled_back = state
            .rollback_to_savepoint(21, "outer")
            .expect("savepoint should exist")
            .expect("connection should have an active transaction");
        assert_eq!(rolled_back.savepoint_xid, outer_xid);
        assert_eq!(rolled_back.aborted_xids, vec![outer_xid, nested_xid]);
        assert_eq!(state.writer_xid(21), Some(parent_xid));
        assert!(state.snapshot_manager().is_aborted(outer_xid));
        assert!(state.snapshot_manager().is_aborted(nested_xid));

        state
            .commit(21, |_| Ok::<_, ()>(()))
            .expect("commit preparation should succeed")
            .expect("connection should have an active transaction");
        assert!(!state.snapshot_manager().is_active(parent_xid));
        assert!(!state.snapshot_manager().is_active(released_xid));
        assert!(!state.snapshot_manager().is_aborted(released_xid));
    }

    #[test]
    fn connections_keep_transaction_and_local_tenant_state_isolated() {
        let state = TransactionState::new();
        let first_xid = state.begin(31, IsolationLevel::SnapshotIsolation);
        let second_xid = state.begin(32, IsolationLevel::ReadCommitted);
        let first_savepoint = state
            .savepoint(31, "first_only")
            .expect("first connection should have an active transaction");

        assert_eq!(state.writer_xid(31), Some(first_savepoint));
        assert_eq!(state.writer_xid(32), Some(second_xid));
        assert_eq!(state.context(31).expect("first context").xid, first_xid);
        assert_eq!(state.context(32).expect("second context").xid, second_xid);

        state.set_local_tenant(31, Some("tenant-a".to_string()));
        state.set_local_tenant(32, Some("tenant-b".to_string()));
        assert_eq!(
            state.local_tenant(31).flatten().as_deref(),
            Some("tenant-a")
        );
        assert_eq!(
            state.local_tenant(32).flatten().as_deref(),
            Some("tenant-b")
        );

        state
            .rollback(31)
            .expect("first transaction should roll back");
        assert!(!state.in_transaction(31));
        assert!(state.in_transaction(32));
        assert_eq!(state.local_tenant(31), None);
        assert_eq!(
            state.local_tenant(32).flatten().as_deref(),
            Some("tenant-b")
        );

        state
            .commit(32, |_| Ok::<_, ()>(()))
            .expect("commit preparation should succeed")
            .expect("connection should have an active transaction");
        assert_eq!(state.local_tenant(32), None);
    }
}