@@ -4,44 +4,115 @@ use super::SessionId;
44use super :: SessionStore ;
55use async_trait:: async_trait;
66use std:: collections:: HashMap ;
7+ use std:: sync:: atomic:: { AtomicU64 , Ordering } ;
78use std:: sync:: Arc ;
9+ use std:: time:: { Duration , SystemTime , UNIX_EPOCH } ;
810use tokio:: sync:: RwLock ;
911
10- /// In-memory session store implementation
12+ /// Default maximum number of concurrent sessions retained by the store.
13+ pub const DEFAULT_MAX_SESSIONS : usize = 10_000 ;
14+
15+ fn now_millis ( ) -> u64 {
16+ SystemTime :: now ( )
17+ . duration_since ( UNIX_EPOCH )
18+ . map ( |d| d. as_millis ( ) as u64 )
19+ . unwrap_or ( 0 )
20+ }
21+
22+ /// A stored session together with the time it was last accessed.
23+ struct SessionEntry {
24+ runtime : Arc < ServerRuntime > ,
25+ last_access_ms : AtomicU64 ,
26+ }
27+
28+ impl SessionEntry {
29+ fn new ( runtime : Arc < ServerRuntime > ) -> Self {
30+ Self {
31+ runtime,
32+ last_access_ms : AtomicU64 :: new ( now_millis ( ) ) ,
33+ }
34+ }
35+
36+ /// Marks the session as accessed now.
37+ fn touch ( & self ) {
38+ self . last_access_ms . store ( now_millis ( ) , Ordering :: Relaxed ) ;
39+ }
40+
41+ /// Returns true if the session has been idle for longer than `ttl_ms`.
42+ fn is_idle ( & self , now_ms : u64 , ttl_ms : u64 ) -> bool {
43+ now_ms. saturating_sub ( self . last_access_ms . load ( Ordering :: Relaxed ) ) > ttl_ms
44+ }
45+ }
46+
47+ /// In-memory session store with a bounded session count and optional idle TTL.
1148///
12- /// Stores session data in a thread-safe HashMap, using a read-write lock for
13- #[ derive( Clone , Default ) ]
49+ /// Idle sessions (older than the configured TTL) are evicted lazily, on access
50+ /// and whenever the store is checked for capacity. Once `max_sessions` is
51+ /// reached the server rejects new sessions with `503 Service Unavailable`,
52+ /// preventing an unauthenticated client from exhausting memory via repeated
53+ /// `initialize` requests.
54+ #[ derive( Clone ) ]
1455pub struct InMemorySessionStore {
15- store : Arc < RwLock < HashMap < String , Arc < ServerRuntime > > > > ,
56+ store : Arc < RwLock < HashMap < String , SessionEntry > > > ,
57+ max_sessions : usize ,
58+ idle_ttl : Option < Duration > ,
59+ }
60+
61+ impl Default for InMemorySessionStore {
62+ fn default ( ) -> Self {
63+ Self :: with_limits ( None , None )
64+ }
1665}
1766
1867impl InMemorySessionStore {
19- /// Creates a new in-memory session store
20- ///
21- /// Initializes an empty HashMap wrapped in a read-write lock for thread-safe access.
22- ///
23- /// # Returns
24- /// * `Self` - A new InMemorySessionStore instance
68+ /// Creates a new in-memory session store with default limits
69+ /// ([`DEFAULT_MAX_SESSIONS`], no idle TTL).
2570pub fn new ( ) -> Self {
71+ Self :: default ( )
72+ }
73+
74+ /// Creates a session store with explicit limits.
75+ ///
76+ /// * `max_sessions` - maximum number of concurrent sessions; `None` uses
77+ /// [`DEFAULT_MAX_SESSIONS`]. Pass `Some(usize::MAX)` for an effectively
78+ /// unbounded store.
79+ /// * `idle_ttl` - sessions idle for longer than this are evicted; `None`
80+ /// disables idle expiry.
81+ pub fn with_limits ( max_sessions : Option < usize > , idle_ttl : Option < Duration > ) -> Self {
2682 Self {
2783 store : Arc :: new ( RwLock :: new ( HashMap :: new ( ) ) ) ,
84+ max_sessions : max_sessions. unwrap_or ( DEFAULT_MAX_SESSIONS ) ,
85+ idle_ttl,
2886 }
2987 }
88+
89+ /// Evicts sessions idle past the configured TTL and returns the resulting
90+ /// session count.
91+ async fn evict_idle ( & self ) -> usize {
92+ let Some ( ttl) = self . idle_ttl else {
93+ return self . store . read ( ) . await . len ( ) ;
94+ } ;
95+ let ttl_ms = ttl. as_millis ( ) as u64 ;
96+ let now = now_millis ( ) ;
97+ let mut store = self . store . write ( ) . await ;
98+ store. retain ( |_, entry| !entry. is_idle ( now, ttl_ms) ) ;
99+ store. len ( )
100+ }
30101}
31102
32103/// Implementation of the SessionStore trait for InMemorySessionStore
33- ///
34- /// Provides asynchronous methods for managing sessions in memory, ensuring
35104#[ async_trait]
36105impl SessionStore for InMemorySessionStore {
37106 async fn get ( & self , key : & SessionId ) -> Option < Arc < ServerRuntime > > {
38107 let store = self . store . read ( ) . await ;
39- store. get ( key) . cloned ( )
108+ let entry = store. get ( key) ?;
109+ entry. touch ( ) ;
110+ Some ( entry. runtime . clone ( ) )
40111 }
41112
42113 async fn set ( & self , key : SessionId , value : Arc < ServerRuntime > ) {
43114 let mut store = self . store . write ( ) . await ;
44- store. insert ( key, value) ;
115+ store. insert ( key, SessionEntry :: new ( value) ) ;
45116 }
46117
47118 async fn delete ( & self , key : & SessionId ) {
@@ -59,10 +130,18 @@ impl SessionStore for InMemorySessionStore {
59130 }
60131 async fn values ( & self ) -> Vec < Arc < ServerRuntime > > {
61132 let store = self . store . read ( ) . await ;
62- store. values ( ) . cloned ( ) . collect :: < Vec < _ > > ( )
133+ store
134+ . values ( )
135+ . map ( |entry| entry. runtime . clone ( ) )
136+ . collect :: < Vec < _ > > ( )
63137 }
64138 async fn has ( & self , session : & SessionId ) -> bool {
65139 let store = self . store . read ( ) . await ;
66140 store. contains_key ( session)
67141 }
142+
143+ async fn is_full ( & self ) -> bool {
144+ let count = self . evict_idle ( ) . await ;
145+ count >= self . max_sessions
146+ }
68147}
0 commit comments