rateguard 0.7.1

//! The key table of the hot path (spec §5.2): fixed slots of atomics, no
//! locks, no `unsafe`, allocated once. A key has two buckets of [`BUCKET`]
//! slots, picked by two independent parts of its hash, and lives in one of
//! their slots; removing it never breaks another key's lookup. One bucket
//! per key sent 25 of 4096 keys to the overflow slot in a table sized for
//! 4096, and over 500 of 7000; with two, a key overflows only when both are
//! full.

use std::hash::{BuildHasher, RandomState};

use rateguard_core::gcra::{AtomicGcra, Quota, ShadowGcra};
use rateguard_core::limiter::Tally;
use xxhash_rust::xxh3::xxh3_64_with_seed;

use crate::sync::{AtomicU64, Ordering, fence};

pub(crate) const BUCKET: usize = 8;
const EMPTY: u64 = 0;
/// A slot being reset: neither a key nor free.
const BUSY: u64 = u64::MAX;

/// One key's state: its GCRA, the attempts and denials since the last tick,
/// its class, and its class's shadow with the shadow's denials. One cache
/// line a slot, exactly 64 bytes, as the README's memory bound (guarantee
/// 4) counts: no two keys share a line, and the shadow's load on the
/// request path is in the line the check has just read.
#[derive(Debug)]
#[repr(align(64))]
pub(crate) struct Slot {
    key: AtomicU64,
    pub gcra: AtomicGcra,
    /// Attempts in the low half, denials in the high: under 2³² of either
    /// in one tick, as six million a second on one key, the most a node
    /// handles, makes 300,000 in 50 ms.
    counts: AtomicU64,
    /// The hash of the key's class, 0 for the default; set by whoever takes
    /// the slot for the key, just after taking it.
    pub class: AtomicU64,
    /// The key's class's shadow limit, scaled to this node's share; off
    /// for a class with none, a key with a limit of its own, and a free
    /// slot.
    pub shadow: ShadowGcra,
    /// Requests the shadow would have denied since the last tick.
    shadow_denials: AtomicU64,
}
const ATTEMPT: u64 = 1;
const DENIAL: u64 = 1 << 32;
impl Slot {
    fn new(quota: Quota) -> Self {
        Self {
            key: AtomicU64::new(EMPTY),
            gcra: AtomicGcra::new(quota),
            counts: AtomicU64::new(0),
            class: AtomicU64::new(0),
            shadow: ShadowGcra::new(),
            shadow_denials: AtomicU64::new(0),
        }
    }

    /// One more request, admitted or not; `Release` to publish what the
    /// caller wrote to the slot before, such as its class.
    pub fn count_attempt(&self, ordering: Ordering) {
        self.counts.fetch_add(ATTEMPT, ordering);
    }

    /// One more denial, besides its attempt.
    pub fn count_denial(&self) {
        self.counts.fetch_add(DENIAL, Ordering::Relaxed);
    }

    /// One more request the shadow would have denied, besides its attempt.
    pub fn count_shadow_denial(&self) {
        self.shadow_denials.fetch_add(1, Ordering::Relaxed);
    }

    /// Whether attempts or denials wait to be taken. Shadow denials come
    /// with an attempt, so attempts tell for them too.
    pub fn has_counts(&self) -> bool {
        self.counts.load(Ordering::Relaxed) != 0
    }

    /// The attempts, denials and shadow denials counted since the last
    /// call, cleared.
    pub fn take_counts(&self) -> Tally {
        // Acquire: a slot's class, written before its first attempt, is
        // seen with it.
        let counts = self.counts.swap(0, Ordering::Acquire);
        Tally {
            attempts: counts & (DENIAL - 1),
            denials: counts >> 32,
            shadow_denials: self.shadow_denials.swap(0, Ordering::Relaxed),
        }
    }
}

pub(crate) struct KeyTable {
    slots: Box<[Slot]>,
    buckets: usize,
    overflow: Slot,
    /// Picks the buckets, and never leaves the process: key hashes are
    /// public, so without it keys could be chosen to fill a victim's.
    seed: u64,
}
// Not derived: up to two million slots would go into a log line.
impl std::fmt::Debug for KeyTable {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("KeyTable")
            .field("slots", &self.slots.len())
            .field("buckets", &self.buckets)
            .finish_non_exhaustive()
    }
}
impl KeyTable {
    /// Room for twice `tracked_keys`, in whole buckets, a power of two of
    /// them. Every slot starts free, carrying `quota`.
    pub fn new(tracked_keys: usize, quota: Quota) -> Self {
        let buckets = (2 * tracked_keys).div_ceil(BUCKET).next_power_of_two();
        Self {
            slots: (0..buckets * BUCKET).map(|_| Slot::new(quota)).collect(),
            buckets,
            overflow: Slot::new(quota),
            // The OS randomness std seeds HashMap with, for the same reason.
            seed: RandomState::new().hash_one(0u64),
        }
    }

    pub fn capacity(&self) -> usize {
        self.slots.len()
    }

    fn bucket(&self, index: usize) -> &[Slot] {
        let start = index * BUCKET;
        &self.slots[start..start + BUCKET]
    }

    /// The key's two buckets, picked by the low and the high half of its
    /// hash rehashed under the table's seed; distinct whenever the table
    /// has more than one.
    fn buckets_of(&self, key: u64) -> [&[Slot]; 2] {
        let mask = self.buckets - 1;
        let mixed = xxh3_64_with_seed(&key.to_le_bytes(), self.seed);
        let first = mixed as usize & mask;
        let mut second = (mixed >> 32) as usize & mask;
        if second == first {
            second = (first + 1) & mask;
        }
        [self.bucket(first), self.bucket(second)]
    }

    /// The key's candidate slots in their canonical order: the first bucket,
    /// then the second (once, if they are the same).
    fn candidates(&self, key: u64) -> impl Iterator<Item = &Slot> + Clone {
        let [first, second] = self.buckets_of(key);
        let second = (!std::ptr::eq(first, second)).then_some(second);
        first.iter().chain(second.into_iter().flatten())
    }

    /// The slot of `key`: found, newly taken in the emptier of its two
    /// buckets, or, both full, the shared overflow slot.
    ///
    /// Taking a slot leaves its GCRA alone: a free slot is fresh and carries
    /// the strictest quota a new key of any class starts with, kept current
    /// by the background task. The caller that took the slot sets its
    /// class's quota just after (see [`claim`](KeyTable::claim)); a thread
    /// admitted through the slot in between has its debt rescaled from that
    /// moment, at most one strict interval's worth: errs on the strict
    /// side.
    #[cfg(test)]
    pub fn slot(&self, key: u64) -> &Slot {
        self.claim(key).0
    }

    /// Like [`slot`](KeyTable::slot), and whether this call took the slot
    /// for the key: then it is the caller's to give the slot the key's
    /// class and quota.
    pub fn claim(&self, key: u64) -> (&Slot, bool) {
        debug_assert!(key != EMPTY && key != BUSY, "key_hash never yields those");
        let candidates = self.candidates(key);
        if let Some(slot) = candidates
            .clone()
            .find(|slot| slot.key.load(Ordering::Acquire) == key)
        {
            return (slot, false);
        }

        let [first, second] = self.buckets_of(key);
        let taken = |bucket: &[Slot]| {
            bucket
                .iter()
                .filter(|slot| slot.key.load(Ordering::Relaxed) != EMPTY)
                .count()
        };
        let order = if taken(second) < taken(first) {
            [second, first]
        } else {
            [first, second]
        };
        let claimed = order.iter().flat_map(|bucket| bucket.iter()).find(|slot| {
            slot.key.load(Ordering::Relaxed) == EMPTY
                && slot
                    .key
                    .compare_exchange(EMPTY, key, Ordering::AcqRel, Ordering::Relaxed)
                    .is_ok()
        });
        let Some(mine) = claimed else {
            return (&self.overflow, false);
        };

        // Our claim must be visible before we look for theirs, and theirs
        // before they look for ours: a store followed by a load of another
        // location (Dekker's pattern), which acquire/release does not order.
        // Without this fence both may miss each other and both copies live;
        // loom finds that interleaving.
        fence(Ordering::SeqCst);

        // Two threads may take two slots for one key at once, in either
        // bucket. The copy earliest in the canonical order wins, and
        // whoever looks last sees both: it gives its own back if it is the
        // later one, or takes the other back if that is. Release goes by
        // compare-and-swap, so a slot already handed to another key is
        // never touched.
        let mut winner = mine;
        let mut mine_seen = false;
        for other in candidates {
            if std::ptr::eq(other, mine) {
                mine_seen = true;
                continue;
            }
            if other.key.load(Ordering::Acquire) != key {
                continue;
            }
            if mine_seen {
                self.release(other, key, other.gcra.quota());
            } else if std::ptr::eq(winner, mine) {
                winner = other;
            }
        }
        if !std::ptr::eq(winner, mine) {
            self.release(mine, key, mine.gcra.quota());
        }
        (winner, std::ptr::eq(winner, mine))
    }

    /// The slot holding `key`, if one does; never takes one.
    pub fn find(&self, key: u64) -> Option<&Slot> {
        self.candidates(key)
            .find(|slot| slot.key.load(Ordering::Acquire) == key)
    }

    /// Whether `slot` holds `key` now.
    pub fn holds(&self, slot: &Slot, key: u64) -> bool {
        slot.key.load(Ordering::Acquire) == key
    }

    pub fn is_overflow(&self, slot: &Slot) -> bool {
        std::ptr::eq(slot, &self.overflow)
    }

    /// The slot shared by every key that finds its bucket full, at the
    /// quota of a new key.
    pub fn overflow(&self) -> &Slot {
        &self.overflow
    }

    /// Every taken slot with its key, as seen at the moment of the visit.
    pub fn for_each(&self, mut f: impl FnMut(u64, &Slot)) {
        for slot in self.slots.iter() {
            let key = slot.key.load(Ordering::Acquire);
            if key != EMPTY && key != BUSY {
                f(key, slot);
            }
        }
    }

    /// Every free slot, for the background task to keep the quota of a new
    /// key current in them.
    pub fn for_each_free(&self, mut f: impl FnMut(&Slot)) {
        for slot in self.slots.iter() {
            if slot.key.load(Ordering::Acquire) == EMPTY {
                f(slot);
            }
        }
    }

    /// Frees the slot of `key`, if it still holds it: claimed for the reset
    /// by compare-and-swap, then fields first and the key last, so that
    /// whoever takes it next finds it fresh. A thread still writing to it
    /// may leave the next key one interval of debt or one attempt: harmless,
    /// and accepted. Whether this call freed it.
    pub fn release(&self, slot: &Slot, key: u64, quota: Quota) -> bool {
        if slot
            .key
            .compare_exchange(key, BUSY, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            return false;
        }
        slot.gcra.reset(quota);
        slot.shadow.reset();
        slot.counts.store(0, Ordering::Relaxed);
        slot.shadow_denials.store(0, Ordering::Relaxed);
        slot.class.store(0, Ordering::Relaxed);
        slot.key.store(EMPTY, Ordering::Release);
        true
    }
}

#[cfg(all(test, not(rateguard_loom)))]
mod tests {
    use super::*;

    fn q() -> Quota {
        Quota::new(100, 10)
    }

    // xxh3 is public: anyone can pick keys whose hashes share the bits the
    // buckets come from, and crowd a victim's key out of both of its
    // buckets. 17 such hashes fill one pair of 16 slots and overflow; a
    // per-node seed scatters them.
    #[test]
    fn keys_picked_to_share_their_buckets_do_not_crowd_each_other_out() {
        let t = KeyTable::new(64, q());
        for i in 1..=17u64 {
            let key = (i << 40) | (5 << 32) | 3;
            assert!(!t.is_overflow(t.slot(key)), "key {i} overflowed");
        }
    }

    // The README's memory bound counts 64 bytes a slot (guarantee 4), one
    // cache line each.
    #[test]
    fn a_slot_is_a_cache_line() {
        assert_eq!(std::mem::size_of::<Slot>(), 64);
        assert_eq!(std::mem::align_of::<Slot>(), 64);
    }

    #[test]
    fn only_the_call_that_takes_a_slot_is_told_so() {
        let t = KeyTable::new(1, q()); // a single bucket of 8
        assert!(t.claim(42).1, "taken");
        assert!(!t.claim(42).1, "found");
        for key in 1..BUCKET as u64 {
            assert!(t.claim(key).1);
        }
        let (slot, taken) = t.claim(99);
        assert!(
            t.is_overflow(slot) && !taken,
            "the overflow slot is nobody's"
        );
    }

    #[test]
    fn a_released_slot_forgets_its_class() {
        let t = KeyTable::new(1, q());
        let slot = t.slot(1);
        slot.class.store(7, Ordering::Relaxed);
        assert!(t.release(slot, 1, q()));
        assert_eq!(t.slot(2).class.load(Ordering::Relaxed), 0);
    }

    #[test]
    fn a_slot_counts_attempts_and_denials_apart() {
        let t = KeyTable::new(1, q());
        let slot = t.slot(1);
        for _ in 0..5 {
            slot.count_attempt(Ordering::Relaxed);
        }
        slot.count_denial();
        slot.count_denial();
        slot.count_shadow_denial();
        assert_eq!(
            slot.take_counts(),
            Tally {
                attempts: 5,
                denials: 2,
                shadow_denials: 1
            }
        );
        assert_eq!(
            slot.take_counts(),
            Tally::default(),
            "taking them clears them"
        );
    }

    #[test]
    fn a_released_slot_forgets_its_shadow() {
        let t = KeyTable::new(1, q());
        let slot = t.slot(1);
        slot.shadow.set(Some(q()), 0);
        slot.count_shadow_denial();
        assert!(t.release(slot, 1, q()));
        let next = t.slot(2);
        assert_eq!(next.shadow.check(0), None);
        assert_eq!(next.take_counts(), Tally::default());
    }

    #[test]
    fn a_key_finds_its_slot_again() {
        let t = KeyTable::new(16, q());
        let a = t.slot(42) as *const Slot;
        assert_eq!(t.slot(42) as *const Slot, a);
        assert_ne!(t.slot(43) as *const Slot, a);
    }

    #[test]
    fn the_size_is_twice_the_tracked_keys_in_whole_buckets() {
        assert_eq!(KeyTable::new(4096, q()).capacity(), 8192);
        assert_eq!(KeyTable::new(1, q()).capacity(), BUCKET);
    }

    #[test]
    fn a_full_bucket_sends_newcomers_to_the_overflow_slot() {
        let t = KeyTable::new(1, q()); // a single bucket of 8
        for key in 1..=BUCKET as u64 {
            assert!(!t.is_overflow(t.slot(key)));
        }
        assert!(t.is_overflow(t.slot(99)));
    }

    #[test]
    fn a_released_slot_is_fresh_for_the_next_key() {
        let t = KeyTable::new(1, q());
        for _ in 0..10 {
            t.slot(1).gcra.check(0);
        }
        t.slot(1).count_attempt(Ordering::Relaxed);
        t.slot(1).count_denial();
        let slot = t.slot(1);
        assert!(t.release(slot, 1, q()));
        assert!(!t.release(slot, 1, q()), "freed once only");
        let next = t.slot(2);
        assert_eq!(
            next as *const Slot, slot as *const Slot,
            "the freed slot is reused"
        );
        assert!(!next.gcra.has_debt(0));
        assert_eq!(next.take_counts(), Tally::default());
    }

    #[test]
    fn a_free_slot_carries_the_quota_a_new_key_starts_with() {
        let t = KeyTable::new(16, q());
        t.for_each_free(|slot| slot.gcra.set_quota(Quota::new(7, 1), 0));
        assert_eq!(t.slot(5).gcra.quota(), Quota::new(7, 1));
    }

    #[test]
    fn for_each_visits_the_taken_slots_only() {
        let t = KeyTable::new(16, q());
        t.slot(5);
        t.slot(6);
        let mut keys = Vec::new();
        t.for_each(|key, _| keys.push(key));
        keys.sort();
        assert_eq!(keys, [5, 6]);
        let mut free = 0;
        t.for_each_free(|_| free += 1);
        assert_eq!(free, t.capacity() - 2);
    }

    fn overflowed(tracked_keys: usize, keys: u64) -> usize {
        let t = KeyTable::new(tracked_keys, q());
        (0..keys)
            .filter(|i| t.is_overflow(t.slot(crate::key::key_hash(&format!("user-{i}")))))
            .count()
    }

    // Measured by the final review: with one bucket per key, 4096 keys in a
    // table for 4096 put ~25 into the shared overflow slot, and 6000 denied
    // 4% of requests from well-behaved users.
    #[test]
    fn as_many_keys_as_tracked_all_get_a_slot_of_their_own() {
        assert_eq!(overflowed(4096, 4096), 0);
    }

    // Two buckets per key, the emptier one taken: up to a quarter again as
    // many keys as tracked, nobody overflows. Measured over 300 seeds: none
    // at 5120 keys; at 6144, one key in 7% of tables; at 7000, 5 to 24.
    #[test]
    fn a_quarter_again_as_many_keys_as_tracked_still_all_get_a_slot() {
        assert_eq!(overflowed(4096, 5120), 0);
    }

    #[test]
    fn half_again_as_many_keys_as_tracked_overflow_a_few_at_most() {
        let n = overflowed(4096, 6144);
        assert!(n <= 3, "{n} keys overflowed");
    }

    // A stampede: eight threads at once on 300 keys none of them has seen.
    // One slot per key, always. A racing loser's attempt is lost with the
    // slot it gives back: ~6% here in a debug build, next to nothing in
    // release, and only ever on a key's first touches.
    #[test]
    fn many_threads_inserting_many_keys_leave_one_slot_per_key() {
        let t = KeyTable::new(512, q());
        let keys: Vec<u64> = (1..=300)
            .map(|i| crate::key::key_hash(&format!("key-{i}")))
            .collect();
        std::thread::scope(|s| {
            for _ in 0..8 {
                s.spawn(|| {
                    for &key in &keys {
                        t.slot(key).count_attempt(Ordering::Relaxed);
                    }
                });
            }
        });
        let mut seen = std::collections::HashMap::new();
        t.for_each(|key, slot| {
            *seen.entry(key).or_insert(0) += slot.take_counts().attempts;
        });
        assert_eq!(seen.len(), 300, "one slot per key");
        let lost = 8 * 300 - seen.values().sum::<u64>();
        assert!(lost < 8 * 300 / 4, "most attempts must count: {lost} lost");
    }
}

#[cfg(all(test, rateguard_loom))]
mod loom_tests {
    use super::*;

    fn q() -> Quota {
        Quota::new(100, 10)
    }

    #[test]
    fn two_threads_inserting_one_key_agree_on_its_slot() {
        loom::model(|| {
            let t = loom::sync::Arc::new(KeyTable::new(1, q()));
            let (a, b) = (t.clone(), t.clone());
            let ha = loom::thread::spawn(move || a.slot(7) as *const Slot as usize);
            let hb = loom::thread::spawn(move || b.slot(7) as *const Slot as usize);
            let (pa, pb) = (ha.join().unwrap(), hb.join().unwrap());
            assert_eq!(pa, pb);
            let mut count = 0;
            t.for_each(|key, _| count += usize::from(key == 7));
            assert_eq!(count, 1);
        });
    }

    // The race the final review found: two threads insert one key while a
    // third frees a slot ahead of theirs. One copy must remain.
    #[test]
    fn two_inserters_and_a_release_leave_one_copy() {
        loom::model(|| {
            let t = loom::sync::Arc::new(KeyTable::new(1, q()));
            t.slot(1);
            let threads: Vec<_> = (0..3)
                .map(|n| {
                    let t = t.clone();
                    loom::thread::spawn(move || {
                        if n == 0 {
                            let slot = t.slot(1);
                            t.release(slot, 1, q());
                        } else {
                            t.slot(7);
                        }
                    })
                })
                .collect();
            for thread in threads {
                thread.join().unwrap();
            }
            let mut copies = 0;
            t.for_each(|key, _| copies += usize::from(key == 7));
            assert_eq!(copies, 1);
        });
    }

    // The background task switches a slot's shadow on and off while a
    // request checks it: the check must never unpack the 0 of "off", and
    // the shadow ends as last set.
    #[test]
    fn a_shadow_set_while_checked_never_panics() {
        loom::model(|| {
            let t = loom::sync::Arc::new(KeyTable::new(1, q()));
            t.slot(1);
            let (a, b) = (t.clone(), t.clone());
            let checks = loom::thread::spawn(move || {
                let slot = a.slot(1);
                for now in 0..2 {
                    if let Some(rateguard_core::gcra::Decision::Deny { .. }) =
                        slot.shadow.check(now)
                    {
                        slot.count_shadow_denial();
                    }
                }
            });
            let sets = loom::thread::spawn(move || {
                let slot = b.slot(1);
                slot.shadow.set(Some(Quota::new(1, 1)), 0);
                slot.shadow.set(None, 1);
                slot.shadow.set(Some(q()), 1);
            });
            checks.join().unwrap();
            sets.join().unwrap();
            assert!(t.slot(1).shadow.check(2).is_some(), "on, as last set");
        });
    }

    #[test]
    fn a_release_racing_an_insert_leaves_a_consistent_bucket() {
        loom::model(|| {
            let t = loom::sync::Arc::new(KeyTable::new(1, q()));
            t.slot(1);
            let (a, b) = (t.clone(), t.clone());
            let ha = loom::thread::spawn(move || {
                let slot = a.slot(1);
                a.release(slot, 1, q());
            });
            let hb = loom::thread::spawn(move || {
                b.slot(2);
            });
            ha.join().unwrap();
            hb.join().unwrap();
            let mut keys = Vec::new();
            t.for_each(|key, _| keys.push(key));
            assert_eq!(keys, vec![2]);
        });
    }
}