reddb_engine/transaction/
state.rs

1//! Per-connection transaction state and lifecycle transitions.
2//!
3//! This module owns the connection-to-transaction and transaction-local
4//! tenant maps. Runtime dispatch supplies commit-time validation and pending
5//! write finalization, while transaction and savepoint state changes stay
6//! behind this interface.
7
8use std::collections::{HashMap, HashSet};
9use std::sync::Arc;
10
11use parking_lot::RwLock;
12
13use crate::transaction::snapshot::{
14    Snapshot, SnapshotManager, SnapshotRetentionGuard, TxnContext, Xid,
15};
16use crate::transaction::IsolationLevel;
17
18/// A requested savepoint is absent from the active transaction.
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct SavepointError {
21    pub name: String,
22}
23
24impl std::fmt::Display for SavepointError {
25    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        write!(formatter, "savepoint {} does not exist", self.name)
27    }
28}
29
30impl std::error::Error for SavepointError {}
31
32pub struct RollbackToSavepoint {
33    pub savepoint_xid: Xid,
34    pub aborted_xids: Vec<Xid>,
35}
36
37/// Volatile connection state for a host transaction dispatcher.
38///
39/// The host serializes lifecycle calls per connection and validates transaction
40/// control before calling this layer. Commit preparation supplies write-conflict
41/// validation and durable publication; this state machine performs neither I/O
42/// nor recovery, and its visibility transitions alone do not acknowledge a write.
43pub struct TransactionState {
44    snapshot_manager: Arc<SnapshotManager>,
45    contexts: RwLock<HashMap<u64, TxnContext>>,
46    local_tenants: RwLock<HashMap<u64, Option<String>>>,
47}
48
49impl TransactionState {
50    pub fn new() -> Self {
51        Self {
52            snapshot_manager: Arc::new(SnapshotManager::new()),
53            contexts: RwLock::new(HashMap::new()),
54            local_tenants: RwLock::new(HashMap::new()),
55        }
56    }
57
58    pub fn begin(&self, connection_id: u64, isolation: IsolationLevel) -> Xid {
59        let xid = self.snapshot_manager.begin();
60        if isolation == IsolationLevel::Serializable {
61            self.snapshot_manager.begin_serializable(xid);
62        }
63        let context = TxnContext {
64            xid,
65            isolation,
66            snapshot: self.snapshot_manager.snapshot(xid),
67            savepoints: Vec::new(),
68            released_sub_xids: Vec::new(),
69        };
70        self.contexts.write().insert(connection_id, context);
71        xid
72    }
73
74    pub fn commit<E, Publication>(
75        &self,
76        connection_id: u64,
77        prepare: impl FnOnce(&TxnContext) -> Result<Publication, E>,
78    ) -> Result<Option<TxnContext>, E> {
79        self.local_tenants.write().remove(&connection_id);
80        let Some(context) = self.contexts.write().remove(&connection_id) else {
81            return Ok(None);
82        };
83
84        // A writer can return its publication guard after validation and WAL
85        // append. Keep it through the transaction's visibility transition.
86        let _publication = match prepare(&context) {
87            Ok(publication) => publication,
88            Err(error) => {
89                self.rollback_context(&context);
90                return Err(error);
91            }
92        };
93
94        for (_, xid) in &context.savepoints {
95            self.snapshot_manager.commit(*xid);
96        }
97        for xid in &context.released_sub_xids {
98            self.snapshot_manager.commit(*xid);
99        }
100        self.snapshot_manager.commit(context.xid);
101        Ok(Some(context))
102    }
103
104    pub fn rollback(&self, connection_id: u64) -> Option<TxnContext> {
105        self.local_tenants.write().remove(&connection_id);
106        let context = self.contexts.write().remove(&connection_id)?;
107        self.rollback_context(&context);
108        Some(context)
109    }
110
111    fn rollback_context(&self, context: &TxnContext) {
112        for (_, xid) in &context.savepoints {
113            self.snapshot_manager.rollback(*xid);
114        }
115        for xid in &context.released_sub_xids {
116            self.snapshot_manager.rollback(*xid);
117        }
118        self.snapshot_manager.rollback(context.xid);
119    }
120
121    pub fn has_savepoint(&self, connection_id: u64, name: &str) -> bool {
122        self.contexts
123            .read()
124            .get(&connection_id)
125            .is_some_and(|context| {
126                context
127                    .savepoints
128                    .iter()
129                    .any(|(existing, _)| existing == name)
130            })
131    }
132
133    pub fn savepoint(&self, connection_id: u64, name: &str) -> Option<Xid> {
134        let mut contexts = self.contexts.write();
135        let context = contexts.get_mut(&connection_id)?;
136        let xid = self.snapshot_manager.begin();
137        context.savepoints.push((name.to_string(), xid));
138        Some(xid)
139    }
140
141    pub fn release_savepoint(
142        &self,
143        connection_id: u64,
144        name: &str,
145    ) -> Result<Option<usize>, SavepointError> {
146        let mut contexts = self.contexts.write();
147        let Some(context) = contexts.get_mut(&connection_id) else {
148            return Ok(None);
149        };
150        let position = context
151            .savepoints
152            .iter()
153            .position(|(savepoint_name, _)| savepoint_name == name)
154            .ok_or_else(|| SavepointError {
155                name: name.to_string(),
156            })?;
157        let released = context.savepoints.len() - position;
158        context.released_sub_xids.extend(
159            context
160                .savepoints
161                .split_off(position)
162                .into_iter()
163                .map(|(_, xid)| xid),
164        );
165        Ok(Some(released))
166    }
167
168    pub fn rollback_to_savepoint(
169        &self,
170        connection_id: u64,
171        name: &str,
172    ) -> Result<Option<RollbackToSavepoint>, SavepointError> {
173        self.rollback_to_savepoint_with(connection_id, name, |_| Ok::<_, SavepointError>(()))
174            .map(|rollback| rollback.map(|(rollback, ())| rollback))
175    }
176
177    /// Keep subtransactions active while the host restores rows and indexes.
178    /// Failed undo still aborts their visibility before returning the error.
179    pub fn rollback_to_savepoint_with<T, E: From<SavepointError>>(
180        &self,
181        connection_id: u64,
182        name: &str,
183        prepare: impl FnOnce(&RollbackToSavepoint) -> Result<T, E>,
184    ) -> Result<Option<(RollbackToSavepoint, T)>, E> {
185        let mut contexts = self.contexts.write();
186        let Some(context) = contexts.get_mut(&connection_id) else {
187            return Ok(None);
188        };
189        let position = context
190            .savepoints
191            .iter()
192            .position(|(savepoint_name, _)| savepoint_name == name)
193            .ok_or_else(|| SavepointError {
194                name: name.to_string(),
195            })?;
196        let savepoint_xid = context.savepoints[position].1;
197        let aborted_xids = context
198            .savepoints
199            .split_off(position)
200            .into_iter()
201            .map(|(_, xid)| xid)
202            .collect::<Vec<_>>();
203        drop(contexts);
204
205        let rollback = RollbackToSavepoint {
206            savepoint_xid,
207            aborted_xids,
208        };
209        let prepared = prepare(&rollback);
210        for xid in &rollback.aborted_xids {
211            self.snapshot_manager.rollback(*xid);
212        }
213        prepared.map(|prepared| Some((rollback, prepared)))
214    }
215
216    pub fn context(&self, connection_id: u64) -> Option<TxnContext> {
217        self.contexts.read().get(&connection_id).cloned()
218    }
219
220    pub fn in_transaction(&self, connection_id: u64) -> bool {
221        self.contexts.read().contains_key(&connection_id)
222    }
223
224    pub fn current_snapshot(&self, connection_id: u64) -> Snapshot {
225        if let Some(context) = self.context(connection_id) {
226            if context.isolation == IsolationLevel::ReadCommitted {
227                return self.snapshot_manager.fresh_read_snapshot();
228            }
229            return context.snapshot;
230        }
231        self.snapshot_manager.fresh_read_snapshot()
232    }
233
234    pub fn retained_current_snapshot(
235        &self,
236        connection_id: u64,
237    ) -> (Snapshot, SnapshotRetentionGuard) {
238        if let Some(context) = self.context(connection_id) {
239            if context.isolation != IsolationLevel::ReadCommitted {
240                return self.snapshot_manager.retain_snapshot(context.snapshot);
241            }
242        }
243        self.snapshot_manager.fresh_retained_snapshot()
244    }
245
246    pub fn writer_xid(&self, connection_id: u64) -> Option<Xid> {
247        self.context(connection_id)
248            .map(|context| context.writer_xid())
249    }
250
251    pub fn own_xids(&self, connection_id: u64) -> HashSet<Xid> {
252        let mut xids = HashSet::new();
253        if let Some(context) = self.context(connection_id) {
254            xids.insert(context.xid);
255            xids.extend(context.savepoints.into_iter().map(|(_, xid)| xid));
256            xids.extend(context.released_sub_xids);
257        }
258        xids
259    }
260
261    pub fn set_local_tenant(&self, connection_id: u64, tenant: Option<String>) {
262        self.local_tenants.write().insert(connection_id, tenant);
263    }
264
265    pub fn local_tenant(&self, connection_id: u64) -> Option<Option<String>> {
266        self.local_tenants.read().get(&connection_id).cloned()
267    }
268
269    pub fn snapshot_manager(&self) -> Arc<SnapshotManager> {
270        Arc::clone(&self.snapshot_manager)
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::SavepointError;
277    use crate::transaction::IsolationLevel;
278
279    use super::TransactionState;
280
281    #[test]
282    fn begin_commit_and_rollback_update_snapshot_state() {
283        let state = TransactionState::new();
284
285        let committed_xid = state.begin(11, IsolationLevel::SnapshotIsolation);
286        assert!(state.in_transaction(11));
287        assert_eq!(state.writer_xid(11), Some(committed_xid));
288        assert!(state.snapshot_manager().is_active(committed_xid));
289
290        let committed = state
291            .commit(11, |_| Ok::<_, ()>(()))
292            .expect("commit preparation should succeed")
293            .expect("connection should have an active transaction");
294        assert_eq!(committed.xid, committed_xid);
295        assert!(!state.in_transaction(11));
296        assert!(!state.snapshot_manager().is_active(committed_xid));
297        assert!(!state.snapshot_manager().is_aborted(committed_xid));
298
299        let aborted_xid = state.begin(11, IsolationLevel::ReadCommitted);
300        let aborted = state
301            .rollback(11)
302            .expect("connection should have an active transaction");
303        assert_eq!(aborted.xid, aborted_xid);
304        assert!(!state.in_transaction(11));
305        assert!(state.snapshot_manager().is_aborted(aborted_xid));
306    }
307
308    #[test]
309    fn savepoint_undo_retains_active_xids_until_preparation_finishes_even_on_error() {
310        for fail in [false, true] {
311            let state = TransactionState::new();
312            state.begin(21, IsolationLevel::SnapshotIsolation);
313            let xid = state.savepoint(21, "undo").expect("savepoint");
314            let result = state.rollback_to_savepoint_with(21, "undo", |rollback| {
315                assert_eq!(rollback.aborted_xids, vec![xid]);
316                assert!(state.snapshot_manager().is_active(xid));
317                assert!(!state.snapshot_manager().is_aborted(xid));
318                if fail {
319                    Err(SavepointError {
320                        name: "undo failed".into(),
321                    })
322                } else {
323                    Ok(())
324                }
325            });
326            assert_eq!(result.is_err(), fail);
327            assert!(!state.snapshot_manager().is_active(xid));
328            assert!(state.snapshot_manager().is_aborted(xid));
329            assert!(state.in_transaction(21));
330        }
331    }
332
333    #[test]
334    fn savepoint_release_and_rollback_follow_nested_sequence() {
335        let state = TransactionState::new();
336        let parent_xid = state.begin(21, IsolationLevel::SnapshotIsolation);
337        let outer_xid = state
338            .savepoint(21, "outer")
339            .expect("transaction should accept a savepoint");
340        let released_xid = state
341            .savepoint(21, "released")
342            .expect("transaction should accept a nested savepoint");
343
344        assert_eq!(
345            state
346                .release_savepoint(21, "released")
347                .expect("savepoint should exist"),
348            Some(1)
349        );
350        assert_eq!(state.writer_xid(21), Some(outer_xid));
351
352        let nested_xid = state
353            .savepoint(21, "nested")
354            .expect("transaction should accept another nested savepoint");
355        let rolled_back = state
356            .rollback_to_savepoint(21, "outer")
357            .expect("savepoint should exist")
358            .expect("connection should have an active transaction");
359        assert_eq!(rolled_back.savepoint_xid, outer_xid);
360        assert_eq!(rolled_back.aborted_xids, vec![outer_xid, nested_xid]);
361        assert_eq!(state.writer_xid(21), Some(parent_xid));
362        assert!(state.snapshot_manager().is_aborted(outer_xid));
363        assert!(state.snapshot_manager().is_aborted(nested_xid));
364
365        state
366            .commit(21, |_| Ok::<_, ()>(()))
367            .expect("commit preparation should succeed")
368            .expect("connection should have an active transaction");
369        assert!(!state.snapshot_manager().is_active(parent_xid));
370        assert!(!state.snapshot_manager().is_active(released_xid));
371        assert!(!state.snapshot_manager().is_aborted(released_xid));
372    }
373
374    #[test]
375    fn connections_keep_transaction_and_local_tenant_state_isolated() {
376        let state = TransactionState::new();
377        let first_xid = state.begin(31, IsolationLevel::SnapshotIsolation);
378        let second_xid = state.begin(32, IsolationLevel::ReadCommitted);
379        let first_savepoint = state
380            .savepoint(31, "first_only")
381            .expect("first connection should have an active transaction");
382
383        assert_eq!(state.writer_xid(31), Some(first_savepoint));
384        assert_eq!(state.writer_xid(32), Some(second_xid));
385        assert_eq!(state.context(31).expect("first context").xid, first_xid);
386        assert_eq!(state.context(32).expect("second context").xid, second_xid);
387
388        state.set_local_tenant(31, Some("tenant-a".to_string()));
389        state.set_local_tenant(32, Some("tenant-b".to_string()));
390        assert_eq!(
391            state.local_tenant(31).flatten().as_deref(),
392            Some("tenant-a")
393        );
394        assert_eq!(
395            state.local_tenant(32).flatten().as_deref(),
396            Some("tenant-b")
397        );
398
399        state
400            .rollback(31)
401            .expect("first transaction should roll back");
402        assert!(!state.in_transaction(31));
403        assert!(state.in_transaction(32));
404        assert_eq!(state.local_tenant(31), None);
405        assert_eq!(
406            state.local_tenant(32).flatten().as_deref(),
407            Some("tenant-b")
408        );
409
410        state
411            .commit(32, |_| Ok::<_, ()>(()))
412            .expect("commit preparation should succeed")
413            .expect("connection should have an active transaction");
414        assert_eq!(state.local_tenant(32), None);
415    }
416}
reddb_engine/transaction/state.rs

reddb_engine/transaction/
state.rs
