bamboo_engine/
session_cache.rs

1//! Lock-free session snapshots. Durable read/modify/write transactions remain
2//! the responsibility of SessionRepository; readers never join their queue.
3
4use std::ops::Deref;
5use std::sync::Arc;
6
7use arc_swap::ArcSwap;
8use bamboo_agent_core::Session;
9use crossbeam_skiplist::SkipMap;
10
11pub type SessionCache = Arc<SessionCacheMap>;
12
13/// An immutable version can remain in use while a writer publishes its successor.
14pub struct SessionSnapshot(ArcSwap<Session>);
15
16impl SessionSnapshot {
17    pub fn new(session: Session) -> Self {
18        Self(ArcSwap::from_pointee(session))
19    }
20
21    pub fn read(&self) -> SessionRead {
22        SessionRead(self.0.load_full())
23    }
24
25    /// Apply a narrow patch to the latest version. The closure may be retried
26    /// after a concurrent publication, so it must have no external side effects.
27    pub fn update(&self, mutate: impl Fn(&mut Session)) {
28        self.0.rcu(|current| {
29            let mut next = (**current).clone();
30            mutate(&mut next);
31            next
32        });
33    }
34}
35
36/// Deliberately does not implement Clone: `read().clone()` clones the Session,
37/// preserving the existing detached-snapshot read contract.
38pub struct SessionRead(Arc<Session>);
39
40impl Deref for SessionRead {
41    type Target = Session;
42    fn deref(&self) -> &Session {
43        &self.0
44    }
45}
46
47/// Both the index and the values use atomic publication. Unlike a DashMap
48/// guard, a returned entry does not hold a shard lock while cloning a transcript.
49#[derive(Default)]
50pub struct SessionCacheMap {
51    entries: SkipMap<String, Arc<SessionSnapshot>>,
52}
53
54pub struct SessionCacheEntry {
55    key: String,
56    value: Arc<SessionSnapshot>,
57}
58
59impl SessionCacheEntry {
60    pub fn key(&self) -> &String {
61        &self.key
62    }
63    pub fn value(&self) -> &Arc<SessionSnapshot> {
64        &self.value
65    }
66}
67
68impl Deref for SessionCacheEntry {
69    type Target = Arc<SessionSnapshot>;
70    fn deref(&self) -> &Self::Target {
71        &self.value
72    }
73}
74
75impl SessionCacheMap {
76    pub fn new() -> Self {
77        Self::default()
78    }
79
80    pub fn insert(&self, id: String, snapshot: Arc<SessionSnapshot>) {
81        // Keep the slot stable until eviction. A concurrent narrow update must
82        // retry against this publication, not succeed on an orphaned old Arc.
83        let entry = self.entries.get_or_insert(id, snapshot.clone());
84        if !Arc::ptr_eq(entry.value(), &snapshot) {
85            entry.value().0.store(snapshot.0.load_full());
86        }
87    }
88
89    pub fn get(&self, id: &str) -> Option<SessionCacheEntry> {
90        self.entries.get(id).map(|entry| SessionCacheEntry {
91            key: entry.key().clone(),
92            value: entry.value().clone(),
93        })
94    }
95
96    pub fn remove(&self, id: &str) -> Option<(String, Arc<SessionSnapshot>)> {
97        self.entries
98            .remove(id)
99            .map(|entry| (entry.key().clone(), entry.value().clone()))
100    }
101
102    pub fn iter(&self) -> impl Iterator<Item = SessionCacheEntry> + '_ {
103        self.entries.iter().map(|entry| SessionCacheEntry {
104            key: entry.key().clone(),
105            value: entry.value().clone(),
106        })
107    }
108
109    pub fn contains_key(&self, id: &str) -> bool {
110        self.entries.contains_key(id)
111    }
112
113    pub fn len(&self) -> usize {
114        self.entries.len()
115    }
116
117    pub fn is_empty(&self) -> bool {
118        self.entries.is_empty()
119    }
120
121    pub fn clear(&self) {
122        self.entries.clear();
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn retained_reader_does_not_block_publication_or_change_its_version() {
132        let snapshot = SessionSnapshot::new(Session::new("s", "model"));
133        let previous = snapshot.read();
134        snapshot.update(|session| {
135            session.metadata.insert("committed".into(), "yes".into());
136        });
137        assert!(!previous.metadata.contains_key("committed"));
138        assert_eq!(snapshot.read().metadata["committed"], "yes");
139    }
140
141    #[test]
142    fn concurrent_narrow_updates_preserve_every_writer() {
143        let snapshot = SessionSnapshot::new(Session::new("shared-root", "model"));
144        std::thread::scope(|scope| {
145            for worker in 0..16 {
146                let snapshot = &snapshot;
147                scope.spawn(move || {
148                    for item in 0..32 {
149                        snapshot.update(|session| {
150                            session
151                                .metadata
152                                .insert(format!("{worker}/{item}"), "yes".into());
153                        });
154                    }
155                });
156            }
157        });
158        assert_eq!(snapshot.read().metadata.len(), 512);
159    }
160
161    #[test]
162    fn narrow_update_retries_when_full_snapshot_is_published_to_same_slot() {
163        use std::sync::{
164            atomic::{AtomicBool, Ordering},
165            Barrier,
166        };
167        let cache = SessionCacheMap::new();
168        cache.insert(
169            "root".into(),
170            Arc::new(SessionSnapshot::new(Session::new("root", "old"))),
171        );
172        let entered = Barrier::new(2);
173        let release = Barrier::new(2);
174        std::thread::scope(|scope| {
175            let updater = scope.spawn(|| {
176                let first = AtomicBool::new(true);
177                cache.get("root").unwrap().update(|session| {
178                    // Deterministic test seam: stop only the first CAS attempt.
179                    if first.swap(false, Ordering::SeqCst) {
180                        entered.wait();
181                        release.wait();
182                    }
183                    session
184                        .metadata
185                        .insert("narrow-patch".into(), "kept".into());
186                });
187            });
188            entered.wait();
189            cache.insert(
190                "root".into(),
191                Arc::new(SessionSnapshot::new(Session::new("root", "new"))),
192            );
193            release.wait();
194            updater.join().unwrap();
195        });
196        let actual = cache.get("root").unwrap().read();
197        assert_eq!(actual.model, "new");
198        assert_eq!(actual.metadata["narrow-patch"], "kept");
199    }
200
201    #[test]
202    fn hundreds_of_independent_sessions_publish_and_remove_with_retained_reads() {
203        let cache = SessionCacheMap::new();
204        std::thread::scope(|scope| {
205            for worker in 0..16 {
206                let cache = &cache;
207                scope.spawn(move || {
208                    for item in 0..32 {
209                        let id = format!("child-{worker}-{item}");
210                        cache.insert(
211                            id.clone(),
212                            Arc::new(SessionSnapshot::new(Session::new(&id, "m"))),
213                        );
214                        let held = cache.get(&id).unwrap();
215                        let old = held.read();
216                        held.update(|session| {
217                            session.metadata.insert("done".into(), "yes".into());
218                        });
219                        assert_eq!(cache.get(&id).unwrap().read().metadata["done"], "yes");
220                        cache.remove(&id).unwrap();
221                        assert!(cache.get(&id).is_none());
222                        assert!(!old.metadata.contains_key("done"));
223                    }
224                });
225            }
226        });
227        assert!(cache.is_empty());
228    }
229}
bamboo_engine/session_cache.rs

bamboo_engine/
session_cache.rs
