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