use std::collections::HashMap;
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;
use thiserror::Error;
use crate::application::session::{Clock, IdSource};
use crate::domain::mcp::McpSessionId;
use crate::domain::session::{Instant, SessionId};
pub struct SessionStore<S> {
sessions: Mutex<HashMap<McpSessionId, Arc<S>>>,
last_touched: Option<Mutex<HashMap<McpSessionId, Instant>>>,
ids: Arc<dyn IdSource + Send + Sync>,
clock: Option<Arc<dyn Clock + Send + Sync>>,
}
type ReaperPair<'a> = (
&'a Mutex<HashMap<McpSessionId, Instant>>,
&'a Arc<dyn Clock + Send + Sync>,
);
impl<S> SessionStore<S> {
#[must_use]
pub fn new(ids: Arc<dyn IdSource + Send + Sync>) -> Self {
Self {
sessions: Mutex::new(HashMap::new()),
last_touched: None,
ids,
clock: None,
}
}
#[must_use]
pub fn with_idle_reaper(
ids: Arc<dyn IdSource + Send + Sync>,
clock: Arc<dyn Clock + Send + Sync>,
) -> Self {
Self {
sessions: Mutex::new(HashMap::new()),
last_touched: Some(Mutex::new(HashMap::new())),
ids,
clock: Some(clock),
}
}
pub fn open(&self, session: Arc<S>) -> Result<McpSessionId, SessionStoreError> {
let id = McpSessionId::new(self.ids.next_session_id());
self.insert(id, session)?;
Ok(id)
}
pub fn open_with_id(
&self,
session: Arc<S>,
id: SessionId,
) -> Result<McpSessionId, SessionStoreError> {
let id = McpSessionId::new(id);
self.insert(id, session)?;
Ok(id)
}
pub fn next_id(&self) -> SessionId {
self.ids.next_session_id()
}
fn insert(&self, id: McpSessionId, session: Arc<S>) -> Result<(), SessionStoreError> {
{
let mut sessions = self.locked()?;
if sessions.contains_key(&id) {
return Err(SessionStoreError::DuplicateId);
}
sessions.insert(id, session);
}
self.record_touched(id)?;
Ok(())
}
pub fn lookup(&self, id: &McpSessionId) -> Result<Option<Arc<S>>, SessionStoreError> {
let found = self.locked()?.get(id).cloned();
if found.is_some() {
self.record_touched(*id)?;
}
Ok(found)
}
pub fn terminate(&self, id: &McpSessionId) -> Result<Option<Arc<S>>, SessionStoreError> {
let removed = self.locked()?.remove(id);
if removed.is_some() {
self.forget_touched(*id)?;
}
Ok(removed)
}
pub fn reap_idle(&self, max_age: Duration) -> Result<Vec<McpSessionId>, SessionStoreError> {
let Some((touched, clock)) = self.reaper_pair() else {
return Ok(Vec::new());
};
let now = clock.now();
let mut stale: Vec<McpSessionId> = Vec::new();
{
let touched = touched.lock().map_err(|_| SessionStoreError::Unavailable)?;
for (id, when) in touched.iter() {
if now.as_duration().saturating_sub(when.as_duration()) > max_age {
stale.push(*id);
}
}
}
let mut reaped: Vec<McpSessionId> = Vec::with_capacity(stale.len());
{
let mut sessions = self.locked()?;
for id in stale {
if sessions.remove(&id).is_some() {
reaped.push(id);
}
}
}
{
let mut touched = touched.lock().map_err(|_| SessionStoreError::Unavailable)?;
for id in &reaped {
touched.remove(id);
}
}
Ok(reaped)
}
pub fn len(&self) -> Result<usize, SessionStoreError> {
Ok(self.locked()?.len())
}
pub fn is_empty(&self) -> Result<bool, SessionStoreError> {
Ok(self.locked()?.is_empty())
}
fn locked(&self) -> Result<MutexGuard<'_, HashMap<McpSessionId, Arc<S>>>, SessionStoreError> {
self.sessions
.lock()
.map_err(|_| SessionStoreError::Unavailable)
}
fn reaper_pair(&self) -> Option<ReaperPair<'_>> {
let (touched, clock) = match (&self.last_touched, &self.clock) {
(Some(t), Some(c)) => (t, c),
(None, None) => return None,
_ => unreachable!("last_touched and clock must agree"),
};
Some((touched, clock))
}
fn record_touched(&self, id: McpSessionId) -> Result<(), SessionStoreError> {
let Some((touched, clock)) = self.reaper_pair() else {
return Ok(());
};
{
let mut touched = touched.lock().map_err(|_| SessionStoreError::Unavailable)?;
touched.insert(id, clock.now());
}
Ok(())
}
fn forget_touched(&self, id: McpSessionId) -> Result<(), SessionStoreError> {
let Some((touched, _)) = self.reaper_pair() else {
return Ok(());
};
{
let mut touched = touched.lock().map_err(|_| SessionStoreError::Unavailable)?;
touched.remove(&id);
}
Ok(())
}
}
impl<S> core::fmt::Debug for SessionStore<S> {
fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
formatter
.debug_struct("SessionStore")
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum SessionStoreError {
#[error("the session store's lock is poisoned")]
Unavailable,
#[error("the session id source returned an id that is already in use")]
DuplicateId,
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use super::{SessionStore, SessionStoreError};
use crate::application::session::{Clock, IdSource};
use crate::domain::mcp::{McpSessionId, VISIBLE_ASCII_RANGE};
use crate::domain::session::{Instant, SessionId};
struct Counting(AtomicU64);
impl Counting {
fn new() -> Self {
Self(AtomicU64::new(0))
}
fn shared() -> Arc<dyn IdSource + Send + Sync> {
Arc::new(Self::new())
}
fn issued(&self) -> u64 {
self.0.load(Ordering::SeqCst)
}
}
impl IdSource for Counting {
fn next_session_id(&self) -> SessionId {
SessionId::new(self.0.fetch_add(1, Ordering::SeqCst))
}
}
struct Stuck(SessionId);
impl IdSource for Stuck {
fn next_session_id(&self) -> SessionId {
self.0
}
}
type TestError = Box<dyn std::error::Error>;
#[derive(Debug)]
enum ThreadFailure {
Store(SessionStoreError),
Vanished,
}
impl From<SessionStoreError> for ThreadFailure {
fn from(error: SessionStoreError) -> Self {
Self::Store(error)
}
}
impl core::fmt::Display for ThreadFailure {
fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Store(error) => write!(formatter, "the store refused a lookup: {error}"),
Self::Vanished => write!(formatter, "a live session was not found"),
}
}
}
impl std::error::Error for ThreadFailure {}
#[test]
fn the_error_variants_are_distinguishable() {
assert_ne!(
SessionStoreError::Unavailable.to_string(),
SessionStoreError::DuplicateId.to_string()
);
assert!(
SessionStoreError::DuplicateId
.to_string()
.contains("already in use")
);
}
#[test]
fn a_minted_id_round_trips_through_the_header_parser() -> Result<(), TestError> {
let store = SessionStore::new(Counting::shared());
let id = store.open(Arc::new(42usize))?;
assert_eq!(id.to_string().parse::<McpSessionId>(), Ok(id));
Ok(())
}
#[test]
fn the_header_value_is_fixed_width_and_visible_ascii() -> Result<(), TestError> {
let store = SessionStore::new(Counting::shared());
for _ in 0..4 {
let id = store.open(Arc::new(()))?;
let rendered = id.to_string();
assert_eq!(
rendered.len(),
16,
"session id {rendered:?} is not sixteen characters"
);
assert!(
rendered
.bytes()
.all(|byte| VISIBLE_ASCII_RANGE.contains(&byte)),
"session id {rendered:?} carries a byte no header may contain"
);
assert!(
rendered.bytes().all(|byte| byte.is_ascii_hexdigit()),
"session id {rendered:?} is not hexadecimal"
);
}
Ok(())
}
#[test]
fn ids_come_from_the_injected_source() -> Result<(), TestError> {
let source = Arc::new(Counting::new());
let store = SessionStore::new(source.clone());
let first = store.open(Arc::new(()))?;
let second = store.open(Arc::new(()))?;
assert_eq!(first, McpSessionId::new(SessionId::new(0)));
assert_eq!(second, McpSessionId::new(SessionId::new(1)));
assert_eq!(source.issued(), 2, "the store asked for a different count");
Ok(())
}
#[test]
fn a_miss_and_a_malformed_header_are_different_outcomes() -> Result<(), TestError> {
let store = SessionStore::new(Counting::shared());
let live = store.open(Arc::new(()))?;
store.terminate(&live)?;
assert_eq!(store.lookup(&live)?, None);
assert!("not-a-session".parse::<McpSessionId>().is_err());
assert!("".parse::<McpSessionId>().is_err());
Ok(())
}
#[test]
fn terminating_an_absent_session_is_a_miss() -> Result<(), TestError> {
let store = SessionStore::new(Counting::shared());
let id = store.open(Arc::new(()))?;
assert!(store.terminate(&id)?.is_some());
assert_eq!(store.terminate(&id)?, None);
assert!(store.is_empty()?);
Ok(())
}
#[test]
fn a_repeated_id_refuses_instead_of_overwriting() -> Result<(), TestError> {
let store = SessionStore::new(Arc::new(Stuck(SessionId::new(7))));
let first = store.open(Arc::new("first"))?;
let second = store.open(Arc::new("second"));
assert_eq!(second, Err(SessionStoreError::DuplicateId));
assert_eq!(store.len()?, 1);
let survivor = store.lookup(&first)?;
assert_eq!(survivor.as_deref(), Some(&"first"));
Ok(())
}
#[test]
fn concurrent_lookups_never_block_one_another() -> Result<(), TestError> {
const SESSIONS: usize = 8;
const THREADS: usize = 4;
let store = Arc::new(SessionStore::new(Counting::shared()));
let mut ids = Vec::new();
for _ in 0..SESSIONS {
ids.push(store.open(Arc::new(SESSIONS))?);
}
let threads: Vec<_> = (0..THREADS)
.map(|_| {
let store = Arc::clone(&store);
let ids = ids.clone();
std::thread::spawn(move || -> Result<(), ThreadFailure> {
for _ in 0..64 {
for id in &ids {
if store.lookup(id)?.is_none() {
return Err(ThreadFailure::Vanished);
}
}
}
Ok(())
})
})
.collect();
for thread in threads {
match thread.join() {
Ok(result) => result?,
Err(_) => return Err("a lookup thread panicked".into()),
}
}
assert_eq!(store.len()?, SESSIONS);
Ok(())
}
#[test]
fn the_counters_track_the_operations() -> Result<(), TestError> {
let store = SessionStore::new(Counting::shared());
assert!(store.is_empty()?);
assert_eq!(store.len()?, 0);
let first = store.open(Arc::new(()))?;
let second = store.open(Arc::new(()))?;
assert_eq!(store.len()?, 2);
assert!(!store.is_empty()?);
store.terminate(&first)?;
assert_eq!(store.len()?, 1);
store.terminate(&second)?;
assert!(store.is_empty()?);
Ok(())
}
struct Stepping {
millis: AtomicU64,
}
impl Stepping {
fn new() -> Arc<Self> {
Arc::new(Self {
millis: AtomicU64::new(0),
})
}
fn advance(&self, by: Duration) {
let added = u64::try_from(by.as_millis()).unwrap_or(u64::MAX);
self.millis.fetch_add(added, Ordering::Relaxed);
}
}
impl Clock for Stepping {
fn now(&self) -> Instant {
Instant::from_duration(Duration::from_millis(self.millis.load(Ordering::Relaxed)))
}
}
impl Clock for Arc<Stepping> {
fn now(&self) -> Instant {
<Stepping as Clock>::now(self.as_ref())
}
}
#[test]
fn a_store_without_a_reaper_is_a_no_op_for_reap_idle() -> Result<(), TestError> {
let store = SessionStore::new(Counting::shared());
let id = store.open(Arc::new(()))?;
assert!(store.reap_idle(Duration::from_millis(0))?.is_empty());
assert_eq!(store.len()?, 1, "no reaper means no eviction");
let _ = id;
Ok(())
}
#[test]
fn reap_idle_evicts_sessions_older_than_the_threshold() -> Result<(), TestError> {
let clock = Stepping::new();
let store = SessionStore::with_idle_reaper(Counting::shared(), clock.clone());
let stale_a = store.open(Arc::new(()))?;
let stale_b = store.open(Arc::new(()))?;
let live = store.open(Arc::new(()))?;
assert!(store.reap_idle(Duration::from_secs(60))?.is_empty());
assert_eq!(store.len()?, 3);
clock.advance(Duration::from_secs(120));
assert!(
store.lookup(&live)?.is_some(),
"the live session should still be in the store"
);
let reaped = store.reap_idle(Duration::from_secs(60))?;
let reaped_set: std::collections::HashSet<McpSessionId> = reaped.into_iter().collect();
assert!(reaped_set.contains(&stale_a));
assert!(reaped_set.contains(&stale_b));
assert!(!reaped_set.contains(&live));
assert_eq!(store.len()?, 1);
assert!(
store.lookup(&live)?.is_some(),
"the live session must survive the sweep"
);
assert!(store.lookup(&stale_a)?.is_none());
assert!(store.lookup(&stale_b)?.is_none());
Ok(())
}
#[test]
fn reap_idle_is_idempotent_under_a_steady_clock() -> Result<(), TestError> {
let clock = Stepping::new();
let store = SessionStore::with_idle_reaper(Counting::shared(), clock.clone());
let _id = store.open(Arc::new(()))?;
clock.advance(Duration::from_secs(120));
let first = store.reap_idle(Duration::from_secs(60))?;
assert_eq!(first.len(), 1);
let second = store.reap_idle(Duration::from_secs(60))?;
assert!(second.is_empty(), "the second sweep has nothing to do");
Ok(())
}
#[test]
fn terminate_clears_the_touched_entry() -> Result<(), TestError> {
let clock = Stepping::new();
let store = SessionStore::with_idle_reaper(Counting::shared(), clock.clone());
let id = store.open(Arc::new(()))?;
store.terminate(&id)?;
clock.advance(Duration::from_secs(120));
let reaped = store.reap_idle(Duration::from_secs(60))?;
assert!(reaped.is_empty());
assert_eq!(store.len()?, 0);
Ok(())
}
#[test]
fn a_looked_up_session_never_becomes_stale() -> Result<(), TestError> {
let clock = Stepping::new();
let store = SessionStore::with_idle_reaper(Counting::shared(), clock.clone());
let id = store.open(Arc::new(()))?;
for _ in 0..16 {
clock.advance(Duration::from_secs(60));
assert!(store.lookup(&id)?.is_some());
}
let reaped = store.reap_idle(Duration::from_secs(0))?;
assert!(
reaped.is_empty(),
"a session touched at `now` is not stale at `now`"
);
Ok(())
}
}