soma_mcp_client/upstream/relay/
cache.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::time::{Duration, Instant};
3
4use super::{RelayCapabilities, RelaySessionId};
5
6#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
7pub struct RelayCacheKey {
8    pub upstream: String,
9    pub session_id: RelaySessionId,
10    pub subject: Option<String>,
11}
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct RelayConnection {
15    pub key: RelayCacheKey,
16    pub created_at: Instant,
17    pub last_used: Instant,
18    pub alive: bool,
19    pub capabilities: RelayCapabilities,
20}
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub enum RelayConnectSlot {
24    CacheHit,
25    Leader,
26    Waiter,
27}
28
29#[derive(Debug, Clone)]
30pub struct RelayCache {
31    entries: BTreeMap<RelayCacheKey, RelayConnection>,
32    connect_locks: BTreeSet<RelayCacheKey>,
33    ttl: Duration,
34    capacity: usize,
35    pending_shutdown: Vec<RelayConnection>,
36}
37
38impl RelayCache {
39    #[must_use]
40    pub fn new(ttl: Duration, capacity: usize) -> Self {
41        Self {
42            entries: BTreeMap::new(),
43            connect_locks: BTreeSet::new(),
44            ttl,
45            capacity: capacity.max(1),
46            pending_shutdown: Vec::new(),
47        }
48    }
49
50    pub fn begin_connect(&mut self, key: RelayCacheKey, now: Instant) -> RelayConnectSlot {
51        self.sweep(now);
52        if let Some(entry) = self.entries.get_mut(&key) {
53            entry.last_used = now;
54            return RelayConnectSlot::CacheHit;
55        }
56        if !self.connect_locks.insert(key) {
57            return RelayConnectSlot::Waiter;
58        }
59        RelayConnectSlot::Leader
60    }
61
62    pub fn complete_connect(&mut self, connection: RelayConnection) {
63        self.connect_locks.remove(&connection.key);
64        self.entries.insert(connection.key.clone(), connection);
65        self.evict_over_capacity();
66    }
67
68    pub fn leader_cancelled(&mut self, key: &RelayCacheKey) {
69        self.connect_locks.remove(key);
70    }
71
72    pub fn connect_failed(&mut self, key: &RelayCacheKey) {
73        self.connect_locks.remove(key);
74    }
75
76    pub fn waiter_cancelled(&mut self, _key: &RelayCacheKey) {}
77
78    pub fn mark_dead(&mut self, key: &RelayCacheKey) {
79        if let Some(entry) = self.entries.get_mut(key) {
80            entry.alive = false;
81        }
82    }
83
84    pub fn sweep(&mut self, now: Instant) {
85        let expired: Vec<RelayCacheKey> = self
86            .entries
87            .iter()
88            .filter(|(_, entry)| !entry.alive || now.duration_since(entry.last_used) >= self.ttl)
89            .map(|(key, _)| key.clone())
90            .collect();
91        for key in expired {
92            if let Some(connection) = self.entries.remove(&key) {
93                self.pending_shutdown.push(connection);
94            }
95        }
96    }
97
98    pub fn lock_count(&self) -> usize {
99        self.connect_locks.len()
100    }
101
102    pub fn len(&self) -> usize {
103        self.entries.len()
104    }
105
106    pub fn is_empty(&self) -> bool {
107        self.entries.is_empty()
108    }
109
110    pub fn take_pending_shutdown(&mut self) -> Vec<RelayConnection> {
111        std::mem::take(&mut self.pending_shutdown)
112    }
113
114    fn evict_over_capacity(&mut self) {
115        while self.entries.len() > self.capacity {
116            let Some(lru_key) = self
117                .entries
118                .iter()
119                .min_by_key(|(_, entry)| entry.last_used)
120                .map(|(key, _)| key.clone())
121            else {
122                return;
123            };
124            if let Some(connection) = self.entries.remove(&lru_key) {
125                self.pending_shutdown.push(connection);
126            }
127        }
128    }
129}
130
131#[cfg(test)]
132#[path = "cache_tests.rs"]
133mod tests;
soma_mcp_client/upstream/relay/cache.rs

soma_mcp_client/upstream/relay/
cache.rs
