Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389use std::collections::HashMap;use std::collections::hash_map::Entry as Slot;use std::hash::Hash;use std::sync::Mutex;
use knot_runtime::UnixMicros;
knot_types::scalar_newtype! { pub struct GroupQuota(usize); pub struct TotalQuota(usize);}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]pub struct Quotas { pub per_group: GroupQuota, pub total: TotalQuota,}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]pub enum Rejected { Total, Group,}
#[derive(Debug, Clone, PartialEq, Eq)]pub enum Admitted<V> { Inserted, Occupied(V),}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]enum Occupancy { Keep, Extend,}
struct Entry<G, V> { group: G, value: V, expires_at: UnixMicros,}
struct Inner<K, G, V> { entries: HashMap<K, Entry<G, V>>, group_counts: HashMap<G, usize>,}
pub struct Expiring<K, G, V> { inner: Mutex<Inner<K, G, V>>, quotas: Quotas,}
impl<K, G, V> Expiring<K, G, V>where K: Eq + Hash + Clone, G: Eq + Hash + Clone,{ pub fn new(quotas: Quotas) -> Self { Self { inner: Mutex::new(Inner { entries: HashMap::new(), group_counts: HashMap::new(), }), quotas, } }
fn lock(&self) -> std::sync::MutexGuard<'_, Inner<K, G, V>> { self.inner .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) }
pub fn prune(&self, now: UnixMicros) -> Vec<K> { prune_locked(&mut self.lock(), now) }
pub fn admit( &self, key: K, group: G, value: V, expires_at: UnixMicros, now: UnixMicros, ) -> Result<Admitted<V>, Rejected> where V: Clone + PartialEq, { self.enter(key, group, value, expires_at, now, Occupancy::Keep) }
pub fn admit_or_renew( &self, key: K, group: G, value: V, expires_at: UnixMicros, now: UnixMicros, ) -> Result<Admitted<V>, Rejected> where V: Clone + PartialEq, { self.enter(key, group, value, expires_at, now, Occupancy::Extend) }
fn enter( &self, key: K, group: G, value: V, expires_at: UnixMicros, now: UnixMicros, occupied: Occupancy, ) -> Result<Admitted<V>, Rejected> where V: Clone + PartialEq, { let mut inner = self.lock(); let key = match inner.entries.entry(key) { Slot::Occupied(mut held) if held.get().expires_at > now => { let entry = held.get_mut(); if occupied == Occupancy::Extend && entry.value == value { entry.expires_at = expires_at; } return Ok(Admitted::Occupied(entry.value.clone())); } Slot::Occupied(held) => { let (key, entry) = held.remove_entry(); release_group(&mut inner.group_counts, &entry.group); key } Slot::Vacant(free) => free.into_key(), }; // This is lazy on purpose! // Pruning traverses every entry so we shouldn't do it on every // admission, better to do this O(1) quota check. if self.rejection(&inner, &group).is_some() { prune_locked(&mut inner, now); } match self.rejection(&inner, &group) { Some(rejected) => Err(rejected), None => { *inner.group_counts.entry(group.clone()).or_insert(0) += 1; inner.entries.insert( key, Entry { group, value, expires_at, }, ); Ok(Admitted::Inserted) } } }
fn rejection(&self, inner: &Inner<K, G, V>, group: &G) -> Option<Rejected> { let held = inner.group_counts.get(group).copied().unwrap_or(0); match ( held >= self.quotas.per_group.get(), inner.entries.len() >= self.quotas.total.get(), ) { (true, _) => Some(Rejected::Group), (_, true) => Some(Rejected::Total), (false, false) => None, } }
pub fn remove(&self, key: &K) -> Option<V> { remove_locked(&mut self.lock(), key) }
pub fn get(&self, key: &K, now: UnixMicros) -> Option<V> where V: Clone, { let inner = self.lock(); inner .entries .get(key) .filter(|entry| entry.expires_at > now) .map(|entry| entry.value.clone()) }
pub fn len(&self) -> usize { self.lock().entries.len() }
pub fn is_empty(&self) -> bool { self.len() == 0 }
pub fn group_len(&self, group: &G) -> usize { self.lock().group_counts.get(group).copied().unwrap_or(0) }}
fn prune_locked<K, G, V>(inner: &mut Inner<K, G, V>, now: UnixMicros) -> Vec<K>where K: Eq + Hash + Clone, G: Eq + Hash,{ let expired: Vec<K> = inner .entries .iter() .filter(|(_, entry)| entry.expires_at <= now) .map(|(key, _)| key.clone()) .collect(); expired.iter().for_each(|key| { remove_locked(inner, key); }); expired}
fn remove_locked<K, G, V>(inner: &mut Inner<K, G, V>, key: &K) -> Option<V>where K: Eq + Hash, G: Eq + Hash,{ let entry = inner.entries.remove(key)?; release_group(&mut inner.group_counts, &entry.group); Some(entry.value)}
fn release_group<G: Eq + Hash>(counts: &mut HashMap<G, usize>, group: &G) { if let Some(count) = counts.get_mut(group) { *count = count.saturating_sub(1); if *count == 0 { counts.remove(group); } }}
#[cfg(test)]mod tests { use super::*;
fn at(micros: u64) -> UnixMicros { UnixMicros::new(micros) }
fn store(per_group: usize, total: usize) -> Expiring<&'static str, &'static str, u8> { Expiring::new(Quotas { per_group: GroupQuota::new(per_group), total: TotalQuota::new(total), }) }
#[test] fn a_slot_stays_occupied_until_it_expires_and_is_then_free_to_take_over() { let store = store(8, 8); assert_eq!( store.admit("uni", "nel", 1, at(100), at(0)), Ok(Admitted::Inserted) ); assert_eq!( store.admit("uni", "olaren", 2, at(900), at(50)), Ok(Admitted::Occupied(1)), "the caller decides whether an occupied slot is a replay or a renewal" ); assert_eq!(store.get(&"uni", at(50)), Some(1)); assert_eq!( store.get(&"uni", at(100)), None, "a replay guard must keep the original expiry or a replayed token renews its own window" ); assert_eq!( store.admit("uni", "olaren", 2, at(300), at(100)), Ok(Admitted::Inserted), "expiry is inclusive so a slot expiring exactly now is available" ); assert_eq!( store.group_len(&"nel"), 0, "the takeover releases the old group's count" ); assert_eq!(store.group_len(&"olaren"), 1); }
#[test] fn admit_or_renew_extends_only_the_current_holder() { let store = store(8, 8); store .admit_or_renew("uni", "nel", 1, at(100), at(0)) .unwrap(); assert_eq!( store.admit_or_renew("uni", "olaren", 2, at(900), at(50)), Ok(Admitted::Occupied(1)) ); assert_eq!( store.get(&"uni", at(150)), None, "a caller that isn't the holder mustn't extend the lease" ); store .admit_or_renew("uni", "nel", 1, at(300), at(200)) .unwrap(); assert_eq!( store.admit_or_renew("uni", "nel", 1, at(600), at(250)), Ok(Admitted::Occupied(1)) ); assert_eq!( store.get(&"uni", at(500)), Some(1), "renewing under the lock that read the entry leaves no window for a release \ to drop the slot between the read and the renewal" ); }
#[test] fn each_quota_bounds_its_own_scope_and_names_itself_when_it_refuses() { let store = store(1, 1); store.admit("uni", "nel", 1, at(100), at(0)).unwrap(); assert_eq!( store.admit("kelp", "nel", 2, at(100), at(0)), Err(Rejected::Group), "a rejection names the group limit first because the caller can act on its own quota" ); assert_eq!( store.admit("kelp", "olaren", 2, at(100), at(0)), Err(Rejected::Total), "a distinct group has its own budget but still shares the total" ); assert_eq!( store.admit("kelp", "olaren", 2, at(400), at(200)), Ok(Admitted::Inserted), "admit prunes the expired entry and takes the slot it freed" ); }
#[test] fn releasing_a_slot_by_expiry_or_by_hand_frees_its_group_budget() { let store = store(2, 8); store.admit("uni", "nel", 7, at(100), at(0)).unwrap(); store.admit("kelp", "nel", 8, at(500), at(0)).unwrap(); assert_eq!(store.prune(at(200)), vec!["uni"]); assert_eq!(store.len(), 1); assert_eq!( store.group_len(&"nel"), 1, "pruning decrements the group count rather than rebuilding it" ); assert_eq!(store.remove(&"kelp"), Some(8)); assert_eq!(store.remove(&"kelp"), None); assert_eq!( store.group_len(&"nel"), 0, "removing the last entry of a group frees its budget" ); assert_eq!( store.admit("whelk", "nel", 9, at(600), at(200)), Ok(Admitted::Inserted) ); }
#[test] fn concurrent_admissions_at_the_quota_keep_group_counts_equal_to_what_is_stored() { let store = std::sync::Arc::new(Expiring::new(Quotas { per_group: GroupQuota::new(2), total: TotalQuota::new(3), })); let keys = ["uni", "kelp", "whelk", "clam", "conch", "limpet"]; let groups = ["nel", "olaren"]; let threads: Vec<_> = (0..8u64) .map(|thread| { let store = std::sync::Arc::clone(&store); std::thread::spawn(move || { (0..2_000u64).for_each(|round| { let key = keys[(thread + round) as usize % keys.len()]; let group = groups[(thread + round) as usize % groups.len()]; let now = at(round * 10); let _ = store.admit(key, group, 1u8, at(round * 10 + 40), now); store.prune(now); }); }) }) .collect(); threads .into_iter() .for_each(|thread| thread.join().unwrap()); let counted: usize = groups.iter().map(|group| store.group_len(group)).sum(); assert_eq!( counted, store.len(), "a group count higher than what is stored locks its group out of every later admission" ); }}