Skip to main content

alopex_embedded/
txn_manager.rs

1//! Transaction manager for CLI-driven KV transactions.
2
3use std::time::{Duration, SystemTime, UNIX_EPOCH};
4
5use alopex_core::kv::any::AnyKVTransaction;
6use alopex_core::kv::{RangeChangeJournalCapability, ReadAtCapability, ReadAtPoint, ReadAtResult};
7use alopex_core::{KVStore, KVTransaction, TxnMode};
8use alopex_sql::storage::{LocalRangeChangeJournal, RangeChangeJournalScope};
9
10use crate::{Database, Error, Result};
11
12const TXN_META_PREFIX: &[u8] = b"__alopex_txn_meta__:";
13const TXN_WRITE_PREFIX: &[u8] = b"__alopex_txn_write__:";
14const TXN_WRITE_DELETE: u8 = 0;
15const TXN_WRITE_PUT: u8 = 1;
16
17/// Metadata for a persisted KV transaction.
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct TransactionInfo {
20    /// Transaction identifier.
21    pub txn_id: String,
22    /// Persisted transaction start time.
23    pub started_at: SystemTime,
24    /// Timeout duration in seconds.
25    pub timeout_secs: u64,
26    /// Whether the transaction is expired.
27    pub is_expired: bool,
28}
29
30#[derive(Debug, Clone, Copy)]
31struct TxnMeta {
32    started_at_secs: u64,
33    timeout_secs: u64,
34}
35
36enum TxnWrite {
37    Put(Vec<u8>),
38    Delete,
39}
40
41/// Manages persisted KV transactions for CLI usage.
42pub struct TransactionManager;
43
44impl TransactionManager {
45    /// Checks the selected backend's read-at eligibility before a remote worker
46    /// attempts to open a fenced SQL session. This is distinct from persisted
47    /// CLI transaction metadata and never uses its wall-clock start time.
48    pub fn validate_read_at(db: &Database, point: &ReadAtPoint) -> ReadAtResult<()> {
49        db.store.read_at_capability().validate(point)
50    }
51
52    /// Returns the selected backend's explicit read-at capability.
53    pub fn read_at_capability(db: &Database) -> ReadAtCapability {
54        db.store.read_at_capability()
55    }
56
57    /// Begins a new transaction with the given timeout and returns its ID.
58    pub fn begin_with_timeout(db: &Database, timeout: Duration) -> Result<String> {
59        let txn_id = generate_txn_id();
60        let meta = TxnMeta {
61            started_at_secs: current_timestamp_secs(),
62            timeout_secs: timeout.as_secs(),
63        };
64        let mut txn = db.store.begin(TxnMode::ReadWrite).map_err(Error::Core)?;
65        txn.put(txn_meta_key(&txn_id), encode_meta(meta))
66            .map_err(Error::Core)?;
67        txn.commit_self().map_err(Error::Core)?;
68        Ok(txn_id)
69    }
70
71    /// Retrieves persisted transaction metadata.
72    pub fn get_info(db: &Database, txn_id: &str) -> Result<TransactionInfo> {
73        let mut txn = db.store.begin(TxnMode::ReadOnly).map_err(Error::Core)?;
74        let result = load_meta(&mut txn, txn_id).map(|meta| {
75            let started_at = UNIX_EPOCH + Duration::from_secs(meta.started_at_secs);
76            let is_expired = is_expired_from_meta(meta, current_timestamp_secs());
77            TransactionInfo {
78                txn_id: txn_id.to_string(),
79                started_at,
80                timeout_secs: meta.timeout_secs,
81                is_expired,
82            }
83        });
84        let commit_result = txn.commit_self().map_err(Error::Core);
85        match (result, commit_result) {
86            (Err(err), _) => Err(err),
87            (Ok(_), Err(err)) => Err(err),
88            (Ok(info), Ok(())) => Ok(info),
89        }
90    }
91
92    /// Checks whether a transaction has expired.
93    pub fn is_expired(db: &Database, txn_id: &str) -> Result<bool> {
94        let mut txn = db.store.begin(TxnMode::ReadOnly).map_err(Error::Core)?;
95        let result = load_meta(&mut txn, txn_id)
96            .map(|meta| is_expired_from_meta(meta, current_timestamp_secs()));
97        let commit_result = txn.commit_self().map_err(Error::Core);
98        match (result, commit_result) {
99            (Err(err), _) => Err(err),
100            (Ok(_), Err(err)) => Err(err),
101            (Ok(value), Ok(())) => Ok(value),
102        }
103    }
104
105    /// Retrieves a key within the specified transaction.
106    pub fn get(db: &Database, txn_id: &str, key: &[u8]) -> Result<Option<Vec<u8>>> {
107        let mut txn = db.store.begin(TxnMode::ReadOnly).map_err(Error::Core)?;
108        let result = (|| {
109            let _ = load_meta(&mut txn, txn_id)?;
110            if let Some(raw) = txn.get(&txn_write_key(txn_id, key)).map_err(Error::Core)? {
111                return Ok(match decode_write(txn_id, &raw)? {
112                    TxnWrite::Put(value) => Some(value),
113                    TxnWrite::Delete => None,
114                });
115            }
116            txn.get(&key.to_vec()).map_err(Error::Core)
117        })();
118        let commit_result = txn.commit_self().map_err(Error::Core);
119        match (result, commit_result) {
120            (Err(err), _) => Err(err),
121            (Ok(_), Err(err)) => Err(err),
122            (Ok(value), Ok(())) => Ok(value),
123        }
124    }
125
126    /// Stages a put operation within the specified transaction.
127    pub fn put(db: &Database, txn_id: &str, key: &[u8], value: &[u8]) -> Result<()> {
128        let mut txn = db.store.begin(TxnMode::ReadWrite).map_err(Error::Core)?;
129        let _ = load_meta(&mut txn, txn_id)?;
130        let encoded = encode_write(TxnWrite::Put(value.to_vec()));
131        txn.put(txn_write_key(txn_id, key), encoded)
132            .map_err(Error::Core)?;
133        txn.commit_self().map_err(Error::Core)?;
134        Ok(())
135    }
136
137    /// Stages a delete operation within the specified transaction.
138    pub fn delete(db: &Database, txn_id: &str, key: &[u8]) -> Result<()> {
139        let mut txn = db.store.begin(TxnMode::ReadWrite).map_err(Error::Core)?;
140        let _ = load_meta(&mut txn, txn_id)?;
141        let encoded = encode_write(TxnWrite::Delete);
142        txn.put(txn_write_key(txn_id, key), encoded)
143            .map_err(Error::Core)?;
144        txn.commit_self().map_err(Error::Core)?;
145        Ok(())
146    }
147
148    /// Commits staged writes and finalizes the transaction.
149    pub fn commit(db: &Database, txn_id: &str) -> Result<()> {
150        let mut txn = db.store.begin(TxnMode::ReadWrite).map_err(Error::Core)?;
151        let _ = load_meta(&mut txn, txn_id)?;
152        let journal = if db.store.range_change_journal_capability()
153            == RangeChangeJournalCapability::Supported
154        {
155            Some(
156                LocalRangeChangeJournal::capture(
157                    &mut txn,
158                    RangeChangeJournalScope::local(Default::default()),
159                )
160                .map_err(Error::Core)?,
161            )
162        } else {
163            None
164        };
165        let prefix = txn_write_prefix(txn_id);
166        let staged: Vec<(Vec<u8>, Vec<u8>)> =
167            txn.scan_prefix(&prefix).map_err(Error::Core)?.collect();
168        for (staged_key, raw) in &staged {
169            let user_key = extract_user_key(txn_id, staged_key)?;
170            match decode_write(txn_id, raw)? {
171                TxnWrite::Put(value) => {
172                    txn.put(user_key, value).map_err(Error::Core)?;
173                }
174                TxnWrite::Delete => {
175                    txn.delete(user_key).map_err(Error::Core)?;
176                }
177            }
178        }
179        for (staged_key, _) in staged {
180            txn.delete(staged_key).map_err(Error::Core)?;
181        }
182        txn.delete(txn_meta_key(txn_id)).map_err(Error::Core)?;
183        if let Some(journal) = journal {
184            journal.stage(&mut txn).map_err(Error::Core)?;
185        }
186        txn.commit_self().map_err(Error::Core)?;
187        Ok(())
188    }
189
190    /// Rolls back staged writes and removes transaction metadata.
191    pub fn rollback(db: &Database, txn_id: &str) -> Result<()> {
192        let mut txn = db.store.begin(TxnMode::ReadWrite).map_err(Error::Core)?;
193        let _ = load_meta(&mut txn, txn_id)?;
194        let prefix = txn_write_prefix(txn_id);
195        let staged: Vec<(Vec<u8>, Vec<u8>)> =
196            txn.scan_prefix(&prefix).map_err(Error::Core)?.collect();
197        for (staged_key, _) in staged {
198            txn.delete(staged_key).map_err(Error::Core)?;
199        }
200        txn.delete(txn_meta_key(txn_id)).map_err(Error::Core)?;
201        txn.commit_self().map_err(Error::Core)?;
202        Ok(())
203    }
204}
205
206fn current_timestamp_secs() -> u64 {
207    SystemTime::now()
208        .duration_since(UNIX_EPOCH)
209        .unwrap_or_default()
210        .as_secs()
211}
212
213fn generate_txn_id() -> String {
214    let nanos = SystemTime::now()
215        .duration_since(UNIX_EPOCH)
216        .unwrap_or_default()
217        .as_nanos();
218    format!("txn-{}-{}", nanos, std::process::id())
219}
220
221fn txn_meta_key(txn_id: &str) -> Vec<u8> {
222    let mut key = Vec::with_capacity(TXN_META_PREFIX.len() + txn_id.len());
223    key.extend_from_slice(TXN_META_PREFIX);
224    key.extend_from_slice(txn_id.as_bytes());
225    key
226}
227
228fn txn_write_prefix(txn_id: &str) -> Vec<u8> {
229    let mut key = Vec::with_capacity(TXN_WRITE_PREFIX.len() + txn_id.len() + 1);
230    key.extend_from_slice(TXN_WRITE_PREFIX);
231    key.extend_from_slice(txn_id.as_bytes());
232    key.push(b':');
233    key
234}
235
236fn txn_write_key(txn_id: &str, key: &[u8]) -> Vec<u8> {
237    let mut full = txn_write_prefix(txn_id);
238    full.extend_from_slice(key);
239    full
240}
241
242fn encode_meta(meta: TxnMeta) -> Vec<u8> {
243    let mut payload = Vec::with_capacity(16);
244    payload.extend_from_slice(&meta.started_at_secs.to_le_bytes());
245    payload.extend_from_slice(&meta.timeout_secs.to_le_bytes());
246    payload
247}
248
249fn decode_meta(txn_id: &str, raw: &[u8]) -> Result<TxnMeta> {
250    if raw.len() < 16 {
251        return Err(Error::Core(alopex_core::Error::InvalidFormat(format!(
252            "transaction metadata invalid: {}",
253            txn_id
254        ))));
255    }
256    let started_at_secs = u64::from_le_bytes(raw[0..8].try_into().unwrap());
257    let timeout_secs = u64::from_le_bytes(raw[8..16].try_into().unwrap());
258    Ok(TxnMeta {
259        started_at_secs,
260        timeout_secs,
261    })
262}
263
264fn load_meta(txn: &mut AnyKVTransaction<'_>, txn_id: &str) -> Result<TxnMeta> {
265    let Some(raw) = txn.get(&txn_meta_key(txn_id)).map_err(Error::Core)? else {
266        return Err(Error::InvalidTransactionId(txn_id.to_string()));
267    };
268    decode_meta(txn_id, &raw)
269}
270
271fn is_expired_from_meta(meta: TxnMeta, now_secs: u64) -> bool {
272    now_secs.saturating_sub(meta.started_at_secs) >= meta.timeout_secs
273}
274
275fn encode_write(entry: TxnWrite) -> Vec<u8> {
276    match entry {
277        TxnWrite::Put(value) => {
278            let mut payload = Vec::with_capacity(1 + value.len());
279            payload.push(TXN_WRITE_PUT);
280            payload.extend_from_slice(&value);
281            payload
282        }
283        TxnWrite::Delete => vec![TXN_WRITE_DELETE],
284    }
285}
286
287fn decode_write(txn_id: &str, raw: &[u8]) -> Result<TxnWrite> {
288    let Some((&tag, rest)) = raw.split_first() else {
289        return Err(Error::Core(alopex_core::Error::InvalidFormat(format!(
290            "transaction write entry invalid: {}",
291            txn_id
292        ))));
293    };
294    match tag {
295        TXN_WRITE_PUT => Ok(TxnWrite::Put(rest.to_vec())),
296        TXN_WRITE_DELETE => Ok(TxnWrite::Delete),
297        _ => Err(Error::Core(alopex_core::Error::InvalidFormat(format!(
298            "transaction write entry invalid: {}",
299            txn_id
300        )))),
301    }
302}
303
304fn extract_user_key(txn_id: &str, staged_key: &[u8]) -> Result<Vec<u8>> {
305    let prefix = txn_write_prefix(txn_id);
306    if !staged_key.starts_with(&prefix) {
307        return Err(Error::Core(alopex_core::Error::InvalidFormat(format!(
308            "transaction write key invalid: {}",
309            txn_id
310        ))));
311    }
312    Ok(staged_key[prefix.len()..].to_vec())
313}
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318    use alopex_core::kv::decode_range_change;
319    use alopex_sql::storage::KeyEncoder;
320
321    fn create_test_db() -> Database {
322        Database::open_in_memory().unwrap()
323    }
324
325    #[test]
326    fn test_txn_put_get_commit() {
327        let db = create_test_db();
328        let txn_id = TransactionManager::begin_with_timeout(&db, Duration::from_secs(60)).unwrap();
329
330        TransactionManager::put(&db, &txn_id, b"alpha", b"beta").unwrap();
331        let value = TransactionManager::get(&db, &txn_id, b"alpha").unwrap();
332        assert_eq!(value, Some(b"beta".to_vec()));
333
334        TransactionManager::commit(&db, &txn_id).unwrap();
335
336        let mut verify_txn = db.begin(TxnMode::ReadOnly).unwrap();
337        let stored = verify_txn.get(b"alpha").unwrap();
338        verify_txn.commit().unwrap();
339        assert_eq!(stored, Some(b"beta".to_vec()));
340
341        let err = TransactionManager::get(&db, &txn_id, b"alpha").unwrap_err();
342        assert!(matches!(err, Error::InvalidTransactionId(_)));
343    }
344
345    #[test]
346    fn test_txn_rollback_discards_writes() {
347        let db = create_test_db();
348        let txn_id = TransactionManager::begin_with_timeout(&db, Duration::from_secs(60)).unwrap();
349
350        TransactionManager::put(&db, &txn_id, b"key", b"value").unwrap();
351        TransactionManager::rollback(&db, &txn_id).unwrap();
352
353        let mut verify_txn = db.begin(TxnMode::ReadOnly).unwrap();
354        let stored = verify_txn.get(b"key").unwrap();
355        verify_txn.commit().unwrap();
356        assert!(stored.is_none());
357
358        let err = TransactionManager::get(&db, &txn_id, b"key").unwrap_err();
359        assert!(matches!(err, Error::InvalidTransactionId(_)));
360    }
361
362    #[test]
363    fn test_txn_delete_marks_missing() {
364        let db = create_test_db();
365        {
366            let mut seed = db.begin(TxnMode::ReadWrite).unwrap();
367            seed.put(b"drop-me", b"payload").unwrap();
368            seed.commit().unwrap();
369        }
370
371        let txn_id = TransactionManager::begin_with_timeout(&db, Duration::from_secs(60)).unwrap();
372        TransactionManager::delete(&db, &txn_id, b"drop-me").unwrap();
373        let value = TransactionManager::get(&db, &txn_id, b"drop-me").unwrap();
374        assert!(value.is_none());
375        TransactionManager::rollback(&db, &txn_id).unwrap();
376    }
377
378    #[test]
379    fn test_txn_is_expired() {
380        let db = create_test_db();
381        let txn_id = TransactionManager::begin_with_timeout(&db, Duration::from_secs(0)).unwrap();
382        assert!(TransactionManager::is_expired(&db, &txn_id).unwrap());
383    }
384
385    #[test]
386    fn failed_commit_can_be_replayed_without_duplicate_journal() {
387        let db = create_test_db();
388        let txn_id = TransactionManager::begin_with_timeout(&db, Duration::from_secs(60)).unwrap();
389        let row_key = KeyEncoder::row_key(1, 1);
390        TransactionManager::put(&db, &txn_id, &row_key, b"row").unwrap();
391
392        db.set_memory_limit(Some(0));
393        assert!(TransactionManager::commit(&db, &txn_id).is_err());
394        db.set_memory_limit(None);
395        TransactionManager::commit(&db, &txn_id).unwrap();
396
397        let mut reader = db.begin(TxnMode::ReadOnly).unwrap();
398        assert_eq!(reader.get(&row_key).unwrap(), Some(b"row".to_vec()));
399        reader.commit().unwrap();
400        let journal_records = db
401            .snapshot()
402            .into_iter()
403            .filter_map(|(_, value)| decode_range_change(&value).ok())
404            .collect::<Vec<_>>();
405        assert_eq!(journal_records.len(), 1);
406        assert_eq!(journal_records[0].epoch, 1);
407    }
408}