use std::sync::atomic::{AtomicU32, AtomicU64, AtomicU8, AtomicUsize, Ordering};
pub(super) const SLOT_BITS: u32 = 17;
pub(super) const SLOTS: usize = 1 << SLOT_BITS;
const SLOT_MASK: usize = SLOTS - 1;
pub(super) const ENTRY_BYTES: usize = 24;
pub(super) const TABLE_BYTES: usize = SLOTS * ENTRY_BYTES;
pub(super) const LOAD_FACTOR_LIMIT: usize = SLOTS / 4 * 3;
pub(super) const MAX_PROBES: usize = 64;
pub(super) const FLAG_SIZE_UNAVAILABLE: u8 = 0b0000_0001;
pub(super) const FLAG_SIZE_FROM_GAP: u8 = 0b0000_0010;
#[repr(C)]
pub(super) struct Slot {
key_hash: AtomicU64,
bytes: AtomicU64,
count: AtomicU32,
flags: AtomicU8,
_pad: [u8; 3],
}
const _: () = assert!(std::mem::size_of::<Slot>() == ENTRY_BYTES);
const _: () = assert!(TABLE_BYTES == 3 * 1024 * 1024);
impl Slot {
fn empty() -> Self {
Self {
key_hash: AtomicU64::new(0),
bytes: AtomicU64::new(0),
count: AtomicU32::new(0),
flags: AtomicU8::new(0),
_pad: [0; 3],
}
}
fn clear(&self) {
self.key_hash.store(0, Ordering::Relaxed);
self.bytes.store(0, Ordering::Relaxed);
self.count.store(0, Ordering::Relaxed);
self.flags.store(0, Ordering::Relaxed);
}
}
#[derive(Clone, Copy, Debug)]
pub(super) struct Entry {
pub(super) count: u32,
pub(super) bytes: u64,
pub(super) size_unavailable: bool,
pub(super) size_from_gap: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Insert {
Recorded,
NotAdmitted,
Full,
}
pub(super) struct Table {
slots: Box<[Slot]>,
occupancy: AtomicUsize,
}
pub(super) fn hash_partition(keyspace: &str, table: &str, key: &[u8]) -> u64 {
const SEP: u8 = 0xFF;
let mut h: u64 = 0xcbf2_9ce4_8422_2325;
let mix = |bytes: &[u8], h: &mut u64| {
for b in bytes {
*h ^= *b as u64;
*h = h.wrapping_mul(0x0000_0100_0000_01b3);
}
};
mix(keyspace.as_bytes(), &mut h);
mix(&[SEP], &mut h);
mix(table.as_bytes(), &mut h);
mix(&[SEP], &mut h);
for b in key {
h ^= *b as u64;
h = h.wrapping_mul(0x0000_0100_0000_01b3);
}
let mut z = h.wrapping_add(0x9e37_79b9_7f4a_7c15);
z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
z ^= z >> 31;
if z == 0 {
1
} else {
z
}
}
#[inline]
pub(super) fn admitted(hash: u64, prefix_bits: u32) -> bool {
if prefix_bits == 0 {
return true;
}
debug_assert!(prefix_bits < 64);
(hash >> (64 - prefix_bits)) == 0
}
impl Table {
pub(super) fn new() -> Self {
let mut slots = Vec::with_capacity(SLOTS);
slots.resize_with(SLOTS, Slot::empty);
Self {
slots: slots.into_boxed_slice(),
occupancy: AtomicUsize::new(0),
}
}
pub(super) fn footprint_bytes(&self) -> usize {
self.slots.len() * ENTRY_BYTES
}
pub(super) fn occupancy(&self) -> usize {
self.occupancy.load(Ordering::Relaxed)
}
#[cfg(test)]
pub(super) fn record(&self, hash: u64, prefix_bits: u32, bytes: Option<u64>) -> Insert {
self.record_with_flags(hash, prefix_bits, bytes, 0)
}
pub(super) fn record_with_flags(
&self,
hash: u64,
prefix_bits: u32,
bytes: Option<u64>,
extra_flags: u8,
) -> Insert {
if !admitted(hash, prefix_bits) {
return Insert::NotAdmitted;
}
let mut idx = (hash as usize) & SLOT_MASK;
for _ in 0..MAX_PROBES {
let slot = &self.slots[idx];
let mut observed = slot.key_hash.load(Ordering::Acquire);
if observed == 0 {
match slot
.key_hash
.compare_exchange(0, hash, Ordering::AcqRel, Ordering::Acquire)
{
Ok(_) => {
self.occupancy.fetch_add(1, Ordering::Relaxed);
observed = hash;
}
Err(actual) => observed = actual,
}
}
if observed == hash {
Self::apply(slot, bytes, extra_flags);
return Insert::Recorded;
}
idx = (idx + 1) & SLOT_MASK;
}
Insert::Full
}
fn apply(slot: &Slot, bytes: Option<u64>, extra_flags: u8) {
if extra_flags != 0 {
slot.flags.fetch_or(extra_flags, Ordering::Relaxed);
}
let mut cur = slot.count.load(Ordering::Relaxed);
loop {
let next = cur.saturating_add(1);
match slot
.count
.compare_exchange_weak(cur, next, Ordering::Relaxed, Ordering::Relaxed)
{
Ok(_) => break,
Err(actual) => cur = actual,
}
}
match bytes {
Some(b) => {
let mut cur = slot.bytes.load(Ordering::Relaxed);
while b > cur {
match slot.bytes.compare_exchange_weak(
cur,
b,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(actual) => cur = actual,
}
}
}
None => {
slot.flags
.fetch_or(FLAG_SIZE_UNAVAILABLE, Ordering::Relaxed);
}
}
}
pub(super) fn for_each_entry(&mut self, mut f: impl FnMut(Entry)) {
for slot in self.slots.iter() {
if slot.key_hash.load(Ordering::Relaxed) == 0 {
continue;
}
f(Entry {
count: slot.count.load(Ordering::Relaxed),
bytes: slot.bytes.load(Ordering::Relaxed),
size_unavailable: slot.flags.load(Ordering::Relaxed) & FLAG_SIZE_UNAVAILABLE != 0,
size_from_gap: slot.flags.load(Ordering::Relaxed) & FLAG_SIZE_FROM_GAP != 0,
});
}
}
pub(super) fn reset(&mut self) {
for slot in self.slots.iter() {
slot.clear();
}
self.occupancy.store(0, Ordering::Relaxed);
}
pub(super) fn downsample(&mut self, prefix_bits: u32) -> usize {
let mut i = 0usize;
while i < SLOTS {
let h = self.slots[i].key_hash.load(Ordering::Relaxed);
if h != 0 && !admitted(h, prefix_bits) {
self.delete_at(i);
continue;
}
i += 1;
}
self.occupancy()
}
fn delete_at(&mut self, at: usize) {
self.slots[at].clear();
self.occupancy.fetch_sub(1, Ordering::Relaxed);
let mut hole = at;
let mut probe = (at + 1) & SLOT_MASK;
while self.slots[probe].key_hash.load(Ordering::Relaxed) != 0 {
let h = self.slots[probe].key_hash.load(Ordering::Relaxed);
let home = (h as usize) & SLOT_MASK;
if !cyclic_in_exclusive_start(hole, probe, home) {
self.move_slot(probe, hole);
hole = probe;
}
probe = (probe + 1) & SLOT_MASK;
if probe == hole {
break;
}
}
}
fn move_slot(&mut self, from: usize, to: usize) {
let (h, b, c, f) = {
let s = &self.slots[from];
(
s.key_hash.load(Ordering::Relaxed),
s.bytes.load(Ordering::Relaxed),
s.count.load(Ordering::Relaxed),
s.flags.load(Ordering::Relaxed),
)
};
let dst = &self.slots[to];
dst.bytes.store(b, Ordering::Relaxed);
dst.count.store(c, Ordering::Relaxed);
dst.flags.store(f, Ordering::Relaxed);
dst.key_hash.store(h, Ordering::Relaxed);
self.slots[from].clear();
}
}
fn cyclic_in_exclusive_start(start: usize, end: usize, x: usize) -> bool {
if start < end {
x > start && x <= end
} else {
x > start || x <= end
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn footprint_is_exactly_three_mib() {
assert_eq!(TABLE_BYTES, 3 * 1024 * 1024);
assert_eq!(std::mem::size_of::<Slot>(), 24);
let t = Table::new();
assert_eq!(t.footprint_bytes(), 3 * 1024 * 1024);
}
#[test]
fn admission_predicate_is_monotone_in_prefix_width() {
for hash in [1u64, 0x0000_0000_dead_beef, u64::MAX, 0x7fff_ffff_ffff_ffff] {
let mut still = true;
for k in 0..=20u32 {
let a = admitted(hash, k);
if !a {
still = false;
}
if !still {
assert!(!a, "admission must be monotone (hash={hash:#x}, k={k})");
}
}
}
}
#[test]
fn downsample_keeps_every_survivor_findable_at_its_exact_count() {
let mut table = Table::new();
let n = 40_000u64;
for i in 0..n {
let h = hash_partition("ks", "t", &i.to_le_bytes());
assert_eq!(table.record(h, 0, Some(10)), Insert::Recorded);
}
let before = table.occupancy();
assert_eq!(before, n as usize);
let survivors = table.downsample(1);
assert!(survivors < before, "a downsample must drop entries");
for i in 0..n {
let h = hash_partition("ks", "t", &i.to_le_bytes());
match table.record(h, 1, Some(10)) {
Insert::Recorded => {}
Insert::NotAdmitted => assert!(!admitted(h, 1)),
Insert::Full => panic!("table must not be full at {survivors} entries"),
}
}
assert_eq!(
table.occupancy(),
survivors,
"a second pass over the same keys must claim no new slots"
);
let mut twos = 0usize;
let mut others = 0usize;
table.for_each_entry(|e| if e.count == 2 { twos += 1 } else { others += 1 });
assert_eq!(others, 0, "every survivor must have folded, not split");
assert_eq!(twos, survivors);
}
#[test]
fn bytes_take_the_maximum_not_the_sum() {
let table = Table::new();
let h = hash_partition("ks", "t", b"partition-a");
for _ in 0..10 {
assert_eq!(table.record(h, 0, Some(4_096)), Insert::Recorded);
}
let mut table = table;
let mut seen = Vec::new();
table.for_each_entry(|e| seen.push(e));
assert_eq!(seen.len(), 1);
assert_eq!(seen[0].count, 10);
assert_eq!(seen[0].bytes, 4_096, "distinct-partition bytes, not a sum");
assert!(!seen[0].size_unavailable);
}
#[test]
fn an_unpriced_access_makes_the_entry_sticky_unavailable() {
let table = Table::new();
let h = hash_partition("ks", "t", b"partition-b");
assert_eq!(table.record(h, 0, Some(8_192)), Insert::Recorded);
assert_eq!(table.record(h, 0, None), Insert::Recorded);
assert_eq!(table.record(h, 0, Some(8_192)), Insert::Recorded);
let mut table = table;
let mut seen = Vec::new();
table.for_each_entry(|e| seen.push(e));
assert_eq!(seen.len(), 1);
assert!(
seen[0].size_unavailable,
"unavailable must be sticky for the window"
);
}
#[test]
fn reset_empties_the_table_without_changing_its_footprint() {
let mut table = Table::new();
for i in 0..1_000u64 {
table.record(hash_partition("ks", "t", &i.to_le_bytes()), 0, Some(1));
}
assert_eq!(table.occupancy(), 1_000);
table.reset();
assert_eq!(table.occupancy(), 0);
let mut n = 0;
table.for_each_entry(|_| n += 1);
assert_eq!(n, 0);
assert_eq!(table.footprint_bytes(), TABLE_BYTES);
}
#[test]
fn a_probe_cluster_reports_full_well_below_the_load_factor() {
let table = Table::new();
let mut home: Option<usize> = None;
let mut hashes = Vec::new();
let mut i = 0u64;
while hashes.len() <= MAX_PROBES && i < 50_000_000 {
let h = hash_partition("ks", "t", &i.to_le_bytes());
let slot = (h as usize) & SLOT_MASK;
match home {
None => {
home = Some(slot);
hashes.push(h);
}
Some(target) if slot == target => hashes.push(h),
_ => {}
}
i += 1;
}
assert!(
hashes.len() > MAX_PROBES,
"needed {} colliding keys, found {}",
MAX_PROBES + 1,
hashes.len()
);
for h in hashes.iter().take(MAX_PROBES) {
assert_eq!(table.record(*h, 0, Some(1)), Insert::Recorded);
}
assert!(table.occupancy() < LOAD_FACTOR_LIMIT / 100);
assert_eq!(
table.record(hashes[MAX_PROBES], 0, Some(1)),
Insert::Full,
"a cluster longer than the probe bound reports Full even on an almost \
empty table — which is why the recorder must widen rather than drop"
);
}
#[test]
fn the_same_key_in_two_tables_is_two_identities() {
let key = b"tenant-42";
assert_ne!(
hash_partition("ks", "users", key),
hash_partition("ks", "orders", key)
);
assert_ne!(
hash_partition("ks_a", "t", key),
hash_partition("ks_b", "t", key)
);
assert_eq!(
hash_partition("ks", "users", key),
hash_partition("ks", "users", key)
);
assert_ne!(
hash_partition("ab", "c", key),
hash_partition("a", "bc", key)
);
}
#[test]
fn cyclic_window_membership() {
assert!(cyclic_in_exclusive_start(2, 5, 3));
assert!(cyclic_in_exclusive_start(2, 5, 5));
assert!(!cyclic_in_exclusive_start(2, 5, 2));
assert!(!cyclic_in_exclusive_start(2, 5, 6));
assert!(cyclic_in_exclusive_start(SLOTS - 2, 1, SLOTS - 1));
assert!(cyclic_in_exclusive_start(SLOTS - 2, 1, 0));
assert!(!cyclic_in_exclusive_start(SLOTS - 2, 1, 2));
}
}