Skip to main content

aster_forge_auth/
coordinator.rs

1//! Bounded, instance-owned refresh serialization with cancellation-safe resource reclamation.
2
3use crate::{AuthError, RefreshJti, Result};
4use std::{
5    collections::HashMap,
6    fmt,
7    sync::{Arc, Mutex as MapMutex, Weak},
8};
9use tokio::sync::{Mutex, OwnedMutexGuard, OwnedSemaphorePermit, Semaphore};
10use tokio_util::sync::CancellationToken;
11
12struct Inner {
13    locks: MapMutex<HashMap<RefreshJti, Weak<Mutex<()>>>>,
14    budget: Arc<Semaphore>,
15    maximum: usize,
16    closed: CancellationToken,
17}
18
19/// Instance resource shared by all refresh requests. No global singleton or historical-JTI leak.
20/// The finite budget includes running requests and same-JTI waiters; saturation rejects promptly.
21/// It complements, never replaces, the product's authoritative cross-instance database CAS.
22#[derive(Clone)]
23pub struct RefreshCoordinator(Arc<Inner>);
24
25impl fmt::Debug for RefreshCoordinator {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        f.debug_struct("RefreshCoordinator")
28            .field("maximum", &self.0.maximum)
29            .field("active_requests", &self.active_requests())
30            .field("closed", &self.is_closed())
31            .finish_non_exhaustive()
32    }
33}
34
35impl Default for RefreshCoordinator {
36    fn default() -> Self {
37        Self::from_limit(1024)
38    }
39}
40
41impl RefreshCoordinator {
42    /// Creates a bounded instance resource.
43    ///
44    /// # Errors
45    /// Returns `InvalidPolicy` for zero or a limit beyond Tokio's semaphore capacity.
46    pub fn new(maximum: usize) -> Result<Self> {
47        if maximum == 0 || maximum > Semaphore::MAX_PERMITS {
48            return Err(AuthError::InvalidPolicy);
49        }
50        Ok(Self::from_limit(maximum))
51    }
52    fn from_limit(maximum: usize) -> Self {
53        Self(Arc::new(Inner {
54            locks: MapMutex::new(HashMap::new()),
55            budget: Arc::new(Semaphore::new(maximum)),
56            maximum,
57            closed: CancellationToken::new(),
58        }))
59    }
60    /// Rejects new requests and wakes queued waiters. Already-started store transactions complete.
61    pub fn close(&self) {
62        self.0.closed.cancel();
63    }
64    /// Reports whether shutdown has started.
65    #[must_use]
66    pub fn is_closed(&self) -> bool {
67        self.0.closed.is_cancelled()
68    }
69    /// Number of admitted running/queued requests, for resource diagnostics and deterministic tests.
70    #[must_use]
71    pub fn active_requests(&self) -> usize {
72        self.0.maximum - self.0.budget.available_permits()
73    }
74    /// Number of live JTI lock entries. No raw identity is exposed to observation.
75    #[must_use]
76    pub fn tracked_ids(&self) -> usize {
77        self.0
78            .locks
79            .lock()
80            .unwrap_or_else(std::sync::PoisonError::into_inner)
81            .len()
82    }
83
84    pub(crate) async fn acquire(
85        &self,
86        jti: RefreshJti,
87        cancel: &CancellationToken,
88    ) -> Result<Reservation> {
89        if self.is_closed() || cancel.is_cancelled() {
90            return Err(AuthError::Cancelled);
91        }
92        let permit = self
93            .0
94            .budget
95            .clone()
96            .try_acquire_owned()
97            .map_err(|_| AuthError::ResourceLimit)?;
98        let lock = {
99            let mut map = self
100                .0
101                .locks
102                .lock()
103                .unwrap_or_else(std::sync::PoisonError::into_inner);
104            if let Some(lock) = map.get(&jti).and_then(Weak::upgrade) {
105                lock
106            } else {
107                let lock = Arc::new(Mutex::new(()));
108                map.insert(jti.clone(), Arc::downgrade(&lock));
109                lock
110            }
111        };
112        // Reserve before awaiting so cancellation/drop during lock acquisition also reclaims entries.
113        let mut reservation = Reservation {
114            jti,
115            owner: self.clone(),
116            lock: Some(lock.clone()),
117            guard: None,
118            _permit: permit,
119        };
120        let guard = tokio::select! {
121            biased;
122            () = cancel.cancelled() => return Err(AuthError::Cancelled),
123            () = self.0.closed.cancelled() => return Err(AuthError::Cancelled),
124            guard = lock.lock_owned() => guard,
125        };
126        reservation.guard = Some(guard);
127        Ok(reservation)
128    }
129}
130
131pub(crate) struct Reservation {
132    jti: RefreshJti,
133    owner: RefreshCoordinator,
134    lock: Option<Arc<Mutex<()>>>,
135    guard: Option<OwnedMutexGuard<()>>,
136    _permit: OwnedSemaphorePermit,
137}
138impl Drop for Reservation {
139    fn drop(&mut self) {
140        drop(self.guard.take());
141        drop(self.lock.take());
142        // Registration and removal share one mutex; removal cannot erase a newly-created live lock.
143        let mut map = self
144            .owner
145            .0
146            .locks
147            .lock()
148            .unwrap_or_else(std::sync::PoisonError::into_inner);
149        if map
150            .get(&self.jti)
151            .is_some_and(|weak| weak.strong_count() == 0)
152        {
153            map.remove(&self.jti);
154        }
155    }
156}