alice.pds.demo.boltless.dev
This repository has no description
11 kB
388 lines
did:plc:3fwecdnvtcscjnrx2p4n7a…
1use std::collections::HashMap;
2use std::collections::hash_map::Entry as Slot;
3use std::hash::Hash;
4use std::sync::Mutex;
5
6use knot_runtime::UnixMicros;
7
8knot_types::scalar_newtype! {
9 pub struct GroupQuota(usize);
10 pub struct TotalQuota(usize);
11}
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub struct Quotas {
15 pub per_group: GroupQuota,
16 pub total: TotalQuota,
17}
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum Rejected {
21 Total,
22 Group,
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub enum Admitted<V> {
27 Inserted,
28 Occupied(V),
29}
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32enum Occupancy {
33 Keep,
34 Extend,
35}
36
37struct Entry<G, V> {
38 group: G,
39 value: V,
40 expires_at: UnixMicros,
41}
42
43struct Inner<K, G, V> {
44 entries: HashMap<K, Entry<G, V>>,
45 group_counts: HashMap<G, usize>,
46}
47
48pub struct Expiring<K, G, V> {
49 inner: Mutex<Inner<K, G, V>>,
50 quotas: Quotas,
51}
52
53impl<K, G, V> Expiring<K, G, V>
54where
55 K: Eq + Hash + Clone,
56 G: Eq + Hash + Clone,
57{
58 pub fn new(quotas: Quotas) -> Self {
59 Self {
60 inner: Mutex::new(Inner {
61 entries: HashMap::new(),
62 group_counts: HashMap::new(),
63 }),
64 quotas,
65 }
66 }
67
68 fn lock(&self) -> std::sync::MutexGuard<'_, Inner<K, G, V>> {
69 self.inner
70 .lock()
71 .unwrap_or_else(|poisoned| poisoned.into_inner())
72 }
73
74 pub fn prune(&self, now: UnixMicros) -> Vec<K> {
75 prune_locked(&mut self.lock(), now)
76 }
77
78 pub fn admit(
79 &self,
80 key: K,
81 group: G,
82 value: V,
83 expires_at: UnixMicros,
84 now: UnixMicros,
85 ) -> Result<Admitted<V>, Rejected>
86 where
87 V: Clone + PartialEq,
88 {
89 self.enter(key, group, value, expires_at, now, Occupancy::Keep)
90 }
91
92 pub fn admit_or_renew(
93 &self,
94 key: K,
95 group: G,
96 value: V,
97 expires_at: UnixMicros,
98 now: UnixMicros,
99 ) -> Result<Admitted<V>, Rejected>
100 where
101 V: Clone + PartialEq,
102 {
103 self.enter(key, group, value, expires_at, now, Occupancy::Extend)
104 }
105
106 fn enter(
107 &self,
108 key: K,
109 group: G,
110 value: V,
111 expires_at: UnixMicros,
112 now: UnixMicros,
113 occupied: Occupancy,
114 ) -> Result<Admitted<V>, Rejected>
115 where
116 V: Clone + PartialEq,
117 {
118 let mut inner = self.lock();
119 let key = match inner.entries.entry(key) {
120 Slot::Occupied(mut held) if held.get().expires_at > now => {
121 let entry = held.get_mut();
122 if occupied == Occupancy::Extend && entry.value == value {
123 entry.expires_at = expires_at;
124 }
125 return Ok(Admitted::Occupied(entry.value.clone()));
126 }
127 Slot::Occupied(held) => {
128 let (key, entry) = held.remove_entry();
129 release_group(&mut inner.group_counts, &entry.group);
130 key
131 }
132 Slot::Vacant(free) => free.into_key(),
133 };
134 // This is lazy on purpose!
135 // Pruning traverses every entry so we shouldn't do it on every
136 // admission, better to do this O(1) quota check.
137 if self.rejection(&inner, &group).is_some() {
138 prune_locked(&mut inner, now);
139 }
140 match self.rejection(&inner, &group) {
141 Some(rejected) => Err(rejected),
142 None => {
143 *inner.group_counts.entry(group.clone()).or_insert(0) += 1;
144 inner.entries.insert(
145 key,
146 Entry {
147 group,
148 value,
149 expires_at,
150 },
151 );
152 Ok(Admitted::Inserted)
153 }
154 }
155 }
156
157 fn rejection(&self, inner: &Inner<K, G, V>, group: &G) -> Option<Rejected> {
158 let held = inner.group_counts.get(group).copied().unwrap_or(0);
159 match (
160 held >= self.quotas.per_group.get(),
161 inner.entries.len() >= self.quotas.total.get(),
162 ) {
163 (true, _) => Some(Rejected::Group),
164 (_, true) => Some(Rejected::Total),
165 (false, false) => None,
166 }
167 }
168
169 pub fn remove(&self, key: &K) -> Option<V> {
170 remove_locked(&mut self.lock(), key)
171 }
172
173 pub fn get(&self, key: &K, now: UnixMicros) -> Option<V>
174 where
175 V: Clone,
176 {
177 let inner = self.lock();
178 inner
179 .entries
180 .get(key)
181 .filter(|entry| entry.expires_at > now)
182 .map(|entry| entry.value.clone())
183 }
184
185 pub fn len(&self) -> usize {
186 self.lock().entries.len()
187 }
188
189 pub fn is_empty(&self) -> bool {
190 self.len() == 0
191 }
192
193 pub fn group_len(&self, group: &G) -> usize {
194 self.lock().group_counts.get(group).copied().unwrap_or(0)
195 }
196}
197
198fn prune_locked<K, G, V>(inner: &mut Inner<K, G, V>, now: UnixMicros) -> Vec<K>
199where
200 K: Eq + Hash + Clone,
201 G: Eq + Hash,
202{
203 let expired: Vec<K> = inner
204 .entries
205 .iter()
206 .filter(|(_, entry)| entry.expires_at <= now)
207 .map(|(key, _)| key.clone())
208 .collect();
209 expired.iter().for_each(|key| {
210 remove_locked(inner, key);
211 });
212 expired
213}
214
215fn remove_locked<K, G, V>(inner: &mut Inner<K, G, V>, key: &K) -> Option<V>
216where
217 K: Eq + Hash,
218 G: Eq + Hash,
219{
220 let entry = inner.entries.remove(key)?;
221 release_group(&mut inner.group_counts, &entry.group);
222 Some(entry.value)
223}
224
225fn release_group<G: Eq + Hash>(counts: &mut HashMap<G, usize>, group: &G) {
226 if let Some(count) = counts.get_mut(group) {
227 *count = count.saturating_sub(1);
228 if *count == 0 {
229 counts.remove(group);
230 }
231 }
232}
233
234#[cfg(test)]
235mod tests {
236 use super::*;
237
238 fn at(micros: u64) -> UnixMicros {
239 UnixMicros::new(micros)
240 }
241
242 fn store(per_group: usize, total: usize) -> Expiring<&'static str, &'static str, u8> {
243 Expiring::new(Quotas {
244 per_group: GroupQuota::new(per_group),
245 total: TotalQuota::new(total),
246 })
247 }
248
249 #[test]
250 fn a_slot_stays_occupied_until_it_expires_and_is_then_free_to_take_over() {
251 let store = store(8, 8);
252 assert_eq!(
253 store.admit("uni", "nel", 1, at(100), at(0)),
254 Ok(Admitted::Inserted)
255 );
256 assert_eq!(
257 store.admit("uni", "olaren", 2, at(900), at(50)),
258 Ok(Admitted::Occupied(1)),
259 "the caller decides whether an occupied slot is a replay or a renewal"
260 );
261 assert_eq!(store.get(&"uni", at(50)), Some(1));
262 assert_eq!(
263 store.get(&"uni", at(100)),
264 None,
265 "a replay guard must keep the original expiry or a replayed token renews its own window"
266 );
267 assert_eq!(
268 store.admit("uni", "olaren", 2, at(300), at(100)),
269 Ok(Admitted::Inserted),
270 "expiry is inclusive so a slot expiring exactly now is available"
271 );
272 assert_eq!(
273 store.group_len(&"nel"),
274 0,
275 "the takeover releases the old group's count"
276 );
277 assert_eq!(store.group_len(&"olaren"), 1);
278 }
279
280 #[test]
281 fn admit_or_renew_extends_only_the_current_holder() {
282 let store = store(8, 8);
283 store
284 .admit_or_renew("uni", "nel", 1, at(100), at(0))
285 .unwrap();
286 assert_eq!(
287 store.admit_or_renew("uni", "olaren", 2, at(900), at(50)),
288 Ok(Admitted::Occupied(1))
289 );
290 assert_eq!(
291 store.get(&"uni", at(150)),
292 None,
293 "a caller that isn't the holder mustn't extend the lease"
294 );
295 store
296 .admit_or_renew("uni", "nel", 1, at(300), at(200))
297 .unwrap();
298 assert_eq!(
299 store.admit_or_renew("uni", "nel", 1, at(600), at(250)),
300 Ok(Admitted::Occupied(1))
301 );
302 assert_eq!(
303 store.get(&"uni", at(500)),
304 Some(1),
305 "renewing under the lock that read the entry leaves no window for a release \
306 to drop the slot between the read and the renewal"
307 );
308 }
309
310 #[test]
311 fn each_quota_bounds_its_own_scope_and_names_itself_when_it_refuses() {
312 let store = store(1, 1);
313 store.admit("uni", "nel", 1, at(100), at(0)).unwrap();
314 assert_eq!(
315 store.admit("kelp", "nel", 2, at(100), at(0)),
316 Err(Rejected::Group),
317 "a rejection names the group limit first because the caller can act on its own quota"
318 );
319 assert_eq!(
320 store.admit("kelp", "olaren", 2, at(100), at(0)),
321 Err(Rejected::Total),
322 "a distinct group has its own budget but still shares the total"
323 );
324 assert_eq!(
325 store.admit("kelp", "olaren", 2, at(400), at(200)),
326 Ok(Admitted::Inserted),
327 "admit prunes the expired entry and takes the slot it freed"
328 );
329 }
330
331 #[test]
332 fn releasing_a_slot_by_expiry_or_by_hand_frees_its_group_budget() {
333 let store = store(2, 8);
334 store.admit("uni", "nel", 7, at(100), at(0)).unwrap();
335 store.admit("kelp", "nel", 8, at(500), at(0)).unwrap();
336 assert_eq!(store.prune(at(200)), vec!["uni"]);
337 assert_eq!(store.len(), 1);
338 assert_eq!(
339 store.group_len(&"nel"),
340 1,
341 "pruning decrements the group count rather than rebuilding it"
342 );
343 assert_eq!(store.remove(&"kelp"), Some(8));
344 assert_eq!(store.remove(&"kelp"), None);
345 assert_eq!(
346 store.group_len(&"nel"),
347 0,
348 "removing the last entry of a group frees its budget"
349 );
350 assert_eq!(
351 store.admit("whelk", "nel", 9, at(600), at(200)),
352 Ok(Admitted::Inserted)
353 );
354 }
355
356 #[test]
357 fn concurrent_admissions_at_the_quota_keep_group_counts_equal_to_what_is_stored() {
358 let store = std::sync::Arc::new(Expiring::new(Quotas {
359 per_group: GroupQuota::new(2),
360 total: TotalQuota::new(3),
361 }));
362 let keys = ["uni", "kelp", "whelk", "clam", "conch", "limpet"];
363 let groups = ["nel", "olaren"];
364 let threads: Vec<_> = (0..8u64)
365 .map(|thread| {
366 let store = std::sync::Arc::clone(&store);
367 std::thread::spawn(move || {
368 (0..2_000u64).for_each(|round| {
369 let key = keys[(thread + round) as usize % keys.len()];
370 let group = groups[(thread + round) as usize % groups.len()];
371 let now = at(round * 10);
372 let _ = store.admit(key, group, 1u8, at(round * 10 + 40), now);
373 store.prune(now);
374 });
375 })
376 })
377 .collect();
378 threads
379 .into_iter()
380 .for_each(|thread| thread.join().unwrap());
381 let counted: usize = groups.iter().map(|group| store.group_len(group)).sum();
382 assert_eq!(
383 counted,
384 store.len(),
385 "a group count higher than what is stored locks its group out of every later admission"
386 );
387 }
388}