Skip to main content

reinhardt_middleware/session/
store.rs

1//! In-memory `SessionStore` with lazy eviction of expired sessions.
2
3use std::collections::HashMap;
4use std::sync::RwLock;
5use std::sync::atomic::{AtomicUsize, Ordering};
6
7use super::data::SessionData;
8
9/// DI key for resolving the middleware-owned session store through
10/// `Depends<SessionStoreKey, Arc<SessionStore>>`.
11#[derive(Debug, Clone, Copy)]
12pub struct SessionStoreKey;
13
14impl reinhardt_di::InjectableKey for SessionStoreKey {}
15
16/// Session store with automatic lazy eviction of expired sessions
17///
18/// Performs threshold-based lazy cleanup of expired sessions to prevent
19/// unbounded memory growth. Cleanup is triggered inside `save` when the
20/// total number of stored sessions crosses an amortized cleanup boundary; it
21/// is not time-based. The threshold defaults to
22/// `SessionStore::DEFAULT_CLEANUP_THRESHOLD` and can be overridden at
23/// construction via `SessionStore::with_cleanup_threshold` or adjusted at
24/// runtime via `SessionStore::set_cleanup_threshold`.
25#[derive(Debug)]
26pub struct SessionStore {
27	/// Sessions
28	pub(super) sessions: RwLock<HashMap<String, SessionData>>,
29	/// Maximum number of sessions before triggering automatic cleanup
30	max_sessions_before_cleanup: AtomicUsize,
31	/// Next session count at which `save` should perform cleanup.
32	next_cleanup_session_count: AtomicUsize,
33}
34
35impl Default for SessionStore {
36	fn default() -> Self {
37		Self::new()
38	}
39}
40
41impl SessionStore {
42	/// Default cleanup threshold: trigger cleanup when session count exceeds 10,000
43	pub const DEFAULT_CLEANUP_THRESHOLD: usize = 10_000;
44
45	/// Create a new store with the default cleanup threshold
46	/// (`SessionStore::DEFAULT_CLEANUP_THRESHOLD`).
47	pub fn new() -> Self {
48		Self::with_cleanup_threshold(Self::DEFAULT_CLEANUP_THRESHOLD)
49	}
50
51	/// Create a new store with a custom cleanup threshold.
52	///
53	/// The store triggers a single `retain` pass inside `save` whenever the
54	/// number of stored sessions crosses an amortized cleanup boundary.
55	pub fn with_cleanup_threshold(threshold: usize) -> Self {
56		Self {
57			sessions: RwLock::new(HashMap::new()),
58			max_sessions_before_cleanup: AtomicUsize::new(threshold),
59			next_cleanup_session_count: AtomicUsize::new(threshold),
60		}
61	}
62
63	/// Update the cleanup threshold at runtime.
64	pub fn set_cleanup_threshold(&self, threshold: usize) {
65		self.max_sessions_before_cleanup
66			.store(threshold, Ordering::Relaxed);
67		self.next_cleanup_session_count
68			.store(threshold, Ordering::Relaxed);
69	}
70
71	/// Get a session
72	pub fn get(&self, id: &str) -> Option<SessionData> {
73		let sessions = self.sessions.read().unwrap_or_else(|e| e.into_inner());
74		sessions.get(id).cloned()
75	}
76
77	/// Save a session, with automatic cleanup when threshold is exceeded.
78	///
79	/// Cleanup runs when the post-save session count exceeds the configured
80	/// threshold and reaches the next amortized cleanup boundary. The boundary
81	/// advances after each pass so a store that remains above the threshold does
82	/// not scan all sessions on every save.
83	pub fn save(&self, session: SessionData) {
84		let mut sessions = self.sessions.write().unwrap_or_else(|e| e.into_inner());
85		sessions.insert(session.id.clone(), session);
86
87		let threshold = self.max_sessions_before_cleanup.load(Ordering::Relaxed);
88		let next_cleanup_session_count = self.next_cleanup_session_count.load(Ordering::Relaxed);
89		if sessions.len() > threshold && sessions.len() >= next_cleanup_session_count {
90			sessions.retain(|_, s| s.is_valid());
91
92			let cleanup_interval = threshold.max(1);
93			self.next_cleanup_session_count.store(
94				sessions.len().saturating_add(cleanup_interval),
95				Ordering::Relaxed,
96			);
97		}
98	}
99
100	/// Delete a session
101	pub fn delete(&self, id: &str) {
102		let mut sessions = self.sessions.write().unwrap_or_else(|e| e.into_inner());
103		sessions.remove(id);
104	}
105
106	/// Clean up expired sessions
107	pub fn cleanup(&self) {
108		let mut sessions = self.sessions.write().unwrap_or_else(|e| e.into_inner());
109		sessions.retain(|_, session| session.is_valid());
110	}
111
112	/// Clear the store
113	pub fn clear(&self) {
114		let mut sessions = self.sessions.write().unwrap_or_else(|e| e.into_inner());
115		sessions.clear();
116	}
117
118	/// Get the number of sessions
119	pub fn len(&self) -> usize {
120		let sessions = self.sessions.read().unwrap_or_else(|e| e.into_inner());
121		sessions.len()
122	}
123
124	/// Check if the store is empty
125	pub fn is_empty(&self) -> bool {
126		let sessions = self.sessions.read().unwrap_or_else(|e| e.into_inner());
127		sessions.is_empty()
128	}
129}