soma_fleet/
cache.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
3
4use crate::{HostId, HostRecord, PoolKey, TopologySnapshot};
5
6/// Thread-safe connection cache bound to exact host topology revisions.
7///
8/// The cache does not open or close connections. Drivers own those lifecycle
9/// actions and receive removed handles from invalidation and eviction methods.
10pub struct ConnectionCache<C> {
11    entries: RwLock<BTreeMap<PoolKey, Arc<C>>>,
12}
13
14impl<C> Default for ConnectionCache<C> {
15    fn default() -> Self {
16        Self {
17            entries: RwLock::new(BTreeMap::new()),
18        }
19    }
20}
21
22impl<C> ConnectionCache<C> {
23    /// Creates an empty cache.
24    #[must_use]
25    pub fn new() -> Self {
26        Self::default()
27    }
28
29    /// Inserts a connection for one exact host revision.
30    ///
31    /// Returns the replaced handle when the same key already existed.
32    pub fn insert(&self, host: &HostRecord, connection: C) -> Option<Arc<C>> {
33        write_unpoisoned(&self.entries).insert(host.pool_key(), Arc::new(connection))
34    }
35
36    /// Returns a connection only when the host revision matches exactly.
37    #[must_use]
38    pub fn get(&self, host: &HostRecord) -> Option<Arc<C>> {
39        read_unpoisoned(&self.entries)
40            .get(&host.pool_key())
41            .map(Arc::clone)
42    }
43
44    /// Removes one exact host revision.
45    pub fn remove(&self, host: &HostRecord) -> Option<Arc<C>> {
46        write_unpoisoned(&self.entries).remove(&host.pool_key())
47    }
48
49    /// Removes every cached revision for one host.
50    pub fn invalidate_host(&self, host: &HostId) -> Vec<Arc<C>> {
51        let mut entries = write_unpoisoned(&self.entries);
52        let keys = entries
53            .keys()
54            .filter(|key| key.host() == host)
55            .cloned()
56            .collect::<Vec<_>>();
57        keys.into_iter()
58            .filter_map(|key| entries.remove(&key))
59            .collect()
60    }
61
62    /// Removes keys absent from the supplied topology snapshot.
63    ///
64    /// A host whose endpoint changed has a new revision and therefore evicts
65    /// the old cached connection even when its stable host identity is unchanged.
66    pub fn retain_snapshot(&self, snapshot: &TopologySnapshot) -> Vec<Arc<C>> {
67        let current = snapshot
68            .hosts()
69            .map(HostRecord::pool_key)
70            .collect::<BTreeSet<_>>();
71        let mut entries = write_unpoisoned(&self.entries);
72        let stale = entries
73            .keys()
74            .filter(|key| !current.contains(*key))
75            .cloned()
76            .collect::<Vec<_>>();
77        stale
78            .into_iter()
79            .filter_map(|key| entries.remove(&key))
80            .collect()
81    }
82
83    /// Returns the number of cached revision keys.
84    #[must_use]
85    pub fn len(&self) -> usize {
86        read_unpoisoned(&self.entries).len()
87    }
88
89    /// Returns whether the cache has no entries.
90    #[must_use]
91    pub fn is_empty(&self) -> bool {
92        read_unpoisoned(&self.entries).is_empty()
93    }
94
95    /// Returns sorted cached keys for metrics or diagnostics.
96    #[must_use]
97    pub fn keys(&self) -> Vec<PoolKey> {
98        read_unpoisoned(&self.entries).keys().cloned().collect()
99    }
100}
101
102fn read_unpoisoned<T>(lock: &RwLock<T>) -> RwLockReadGuard<'_, T> {
103    lock.read()
104        .unwrap_or_else(std::sync::PoisonError::into_inner)
105}
106
107fn write_unpoisoned<T>(lock: &RwLock<T>) -> RwLockWriteGuard<'_, T> {
108    lock.write()
109        .unwrap_or_else(std::sync::PoisonError::into_inner)
110}
111
112#[cfg(test)]
113#[path = "cache_tests.rs"]
114mod tests;
soma_fleet/cache.rs

soma_fleet/
cache.rs
