rust_mcp_sdk/hyper_servers/session_store/
in_memory.rs

1use super::SessionId;
2use super::{SessionStore, TxServer};
3use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::io::DuplexStream;
7use tokio::sync::Mutex;
8use tokio::sync::RwLock;
9
10/// In-memory session store implementation
11///
12/// Stores session data in a thread-safe HashMap, using a read-write lock for
13/// concurrent access and mutexes for individual session streams.
14#[derive(Clone, Default)]
15pub struct InMemorySessionStore {
16    store: Arc<RwLock<HashMap<String, Arc<Mutex<TxServer>>>>>,
17}
18
19impl InMemorySessionStore {
20    /// Creates a new in-memory session store
21    ///
22    /// Initializes an empty HashMap wrapped in a read-write lock for thread-safe access.
23    ///
24    /// # Returns
25    /// * `Self` - A new InMemorySessionStore instance
26    pub fn new() -> Self {
27        Self {
28            store: Arc::new(RwLock::new(HashMap::new())),
29        }
30    }
31}
32
33/// Implementation of the SessionStore trait for InMemorySessionStore
34///
35/// Provides asynchronous methods for managing sessions in memory, ensuring
36/// thread-safety through read-write locks and mutexes.
37#[async_trait]
38impl SessionStore for InMemorySessionStore {
39    async fn get(&self, key: &SessionId) -> Option<Arc<Mutex<TxServer>>> {
40        let store = self.store.read().await;
41        store.get(key).cloned()
42    }
43
44    async fn set(&self, key: SessionId, value: TxServer) {
45        let mut store = self.store.write().await;
46        store.insert(key, Arc::new(Mutex::new(value)));
47    }
48
49    async fn delete(&self, key: &SessionId) {
50        let mut store = self.store.write().await;
51        store.remove(key);
52    }
53
54    async fn clear(&self) {
55        let mut store = self.store.write().await;
56        store.clear();
57    }
58    async fn keys(&self) -> Vec<SessionId> {
59        let store = self.store.read().await;
60        store.keys().cloned().collect::<Vec<_>>()
61    }
62    async fn values(&self) -> Vec<Arc<Mutex<DuplexStream>>> {
63        let store = self.store.read().await;
64        store.values().cloned().collect::<Vec<_>>()
65    }
66}