tea_session/
memory.rs

1use std::collections::BTreeMap;
2use std::sync::RwLock;
3
4use tea_policy::{ActorId, PolicyGrant};
5use tea_protocol::SessionId;
6
7use crate::{
8    AppendOutcome, AppendTransaction, GrantJournalEntry, SessionStore, SessionStoreError,
9    SessionStoreErrorCode, SessionStoreFuture, StoredSession,
10};
11
12/// In-memory semantic reference implementation of [`SessionStore`].
13///
14/// All append validation and materialization is shared with durable stores via
15/// [`crate::apply_transaction`].
16#[derive(Debug, Default)]
17pub struct InMemorySessionStore {
18    sessions: RwLock<BTreeMap<SessionId, StoredSession>>,
19    names: RwLock<BTreeMap<SessionId, crate::SessionName>>,
20}
21
22impl InMemorySessionStore {
23    /// Creates an empty reference store.
24    #[must_use]
25    pub const fn new() -> Self {
26        Self {
27            sessions: RwLock::new(BTreeMap::new()),
28            names: RwLock::new(BTreeMap::new()),
29        }
30    }
31
32    fn append_sync(
33        &self,
34        transaction: &AppendTransaction,
35    ) -> Result<AppendOutcome, SessionStoreError> {
36        let mut sessions = self.sessions.write().map_err(|_| {
37            SessionStoreError::new(
38                SessionStoreErrorCode::StorageUnavailable,
39                "in-memory session store lock is poisoned",
40            )
41        })?;
42        let known_grant_ids: std::collections::HashSet<tea_policy::GrantId> = sessions
43            .values()
44            .flat_map(|stored| stored.grant_journal.iter())
45            .map(GrantJournalEntry::grant_id)
46            .collect();
47        let existed = sessions.contains_key(&transaction.session_id());
48        let stored = sessions.entry(transaction.session_id()).or_default();
49        let outcome = crate::apply_transaction_in_place(transaction, stored, existed, |grant_id| {
50            known_grant_ids.contains(&grant_id)
51        });
52        if outcome.is_err() && !existed {
53            sessions.remove(&transaction.session_id());
54        }
55        outcome
56    }
57}
58
59impl crate::SessionCatalog for InMemorySessionStore {
60    fn list_sessions(&self) -> SessionStoreFuture<'_, Vec<crate::SessionCatalogEntry>> {
61        Box::pin(async move {
62            let sessions = self.sessions.read().map_err(|_| {
63                SessionStoreError::new(
64                    SessionStoreErrorCode::StorageUnavailable,
65                    "in-memory session store lock is poisoned",
66                )
67            })?;
68            let names = self.names.read().map_err(|_| {
69                SessionStoreError::new(
70                    SessionStoreErrorCode::StorageUnavailable,
71                    "in-memory session name lock is poisoned",
72                )
73            })?;
74            let mut entries = sessions
75                .iter()
76                .map(|(session_id, stored)| {
77                    crate::catalog::catalog_entry(
78                        &stored.snapshot(),
79                        names.get(session_id).cloned(),
80                    )
81                })
82                .collect::<Result<Vec<_>, _>>()?;
83            crate::catalog::sort_catalog(&mut entries);
84            Ok(entries)
85        })
86    }
87
88    fn set_session_name(
89        &self,
90        session_id: SessionId,
91        name: Option<crate::SessionName>,
92    ) -> SessionStoreFuture<'_, ()> {
93        Box::pin(async move {
94            if !self
95                .sessions
96                .read()
97                .map_err(|_| {
98                    SessionStoreError::new(
99                        SessionStoreErrorCode::StorageUnavailable,
100                        "in-memory session store lock is poisoned",
101                    )
102                })?
103                .contains_key(&session_id)
104            {
105                return Err(SessionStoreError::new(
106                    SessionStoreErrorCode::SessionNotFound,
107                    "session does not exist",
108                ));
109            }
110            let mut names = self.names.write().map_err(|_| {
111                SessionStoreError::new(
112                    SessionStoreErrorCode::StorageUnavailable,
113                    "in-memory session name lock is poisoned",
114                )
115            })?;
116            if let Some(name) = name {
117                names.insert(session_id, name);
118            } else {
119                names.remove(&session_id);
120            }
121            Ok(())
122        })
123    }
124
125    fn session_name(
126        &self,
127        session_id: SessionId,
128    ) -> SessionStoreFuture<'_, Option<crate::SessionName>> {
129        Box::pin(async move {
130            if !self
131                .sessions
132                .read()
133                .map_err(|_| {
134                    SessionStoreError::new(
135                        SessionStoreErrorCode::StorageUnavailable,
136                        "in-memory session store lock is poisoned",
137                    )
138                })?
139                .contains_key(&session_id)
140            {
141                return Err(SessionStoreError::new(
142                    SessionStoreErrorCode::SessionNotFound,
143                    "session does not exist",
144                ));
145            }
146            Ok(self
147                .names
148                .read()
149                .map_err(|_| {
150                    SessionStoreError::new(
151                        SessionStoreErrorCode::StorageUnavailable,
152                        "in-memory session name lock is poisoned",
153                    )
154                })?
155                .get(&session_id)
156                .cloned())
157        })
158    }
159}
160
161impl SessionStore for InMemorySessionStore {
162    fn load(&self, session_id: SessionId) -> SessionStoreFuture<'_, crate::SessionSnapshot> {
163        Box::pin(async move {
164            self.sessions
165                .read()
166                .map_err(|_| {
167                    SessionStoreError::new(
168                        SessionStoreErrorCode::StorageUnavailable,
169                        "in-memory session store lock is poisoned",
170                    )
171                })?
172                .get(&session_id)
173                .map(StoredSession::snapshot)
174                .ok_or_else(|| {
175                    SessionStoreError::new(
176                        SessionStoreErrorCode::SessionNotFound,
177                        "session does not exist",
178                    )
179                })
180        })
181    }
182
183    fn append(&self, transaction: AppendTransaction) -> SessionStoreFuture<'_, AppendOutcome> {
184        Box::pin(async move { self.append_sync(&transaction) })
185    }
186
187    fn active_grants_for_actor(
188        &self,
189        actor_id: ActorId,
190    ) -> SessionStoreFuture<'_, Vec<PolicyGrant>> {
191        Box::pin(async move {
192            let sessions = self.sessions.read().map_err(|_| {
193                SessionStoreError::new(
194                    SessionStoreErrorCode::StorageUnavailable,
195                    "in-memory session store lock is poisoned",
196                )
197            })?;
198            Ok(crate::active_grants_for_actor(sessions.values(), &actor_id))
199        })
200    }
201}
tea_session/memory.rs

tea_session/
memory.rs
