alopex_server/
session.rs

1use std::sync::Arc;
2use std::time::{Duration, SystemTime};
3
4use alopex_cluster::TableLifecycleEffect;
5use alopex_core::async_runtime::{BoxFuture, BoxStream};
6use alopex_sql::catalog::TableMetadata;
7use alopex_sql::executor::{ExecutionResult, ExecutorError, Row};
8use alopex_sql::planner::PlannedStatement;
9use alopex_sql::storage::erased::ErasedAsyncSqlTransaction;
10use dashmap::DashMap;
11use futures::StreamExt;
12use tokio::sync::mpsc;
13use tokio_stream::wrappers::ReceiverStream;
14use uuid::Uuid;
15
16use crate::error::{Result, ServerError};
17
18/// Session identifier.
19#[derive(Clone, Debug, Eq, Hash, PartialEq, serde::Serialize, serde::Deserialize)]
20pub struct SessionId(Uuid);
21
22impl SessionId {
23    pub fn new() -> Self {
24        Self(Uuid::new_v4())
25    }
26}
27
28impl Default for SessionId {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34impl std::fmt::Display for SessionId {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        write!(f, "{}", self.0)
37    }
38}
39
40impl std::str::FromStr for SessionId {
41    type Err = uuid::Error;
42
43    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
44        Ok(Self(Uuid::parse_str(s)?))
45    }
46}
47
48/// Session lifecycle state.
49#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
50pub enum SessionState {
51    Idle,
52    InTransaction,
53    Committing,
54    RollingBack,
55}
56
57/// Snapshot of a session for safe sharing.
58#[derive(Clone, Debug, serde::Serialize)]
59pub struct SessionSnapshot {
60    pub id: SessionId,
61    pub has_transaction: bool,
62    pub created_at: SystemTime,
63    pub last_active: SystemTime,
64    pub expires_at: SystemTime,
65    pub state: SessionState,
66}
67
68/// Transaction handle for a session.
69#[derive(Clone)]
70pub struct TxnHandle {
71    inner: Arc<TxnHandleInner>,
72}
73
74#[derive(Clone)]
75pub enum CatalogRollbackEffect {
76    DropTable { table_name: String },
77    CreateTable { table: Box<TableMetadata> },
78}
79
80struct TxnHandleInner {
81    txn: tokio::sync::Mutex<Option<Box<dyn ErasedAsyncSqlTransaction>>>,
82    pending_table_lifecycle_effects: tokio::sync::Mutex<Vec<TableLifecycleEffect>>,
83    pending_catalog_rollback_effects: tokio::sync::Mutex<Vec<CatalogRollbackEffect>>,
84    created_at: SystemTime,
85}
86
87impl TxnHandle {
88    pub fn new(txn: Box<dyn ErasedAsyncSqlTransaction>) -> Self {
89        Self {
90            inner: Arc::new(TxnHandleInner {
91                txn: tokio::sync::Mutex::new(Some(txn)),
92                pending_table_lifecycle_effects: tokio::sync::Mutex::new(Vec::new()),
93                pending_catalog_rollback_effects: tokio::sync::Mutex::new(Vec::new()),
94                created_at: SystemTime::now(),
95            }),
96        }
97    }
98
99    pub fn created_at(&self) -> SystemTime {
100        self.inner.created_at
101    }
102
103    pub fn execute<'a>(
104        &'a self,
105        sql: &'a str,
106    ) -> BoxFuture<'a, alopex_sql::executor::Result<ExecutionResult>> {
107        Box::pin(async move {
108            let mut guard = self.inner.txn.lock().await;
109            let txn = guard
110                .as_mut()
111                .ok_or_else(|| ExecutorError::InvalidOperation {
112                    operation: "execute".into(),
113                    reason: "transaction is closed".into(),
114                })?;
115            txn.execute(sql).await
116        })
117    }
118
119    pub fn query<'a>(&'a self, sql: &'a str) -> BoxStream<'a, alopex_sql::executor::Result<Row>> {
120        let (sender, receiver) = mpsc::channel(32);
121        let sql = sql.to_string();
122        let inner = Arc::clone(&self.inner);
123
124        tokio::spawn(async move {
125            let guard = inner.txn.lock().await;
126            let Some(txn) = guard.as_ref() else {
127                let _ = sender
128                    .send(Err(ExecutorError::InvalidOperation {
129                        operation: "query".into(),
130                        reason: "transaction is closed".into(),
131                    }))
132                    .await;
133                return;
134            };
135            let mut stream = txn.query(&sql);
136            while let Some(item) = stream.next().await {
137                if sender.send(item).await.is_err() {
138                    break;
139                }
140            }
141        });
142
143        Box::pin(ReceiverStream::new(receiver))
144    }
145
146    pub fn plan_for_routing<'a>(
147        &'a self,
148        sql: &'a str,
149    ) -> BoxFuture<'a, alopex_sql::executor::Result<Vec<PlannedStatement>>> {
150        Box::pin(async move {
151            let guard = self.inner.txn.lock().await;
152            let txn = guard
153                .as_ref()
154                .ok_or_else(|| ExecutorError::InvalidOperation {
155                    operation: "plan_for_routing".into(),
156                    reason: "transaction is closed".into(),
157                })?;
158            txn.plan_for_routing(sql).await
159        })
160    }
161
162    pub async fn buffer_table_lifecycle_effects(&self, effects: Vec<TableLifecycleEffect>) {
163        if effects.is_empty() {
164            return;
165        }
166        self.inner
167            .pending_table_lifecycle_effects
168            .lock()
169            .await
170            .extend(effects);
171    }
172
173    pub async fn buffer_catalog_rollback_effects(&self, effects: Vec<CatalogRollbackEffect>) {
174        if effects.is_empty() {
175            return;
176        }
177        self.inner
178            .pending_catalog_rollback_effects
179            .lock()
180            .await
181            .extend(effects);
182    }
183
184    pub async fn commit(self) -> alopex_sql::executor::Result<Vec<TableLifecycleEffect>> {
185        let mut guard = self.inner.txn.lock().await;
186        let txn = guard
187            .take()
188            .ok_or_else(|| ExecutorError::InvalidOperation {
189                operation: "commit".into(),
190                reason: "transaction is closed".into(),
191            })?;
192        txn.commit_boxed().await?;
193        let mut effects = self.inner.pending_table_lifecycle_effects.lock().await;
194        Ok(std::mem::take(&mut *effects))
195    }
196
197    pub async fn rollback(self) -> alopex_sql::executor::Result<Vec<CatalogRollbackEffect>> {
198        let mut guard = self.inner.txn.lock().await;
199        let txn = guard
200            .take()
201            .ok_or_else(|| ExecutorError::InvalidOperation {
202                operation: "rollback".into(),
203                reason: "transaction is closed".into(),
204            })?;
205        let result = txn.rollback_boxed().await;
206        self.inner
207            .pending_table_lifecycle_effects
208            .lock()
209            .await
210            .clear();
211        result?;
212        let mut effects = self.inner.pending_catalog_rollback_effects.lock().await;
213        Ok(std::mem::take(&mut *effects))
214    }
215}
216
217/// Session configuration.
218#[derive(Clone, Copy, Debug)]
219pub struct SessionConfig {
220    pub ttl: Duration,
221}
222
223/// Transaction factory for session manager.
224pub type TransactionFactory =
225    Arc<dyn Fn() -> BoxFuture<'static, Result<Box<dyn ErasedAsyncSqlTransaction>>> + Send + Sync>;
226
227/// Session manager for server.
228pub struct SessionManager {
229    sessions: DashMap<SessionId, Session>,
230    config: SessionConfig,
231    txn_factory: TransactionFactory,
232}
233
234struct Session {
235    id: SessionId,
236    txn_handle: Option<TxnHandle>,
237    created_at: SystemTime,
238    last_active: SystemTime,
239    expires_at: SystemTime,
240    state: SessionState,
241}
242
243impl SessionManager {
244    pub fn new(config: SessionConfig, txn_factory: TransactionFactory) -> Self {
245        Self {
246            sessions: DashMap::new(),
247            config,
248            txn_factory,
249        }
250    }
251
252    pub async fn create_session(&self) -> Result<SessionId> {
253        let now = SystemTime::now();
254        let id = SessionId::new();
255        let session = Session {
256            id: id.clone(),
257            txn_handle: None,
258            created_at: now,
259            last_active: now,
260            expires_at: now + self.config.ttl,
261            state: SessionState::Idle,
262        };
263        self.sessions.insert(id.clone(), session);
264        Ok(id)
265    }
266
267    pub async fn get_session(&self, id: &SessionId) -> Result<SessionSnapshot> {
268        let entry = self
269            .sessions
270            .get(id)
271            .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
272        if entry.expires_at <= SystemTime::now() {
273            drop(entry);
274            self.sessions.remove(id);
275            return Err(ServerError::SessionExpired("session expired".into()));
276        }
277        Ok(SessionSnapshot {
278            id: entry.id.clone(),
279            has_transaction: entry.txn_handle.is_some(),
280            created_at: entry.created_at,
281            last_active: entry.last_active,
282            expires_at: entry.expires_at,
283            state: entry.state,
284        })
285    }
286
287    pub async fn begin_transaction(&self, id: &SessionId) -> Result<TxnHandle> {
288        let mut entry = self
289            .sessions
290            .get_mut(id)
291            .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
292        if entry.expires_at <= SystemTime::now() {
293            drop(entry);
294            self.sessions.remove(id);
295            return Err(ServerError::SessionExpired("session expired".into()));
296        }
297        if entry.txn_handle.is_some() {
298            return Err(ServerError::Conflict("transaction already active".into()));
299        }
300        let txn = (self.txn_factory)().await?;
301        let handle = TxnHandle::new(txn);
302        entry.txn_handle = Some(handle.clone());
303        entry.last_active = SystemTime::now();
304        entry.state = SessionState::InTransaction;
305        Ok(handle)
306    }
307
308    pub async fn get_transaction(&self, id: &SessionId) -> Result<TxnHandle> {
309        let mut entry = self
310            .sessions
311            .get_mut(id)
312            .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
313        if entry.expires_at <= SystemTime::now() {
314            drop(entry);
315            self.sessions.remove(id);
316            return Err(ServerError::SessionExpired("session expired".into()));
317        }
318        let handle = entry
319            .txn_handle
320            .clone()
321            .ok_or_else(|| ServerError::BadRequest("transaction not started".into()))?;
322        entry.last_active = SystemTime::now();
323        entry.state = SessionState::InTransaction;
324        Ok(handle)
325    }
326
327    pub async fn execute_in_session(&self, id: &SessionId, sql: &str) -> Result<ExecutionResult> {
328        let handle = {
329            let mut entry = self
330                .sessions
331                .get_mut(id)
332                .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
333            if entry.expires_at <= SystemTime::now() {
334                drop(entry);
335                self.sessions.remove(id);
336                return Err(ServerError::SessionExpired("session expired".into()));
337            }
338            let handle = entry
339                .txn_handle
340                .clone()
341                .ok_or_else(|| ServerError::BadRequest("transaction not started".into()))?;
342            entry.last_active = SystemTime::now();
343            handle
344        };
345
346        handle
347            .execute(sql)
348            .await
349            .map_err(|err| ServerError::Sql(err.into()))
350    }
351
352    pub async fn commit(&self, id: &SessionId) -> Result<Vec<TableLifecycleEffect>> {
353        let handle = self.take_handle(id, SessionState::Committing)?;
354        let effects = handle
355            .commit()
356            .await
357            .map_err(|err| ServerError::Sql(err.into()))?;
358        Ok(effects)
359    }
360
361    pub async fn rollback(&self, id: &SessionId) -> Result<Vec<CatalogRollbackEffect>> {
362        let handle = self.take_handle(id, SessionState::RollingBack)?;
363        let effects = handle
364            .rollback()
365            .await
366            .map_err(|err| ServerError::Sql(err.into()))?;
367        Ok(effects)
368    }
369
370    pub fn cleanup_expired(&self) {
371        let now = SystemTime::now();
372        let expired: Vec<SessionId> = self
373            .sessions
374            .iter()
375            .filter(|entry| entry.expires_at <= now)
376            .map(|entry| entry.id.clone())
377            .collect();
378        for id in expired {
379            self.sessions.remove(&id);
380        }
381    }
382
383    fn take_handle(&self, id: &SessionId, state: SessionState) -> Result<TxnHandle> {
384        let mut entry = self
385            .sessions
386            .get_mut(id)
387            .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
388        if entry.expires_at <= SystemTime::now() {
389            drop(entry);
390            self.sessions.remove(id);
391            return Err(ServerError::SessionExpired("session expired".into()));
392        }
393        let handle = entry
394            .txn_handle
395            .take()
396            .ok_or_else(|| ServerError::BadRequest("transaction not started".into()))?;
397        entry.state = state;
398        entry.last_active = SystemTime::now();
399        Ok(handle)
400    }
401}
alopex_server/session.rs

alopex_server/
session.rs
