uqa_storage/
transaction.rs1use std::collections::BTreeMap;
10
11#[derive(Debug, thiserror::Error)]
12pub enum TransactionError {
13 #[error("transaction already finished")]
14 Finished,
15 #[error("savepoint `{0}` does not exist")]
16 UnknownSavepoint(String),
17 #[error(transparent)]
18 Storage(#[from] crate::StorageBackendError),
19}
20
21pub type TxResult<T> = std::result::Result<T, TransactionError>;
22
23pub trait Snapshotable {
26 type Snapshot;
27 fn snapshot(&self) -> Self::Snapshot;
28 fn restore(&self, snapshot: &Self::Snapshot);
29}
30
31pub struct InMemoryTransaction<S: Snapshotable> {
35 source: S,
36 snapshot: Option<S::Snapshot>,
37 savepoints: BTreeMap<String, S::Snapshot>,
38 finished: bool,
39}
40
41impl<S: Snapshotable> InMemoryTransaction<S> {
42 pub fn begin(source: S) -> Self {
43 let snapshot = source.snapshot();
44 Self {
45 source,
46 snapshot: Some(snapshot),
47 savepoints: BTreeMap::new(),
48 finished: false,
49 }
50 }
51
52 pub fn active(&self) -> bool {
53 !self.finished
54 }
55
56 pub fn commit(&mut self) -> TxResult<()> {
57 if self.finished {
58 return Err(TransactionError::Finished);
59 }
60 self.snapshot = None;
61 self.savepoints.clear();
62 self.finished = true;
63 Ok(())
64 }
65
66 pub fn rollback(&mut self) -> TxResult<()> {
67 if self.finished {
68 return Err(TransactionError::Finished);
69 }
70 if let Some(snap) = self.snapshot.take() {
71 self.source.restore(&snap);
72 }
73 self.savepoints.clear();
74 self.finished = true;
75 Ok(())
76 }
77
78 pub fn savepoint(&mut self, name: impl Into<String>) -> TxResult<()> {
79 if self.finished {
80 return Err(TransactionError::Finished);
81 }
82 self.savepoints.insert(name.into(), self.source.snapshot());
83 Ok(())
84 }
85
86 pub fn release_savepoint(&mut self, name: &str) -> TxResult<()> {
87 if self.finished {
88 return Err(TransactionError::Finished);
89 }
90 self.savepoints.remove(name);
91 Ok(())
92 }
93
94 pub fn rollback_to(&mut self, name: &str) -> TxResult<()> {
95 if self.finished {
96 return Err(TransactionError::Finished);
97 }
98 let snap = self
99 .savepoints
100 .get(name)
101 .ok_or_else(|| TransactionError::UnknownSavepoint(name.to_string()))?;
102 self.source.restore(snap);
103 Ok(())
104 }
105}
106
107impl<S: Snapshotable> Drop for InMemoryTransaction<S> {
108 fn drop(&mut self) {
109 if !self.finished {
110 if let Some(snap) = self.snapshot.take() {
111 self.source.restore(&snap);
112 }
113 }
114 }
115}
116
117#[cfg(test)]
118mod tests {
119 use super::*;
120 use std::cell::RefCell;
121 use std::rc::Rc;
122
123 #[derive(Clone)]
126 struct CounterSource {
127 v: Rc<RefCell<i64>>,
128 }
129 impl Snapshotable for CounterSource {
130 type Snapshot = i64;
131 fn snapshot(&self) -> i64 {
132 *self.v.borrow()
133 }
134 fn restore(&self, snap: &i64) {
135 *self.v.borrow_mut() = *snap;
136 }
137 }
138
139 #[test]
140 fn rollback_restores_state() {
141 let v = Rc::new(RefCell::new(10));
142 let src = CounterSource { v: v.clone() };
143 let mut tx = InMemoryTransaction::begin(src);
144 *v.borrow_mut() = 99;
145 tx.rollback().unwrap();
146 assert_eq!(*v.borrow(), 10);
147 }
148
149 #[test]
150 fn commit_keeps_changes() {
151 let v = Rc::new(RefCell::new(10));
152 let src = CounterSource { v: v.clone() };
153 let mut tx = InMemoryTransaction::begin(src);
154 *v.borrow_mut() = 99;
155 tx.commit().unwrap();
156 assert_eq!(*v.borrow(), 99);
157 }
158
159 #[test]
160 fn savepoint_rollback_undoes_partial_writes() {
161 let v = Rc::new(RefCell::new(0));
162 let src = CounterSource { v: v.clone() };
163 let mut tx = InMemoryTransaction::begin(src);
164 *v.borrow_mut() = 1;
165 tx.savepoint("sp").unwrap();
166 *v.borrow_mut() = 2;
167 tx.rollback_to("sp").unwrap();
168 assert_eq!(*v.borrow(), 1);
169 tx.commit().unwrap();
170 }
171
172 #[test]
173 fn drop_without_commit_rolls_back() {
174 let v = Rc::new(RefCell::new(10));
175 {
176 let src = CounterSource { v: v.clone() };
177 let _tx = InMemoryTransaction::begin(src);
178 *v.borrow_mut() = 50;
179 }
180 assert_eq!(*v.borrow(), 10);
181 }
182
183 #[test]
184 fn unknown_savepoint_errors() {
185 let v = Rc::new(RefCell::new(0));
186 let src = CounterSource { v };
187 let mut tx = InMemoryTransaction::begin(src);
188 let err = tx.rollback_to("missing").unwrap_err();
189 assert!(matches!(err, TransactionError::UnknownSavepoint(_)));
190 tx.commit().unwrap();
191 }
192}