aster_forge_auth/
coordinator.rs1use 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#[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 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 pub fn close(&self) {
62 self.0.closed.cancel();
63 }
64 #[must_use]
66 pub fn is_closed(&self) -> bool {
67 self.0.closed.is_cancelled()
68 }
69 #[must_use]
71 pub fn active_requests(&self) -> usize {
72 self.0.maximum - self.0.budget.available_permits()
73 }
74 #[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 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 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}