Skip to main content

mentedb_core/
mvcc.rs

1//! MVCC Version Tracking — simple multi-version concurrency control for memories.
2
3use std::collections::HashMap;
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::time::{SystemTime, UNIX_EPOCH};
6
7use serde::{Deserialize, Serialize};
8
9use crate::types::{AgentId, MemoryId, Timestamp};
10
11/// A single versioned snapshot of a memory.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct Version {
14    pub version_id: u64,
15    pub memory_id: MemoryId,
16    pub agent_id: AgentId,
17    pub timestamp: Timestamp,
18    pub data_hash: u64,
19}
20
21/// Tracks version history for all memories.
22#[derive(Debug)]
23pub struct VersionStore {
24    versions: HashMap<MemoryId, Vec<Version>>,
25    next_version: AtomicU64,
26}
27
28impl Default for VersionStore {
29    fn default() -> Self {
30        Self {
31            versions: HashMap::new(),
32            next_version: AtomicU64::new(1),
33        }
34    }
35}
36
37fn now_micros() -> Timestamp {
38    SystemTime::now()
39        .duration_since(UNIX_EPOCH)
40        .unwrap_or_default()
41        .as_micros() as Timestamp
42}
43
44impl VersionStore {
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    /// Record a write and return the new version ID.
50    pub fn record_write(&mut self, memory_id: MemoryId, agent_id: AgentId, data_hash: u64) -> u64 {
51        let vid = self.next_version.fetch_add(1, Ordering::Relaxed);
52        let version = Version {
53            version_id: vid,
54            memory_id,
55            agent_id,
56            timestamp: now_micros(),
57            data_hash,
58        };
59        self.versions.entry(memory_id).or_default().push(version);
60        vid
61    }
62
63    /// Get the latest version for a memory.
64    pub fn get_latest(&self, memory_id: MemoryId) -> Option<&Version> {
65        self.versions.get(&memory_id).and_then(|v| v.last())
66    }
67
68    /// Get the full version history for a memory.
69    pub fn get_history(&self, memory_id: MemoryId) -> Vec<&Version> {
70        self.versions
71            .get(&memory_id)
72            .map(|v| v.iter().collect())
73            .unwrap_or_default()
74    }
75
76    /// Get the version that was current at or before `timestamp`.
77    pub fn get_version_at(&self, memory_id: MemoryId, timestamp: Timestamp) -> Option<&Version> {
78        self.versions
79            .get(&memory_id)
80            .and_then(|v| v.iter().rev().find(|ver| ver.timestamp <= timestamp))
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87    use uuid::Uuid;
88
89    #[test]
90    fn record_and_get_latest() {
91        let mut store = VersionStore::new();
92        let mid = Uuid::new_v4();
93        let aid = Uuid::new_v4();
94        store.record_write(mid, aid, 111);
95        store.record_write(mid, aid, 222);
96        assert_eq!(store.get_latest(mid).unwrap().data_hash, 222);
97    }
98
99    #[test]
100    fn version_ids_increment() {
101        let mut store = VersionStore::new();
102        let mid = Uuid::new_v4();
103        let aid = Uuid::new_v4();
104        let v1 = store.record_write(mid, aid, 1);
105        let v2 = store.record_write(mid, aid, 2);
106        assert_eq!(v2, v1 + 1);
107    }
108
109    #[test]
110    fn get_history() {
111        let mut store = VersionStore::new();
112        let mid = Uuid::new_v4();
113        let aid = Uuid::new_v4();
114        store.record_write(mid, aid, 10);
115        store.record_write(mid, aid, 20);
116        store.record_write(mid, aid, 30);
117        assert_eq!(store.get_history(mid).len(), 3);
118    }
119
120    #[test]
121    fn get_version_at() {
122        let mut store = VersionStore::new();
123        let mid = Uuid::new_v4();
124        let aid = Uuid::new_v4();
125        store.record_write(mid, aid, 1);
126        // The latest write should be findable at a very large timestamp.
127        let ver = store.get_version_at(mid, u64::MAX);
128        assert!(ver.is_some());
129    }
130
131    #[test]
132    fn empty_history() {
133        let store = VersionStore::new();
134        assert!(store.get_latest(Uuid::new_v4()).is_none());
135        assert!(store.get_history(Uuid::new_v4()).is_empty());
136    }
137}