behest_runtime/
session_gate.rs

1//! Per-session concurrency gate.
2//!
3//! Prevents concurrent runs from interleaving writes to the same session.
4//! Each session is protected by its own [`tokio::sync::Mutex`]; a run
5//! acquires the lock at entry and holds it until the run completes or
6//! is cancelled.
7
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use tokio::sync::{Mutex, RwLock};
12use uuid::Uuid;
13
14/// RAII guard that holds a per-session lock for the duration of a run.
15///
16/// Returned by [`SessionGate::acquire`]. Dropping this guard releases the
17/// session lock, allowing the next queued run to proceed.
18pub struct SessionGuard {
19    _inner: tokio::sync::OwnedMutexGuard<()>,
20}
21
22impl SessionGuard {
23    /// Creates a new guard wrapping an acquired mutex lock.
24    fn new(inner: tokio::sync::OwnedMutexGuard<()>) -> Self {
25        Self { _inner: inner }
26    }
27}
28
29impl std::fmt::Debug for SessionGuard {
30    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        f.debug_struct("SessionGuard").finish_non_exhaustive()
32    }
33}
34
35/// Registry of per-session mutexes.
36///
37/// Uses a `RwLock<HashMap>` to lazily create mutex entries.  The outer
38/// `RwLock` is held briefly (read or write) to look up / insert a
39/// mutex *handle*; the actual session serialization happens via the
40/// inner `Mutex`.
41#[derive(Debug, Clone, Default)]
42pub struct SessionGate {
43    locks: Arc<RwLock<HashMap<Uuid, Arc<Mutex<()>>>>>,
44}
45
46impl SessionGate {
47    /// Creates a new empty gate.
48    #[must_use]
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    /// Acquires a session lock, returning an RAII guard.
54    ///
55    /// If the session is already held by another run, this returns
56    /// immediately with an error rather than waiting — concurrent
57    /// writes to the same session are a programming mistake, not a
58    /// normal contention scenario.
59    ///
60    /// # Errors
61    ///
62    /// Returns [`SessionBusy`] if the lock is already held.
63    pub async fn acquire(&self, session_id: Uuid) -> Result<SessionGuard, SessionBusy> {
64        let mutex = {
65            let read = self.locks.read().await;
66            read.get(&session_id).cloned()
67        };
68
69        let mutex = if let Some(m) = mutex {
70            m
71        } else {
72            let mut write = self.locks.write().await;
73            write
74                .entry(session_id)
75                .or_insert_with(|| Arc::new(Mutex::new(())))
76                .clone()
77        };
78
79        // Try to acquire the lock without blocking.  If another run
80        // already holds it, return an error immediately.
81        let guard = mutex
82            .try_lock_owned()
83            .map_err(|_| SessionBusy { session_id })?;
84
85        Ok(SessionGuard::new(guard))
86    }
87}
88
89/// Error returned when a session is already being processed by another run.
90#[derive(Debug, Clone, PartialEq, Eq)]
91pub struct SessionBusy {
92    /// The session that was already locked.
93    pub session_id: Uuid,
94}
95
96impl std::fmt::Display for SessionBusy {
97    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
98        write!(
99            f,
100            "session {} is busy — another run is active",
101            self.session_id
102        )
103    }
104}
105
106impl std::error::Error for SessionBusy {}
107
108#[cfg(test)]
109#[allow(clippy::unwrap_used, clippy::expect_used)]
110mod tests {
111    use super::*;
112
113    #[tokio::test]
114    async fn acquire_and_release() {
115        let gate = SessionGate::new();
116        let sid = Uuid::new_v4();
117
118        let guard = gate.acquire(sid).await.expect("should acquire");
119        // Guard is alive — session is locked.
120        drop(guard);
121
122        // After drop, the same session can be acquired again.
123        let _guard2 = gate
124            .acquire(sid)
125            .await
126            .expect("should re-acquire after drop");
127    }
128
129    #[tokio::test]
130    async fn concurrent_acquire_fails() {
131        let gate = SessionGate::new();
132        let sid = Uuid::new_v4();
133
134        let guard = gate.acquire(sid).await.expect("first acquire");
135        let result = gate.acquire(sid).await;
136        assert!(result.is_err(), "second acquire should fail");
137        assert_eq!(result.unwrap_err().session_id, sid);
138
139        drop(guard);
140
141        // After release, acquire succeeds.
142        let _guard2 = gate
143            .acquire(sid)
144            .await
145            .expect("should succeed after release");
146    }
147
148    #[tokio::test]
149    async fn independent_sessions_no_contention() {
150        let gate = SessionGate::new();
151        let sid1 = Uuid::new_v4();
152        let sid2 = Uuid::new_v4();
153
154        let guard1 = gate.acquire(sid1).await.expect("acquire sid1");
155        let guard2 = gate.acquire(sid2).await.expect("acquire sid2");
156
157        // Both acquired — no contention.
158        drop(guard1);
159        drop(guard2);
160    }
161
162    #[tokio::test]
163    async fn sequential_runs_serialize() {
164        let gate = SessionGate::new();
165        let sid = Uuid::new_v4();
166
167        let g1 = gate.acquire(sid).await.expect("first");
168        // Simulate first run finishing.
169        drop(g1);
170
171        let g2 = gate.acquire(sid).await.expect("second");
172        drop(g2);
173    }
174}
behest_runtime/session_gate.rs

behest_runtime/
session_gate.rs
