use crate::ntp::NtpTimestamp;
use std::hash::{Hash, Hasher};
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct RateLimitConfig {
pub interval_log2: i8,
pub burst: u32,
pub leak_shift: u8,
pub global_rate_hz: f64,
pub global_burst: f64,
}
impl Default for RateLimitConfig {
fn default() -> Self {
RateLimitConfig {
interval_log2: 3,
burst: 8,
leak_shift: 2,
global_rate_hz: 20_000.0,
global_burst: 40_000.0,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Disposition {
Respond,
KissOfDeath,
Drop,
}
pub fn mark_server_timestamps(receive: &mut NtpTimestamp, transmit: &mut NtpTimestamp) {
receive.0 |= 1;
transmit.0 &= !1;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ResponseMode {
Basic,
Interleaved {
prev_transmit: NtpTimestamp,
},
}
#[derive(Clone, Copy, Debug)]
pub struct ClientRecord {
pub last_seen: f64,
tokens: f64,
pub requests: u64,
pub responses: u64,
pub dropped: u64,
drops_since_kod: u32,
pub last_receive: NtpTimestamp,
pub last_transmit: NtpTimestamp,
pub last_receive_sent: NtpTimestamp,
pub interleaved_now: bool,
}
impl ClientRecord {
fn new(now: f64, burst: u32) -> Self {
ClientRecord {
last_seen: now,
tokens: burst as f64,
requests: 0,
responses: 0,
dropped: 0,
drops_since_kod: 0,
last_receive: NtpTimestamp::ZERO,
last_transmit: NtpTimestamp::ZERO,
last_receive_sent: NtpTimestamp::ZERO,
interleaved_now: false,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ServerStats {
pub requests: u64,
pub responses: u64,
pub dropped_rate_limit: u64,
pub kiss_of_death: u64,
pub interleaved_responses: u64,
pub refused: u64,
pub evicted: u64,
}
const NIL: u32 = u32::MAX;
const MIX: u64 = 0x9e37_79b1_85eb_ca87;
#[derive(Clone, Copy)]
pub struct ClientHashBuilder {
seed: u64,
}
impl Default for ClientHashBuilder {
fn default() -> Self {
if let Ok(pinned) = std::env::var("RUSTY_TIME_HASH_SEED")
&& let Ok(seed) = pinned.parse::<u64>()
{
return ClientHashBuilder { seed };
}
use std::hash::BuildHasher as _;
let seed = std::collections::hash_map::RandomState::new().hash_one(0xA5A5_5A5Au64);
ClientHashBuilder { seed }
}
}
impl std::hash::BuildHasher for ClientHashBuilder {
type Hasher = ClientHasher;
fn build_hasher(&self) -> ClientHasher {
ClientHasher { state: self.seed }
}
}
pub struct ClientHasher {
state: u64,
}
impl ClientHasher {
#[inline]
fn mix(&mut self, value: u64) {
self.state = (self.state ^ value).wrapping_mul(MIX);
}
}
impl Hasher for ClientHasher {
#[inline]
fn write(&mut self, bytes: &[u8]) {
let (chunks, rest) = bytes.as_chunks::<8>();
for chunk in chunks {
self.mix(u64::from_le_bytes(*chunk));
}
if !rest.is_empty() {
let mut buf = [0u8; 8];
buf[..rest.len()].copy_from_slice(rest);
self.mix(u64::from_le_bytes(buf) ^ (rest.len() as u64) << 56);
}
}
#[inline]
fn write_u8(&mut self, value: u8) {
self.mix(value as u64);
}
#[inline]
fn write_u32(&mut self, value: u32) {
self.mix(value as u64);
}
#[inline]
fn write_u64(&mut self, value: u64) {
self.mix(value);
}
#[inline]
fn write_usize(&mut self, value: usize) {
self.mix(value as u64);
}
#[inline]
fn finish(&self) -> u64 {
let mut z = self.state;
z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
z ^ (z >> 31)
}
}
struct Index {
buckets: Box<[Bucket]>,
mask: usize,
len: usize,
hasher: ClientHashBuilder,
}
#[derive(Clone, Copy)]
struct Bucket {
hash: u32,
handle: u32,
}
const EMPTY: u32 = u32::MAX;
impl Index {
fn with_capacity(capacity: usize, hasher: ClientHashBuilder) -> Self {
let n = capacity.max(1).saturating_mul(4).next_power_of_two();
assert!(
n - 1 <= u32::MAX as usize,
"bucket count exceeds the 32-bit hash word"
);
Index {
buckets: vec![
Bucket {
hash: 0,
handle: EMPTY
};
n
]
.into_boxed_slice(),
mask: n - 1,
len: 0,
hasher,
}
}
#[inline]
fn hash_of<K: Hash>(&self, key: &K) -> u32 {
use std::hash::BuildHasher as _;
self.hasher.hash_one(key) as u32
}
#[inline]
fn probe<K: Eq>(&self, hash: u32, key: &K, slots: &[Slot<K>]) -> Option<usize> {
let mut i = (hash as usize) & self.mask;
loop {
let b = self.buckets[i];
if b.handle == EMPTY {
return None;
}
if b.hash == hash
&& slots
.get(b.handle as usize)
.is_some_and(|sl| sl.key == *key)
{
return Some(i);
}
i = (i + 1) & self.mask;
}
}
#[inline]
fn get<K: Eq + Hash>(&self, key: &K, slots: &[Slot<K>]) -> Option<u32> {
let h = self.hash_of(key);
let i = self.probe(h, key, slots)?;
Some(self.buckets[i].handle)
}
fn insert<K: Eq + Hash>(&mut self, key: &K, handle: u32, slots: &[Slot<K>]) {
let hash = self.hash_of(key);
if let Some(i) = self.probe(hash, key, slots) {
self.buckets[i].handle = handle;
return;
}
let mut i = (hash as usize) & self.mask;
while self.buckets[i].handle != EMPTY {
i = (i + 1) & self.mask;
}
self.buckets[i] = Bucket { hash, handle };
self.len += 1;
}
fn remove<K: Eq + Hash>(&mut self, key: &K, slots: &[Slot<K>]) {
let hash = self.hash_of(key);
let Some(mut hole) = self.probe(hash, key, slots) else {
return;
};
self.buckets[hole].handle = EMPTY;
self.len -= 1;
let mut i = (hole + 1) & self.mask;
loop {
let b = self.buckets[i];
if b.handle == EMPTY {
break;
}
let ideal = (b.hash as usize) & self.mask;
let to_hole = hole.wrapping_sub(ideal) & self.mask;
let to_here = i.wrapping_sub(ideal) & self.mask;
if to_hole < to_here {
self.buckets[hole] = b;
self.buckets[i].handle = EMPTY;
hole = i;
}
i = (i + 1) & self.mask;
}
}
#[inline]
fn len(&self) -> usize {
self.len
}
#[inline]
fn is_empty(&self) -> bool {
self.len == 0
}
}
struct Slot<K> {
key: K,
record: ClientRecord,
prev: u32,
next: u32,
generation: u32,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ClientHandle {
slot: u32,
generation: u32,
}
impl ClientHandle {
pub const INVALID: ClientHandle = ClientHandle {
slot: NIL,
generation: 0,
};
}
pub struct ClientTable<K: Eq + Hash + Ord + Clone> {
index: Index,
slots: Vec<Slot<K>>,
free: Vec<u32>,
mru: u32,
lru: u32,
global_tokens: f64,
global_last: Option<f64>,
global_drops_since_kod: u32,
capacity: usize,
config: RateLimitConfig,
refill_per_s: f64,
pub stats: ServerStats,
}
impl<K: Eq + Hash + Ord + Clone> ClientTable<K> {
pub fn new(capacity: usize, config: RateLimitConfig) -> Self {
ClientTable {
index: Index::with_capacity(capacity.max(1), ClientHashBuilder::default()),
slots: Vec::with_capacity(capacity.max(1)),
free: Vec::new(),
mru: NIL,
lru: NIL,
global_tokens: config.global_burst,
global_last: None,
global_drops_since_kod: 0,
capacity: capacity.max(1),
config,
refill_per_s: 2f64.powi(-(config.interval_log2 as i32)),
stats: ServerStats::default(),
}
}
fn unlink(&mut self, i: u32) {
let (prev, next) = {
let slot = &self.slots[i as usize];
(slot.prev, slot.next)
};
if prev == NIL {
self.mru = next;
} else {
self.slots[prev as usize].next = next;
}
if next == NIL {
self.lru = prev;
} else {
self.slots[next as usize].prev = prev;
}
}
fn link_front(&mut self, i: u32) {
let old = self.mru;
{
let slot = &mut self.slots[i as usize];
slot.prev = NIL;
slot.next = old;
}
if old == NIL {
self.lru = i;
} else {
self.slots[old as usize].prev = i;
}
self.mru = i;
}
fn touch(&mut self, i: u32) {
if self.mru == i {
return; }
self.unlink(i);
self.link_front(i);
}
pub fn bytes_per_client() -> usize {
core::mem::size_of::<Slot<K>>() + 4 * core::mem::size_of::<Bucket>()
}
pub fn len(&self) -> usize {
self.index.len()
}
pub fn is_empty(&self) -> bool {
self.index.is_empty()
}
pub fn get(&self, key: &K) -> Option<&ClientRecord> {
let i = self.index.get(key, &self.slots)?;
Some(&self.slots[i as usize].record)
}
pub fn most_recent(&self, limit: usize) -> Vec<(K, ClientRecord)> {
let mut out = Vec::with_capacity(limit.min(self.index.len()));
let mut at = self.mru;
while at != NIL && out.len() < limit {
let slot = &self.slots[at as usize];
out.push((slot.key.clone(), slot.record));
at = slot.next;
}
out
}
fn evict_one(&mut self) {
let victim = self.lru;
if victim == NIL {
return;
}
self.unlink(victim);
let ClientTable { index, slots, .. } = self;
index.remove(&slots[victim as usize].key, slots);
self.free.push(victim);
self.stats.evicted += 1;
}
fn alloc_slot(&mut self, key: K, record: ClientRecord) -> u32 {
let i = match self.free.pop() {
Some(i) => {
let slot = &mut self.slots[i as usize];
slot.key = key;
slot.record = record;
slot.prev = NIL;
slot.next = NIL;
slot.generation = slot.generation.wrapping_add(1);
i
}
None => {
self.slots.push(Slot {
key,
record,
prev: NIL,
next: NIL,
generation: 0,
});
(self.slots.len() - 1) as u32
}
};
self.link_front(i);
i
}
fn resolve(&self, handle: ClientHandle) -> Option<usize> {
let i = handle.slot as usize;
let slot = self.slots.get(i)?;
(slot.generation == handle.generation).then_some(i)
}
fn handle_for(&self, slot: u32) -> ClientHandle {
ClientHandle {
slot,
generation: self.slots[slot as usize].generation,
}
}
pub fn admit(&mut self, key: &K, now: f64) -> Disposition {
self.admit_handle(key, now).0
}
pub fn admit_handle(&mut self, key: &K, now: f64) -> (Disposition, ClientHandle) {
self.stats.requests += 1;
if self.config.global_rate_hz > 0.0 {
let last = self.global_last.unwrap_or(now);
let elapsed = (now - last).max(0.0);
self.global_tokens = (self.global_tokens + elapsed * self.config.global_rate_hz)
.min(self.config.global_burst);
self.global_last = Some(now);
if self.global_tokens < 1.0 {
self.stats.dropped_rate_limit += 1;
self.global_drops_since_kod += 1;
let period = 1u32 << self.config.leak_shift.min(16);
if self.global_drops_since_kod >= period {
self.global_drops_since_kod = 0;
self.stats.kiss_of_death += 1;
return (Disposition::KissOfDeath, ClientHandle::INVALID);
}
return (Disposition::Drop, ClientHandle::INVALID);
}
}
let slot = match self.index.get(key, &self.slots) {
Some(i) => {
self.touch(i);
i
}
None => {
if self.index.len() >= self.capacity {
self.evict_one();
}
let i = self.alloc_slot(key.clone(), ClientRecord::new(now, self.config.burst));
let ClientTable { index, slots, .. } = self;
index.insert(key, i, slots);
i
}
};
let config = self.config;
let rate = self.refill_per_s;
let handle = self.handle_for(slot);
let record = &mut self.slots[slot as usize].record;
let elapsed = (now - record.last_seen).max(0.0);
record.tokens = (record.tokens + elapsed * rate).min(config.burst as f64);
record.last_seen = now;
record.requests += 1;
if record.tokens >= 1.0 {
record.tokens -= 1.0;
record.responses += 1;
self.stats.responses += 1;
self.global_tokens -= 1.0;
return (Disposition::Respond, handle);
}
record.dropped += 1;
record.drops_since_kod += 1;
self.stats.dropped_rate_limit += 1;
let period = 1u32 << config.leak_shift.min(16);
if record.drops_since_kod >= period {
record.drops_since_kod = 0;
self.stats.kiss_of_death += 1;
(Disposition::KissOfDeath, handle)
} else {
(Disposition::Drop, handle)
}
}
pub fn response_mode(&mut self, key: &K, request_origin: NtpTimestamp) -> ResponseMode {
match self.index.get(key, &self.slots) {
Some(i) => {
let handle = self.handle_for(i);
self.response_mode_at(handle, request_origin)
}
None => ResponseMode::Basic,
}
}
pub fn response_mode_at(
&mut self,
handle: ClientHandle,
request_origin: NtpTimestamp,
) -> ResponseMode {
let Some(i) = self.resolve(handle) else {
return ResponseMode::Basic;
};
let record = &mut self.slots[i].record;
let (sent_receive, prev_transmit) = (record.last_receive_sent, record.last_transmit);
if sent_receive.is_zero() || prev_transmit.is_zero() {
return ResponseMode::Basic;
}
let interleaved = request_origin == sent_receive && !request_origin.is_zero();
record.interleaved_now = interleaved;
if interleaved {
self.stats.interleaved_responses += 1;
ResponseMode::Interleaved { prev_transmit }
} else {
ResponseMode::Basic
}
}
pub fn note_response(&mut self, key: &K, receive: NtpTimestamp, receive_sent: NtpTimestamp) {
if let Some(i) = self.index.get(key, &self.slots) {
let handle = self.handle_for(i);
self.note_response_at(handle, receive, receive_sent);
}
}
pub fn note_response_at(
&mut self,
handle: ClientHandle,
receive: NtpTimestamp,
receive_sent: NtpTimestamp,
) {
if let Some(i) = self.resolve(handle) {
let record = &mut self.slots[i].record;
record.last_receive = receive;
record.last_receive_sent = receive_sent;
}
}
pub fn note_transmit(&mut self, key: &K, transmit: NtpTimestamp) {
if let Some(i) = self.index.get(key, &self.slots) {
self.slots[i as usize].record.last_transmit = transmit;
}
}
pub fn note_transmit_at(&mut self, handle: ClientHandle, transmit: NtpTimestamp) {
if let Some(i) = self.resolve(handle) {
self.slots[i].record.last_transmit = transmit;
}
}
pub fn note_refused(&mut self) {
self.stats.refused += 1;
}
}
#[cfg(test)]
mod tests {
use super::*;
fn table() -> ClientTable<u32> {
ClientTable::new(
4,
RateLimitConfig {
interval_log2: 3, burst: 2,
leak_shift: 2, global_rate_hz: 0.0,
global_burst: 0.0,
},
)
}
#[test]
fn address_churn_cannot_defeat_the_limiter() {
let mut t = ClientTable::<u32>::new(
1_024,
RateLimitConfig {
interval_log2: 3,
burst: 8,
leak_shift: 4,
global_rate_hz: 100.0,
global_burst: 100.0,
},
);
let mut answered = 0u64;
for client in 0..20_000u32 {
if t.admit(&client, client as f64 * 5e-5) == Disposition::Respond {
answered += 1;
}
}
assert!(
answered <= 250,
"address churn produced {answered} answers against a 100/s ceiling"
);
assert!(answered > 0, "the ceiling must not block everything");
}
#[test]
fn the_global_ceiling_refills_over_time() {
let mut t = ClientTable::<u32>::new(
16,
RateLimitConfig {
interval_log2: -10, burst: 1_000_000,
leak_shift: 8,
global_rate_hz: 10.0,
global_burst: 10.0,
},
);
let mut answered = 0;
for i in 0..100u32 {
if t.admit(&(i % 4), 0.0) == Disposition::Respond {
answered += 1;
}
}
assert!(answered <= 11, "burst exceeded: {answered}");
assert_eq!(t.admit(&0, 10.0), Disposition::Respond);
}
#[test]
fn burst_is_allowed_then_the_limiter_bites() {
let mut t = table();
assert_eq!(t.admit(&1, 0.0), Disposition::Respond);
assert_eq!(t.admit(&1, 0.0), Disposition::Respond);
assert!(matches!(
t.admit(&1, 0.0),
Disposition::Drop | Disposition::KissOfDeath
));
assert_eq!(t.stats.responses, 2);
}
#[test]
fn tokens_refill_over_time() {
let mut t = table();
let _ = t.admit(&1, 0.0);
let _ = t.admit(&1, 0.0);
assert_ne!(t.admit(&1, 0.0), Disposition::Respond);
assert_eq!(t.admit(&1, 8.0), Disposition::Respond);
assert_ne!(t.admit(&1, 8.0), Disposition::Respond);
}
#[test]
fn kiss_of_death_leaks_at_the_configured_rate_not_every_drop() {
let mut t = table();
let _ = t.admit(&1, 0.0);
let _ = t.admit(&1, 0.0);
let mut kods = 0;
for _ in 0..12 {
if t.admit(&1, 0.0) == Disposition::KissOfDeath {
kods += 1;
}
}
assert_eq!(kods, 3, "KoD leak rate wrong");
assert!(
(kods as u64) < t.stats.dropped_rate_limit,
"KoD count must stay below the drop count"
);
}
#[test]
fn one_client_cannot_starve_another() {
let mut t = table();
for _ in 0..50 {
let _ = t.admit(&1, 0.0);
}
assert_eq!(t.admit(&2, 0.0), Disposition::Respond);
assert_eq!(t.admit(&2, 0.0), Disposition::Respond);
}
#[test]
fn table_is_bounded_and_evicts_the_stalest() {
let mut t = table(); for client in 0..4u32 {
let _ = t.admit(&client, client as f64);
}
assert_eq!(t.len(), 4);
let _ = t.admit(&99, 10.0);
assert_eq!(t.len(), 4, "table exceeded its bound");
assert!(t.get(&0).is_none(), "stalest client was not evicted");
assert!(t.get(&99).is_some());
assert_eq!(t.stats.evicted, 1);
}
#[test]
fn interleaved_requires_the_client_to_echo_our_receive_timestamp() {
let mut t = table();
let _ = t.admit(&1, 0.0);
let rx1 = NtpTimestamp::from_unix(1_756_224_000, 0);
assert_eq!(
t.response_mode(&1, NtpTimestamp(0x1111)),
ResponseMode::Basic
);
t.note_response(&1, rx1, rx1);
let tx1 = NtpTimestamp::from_unix(1_756_224_000, 500);
t.note_transmit(&1, tx1);
let _ = t.admit(&1, 8.0);
match t.response_mode(&1, rx1) {
ResponseMode::Interleaved { prev_transmit } => {
assert_eq!(prev_transmit, tx1);
}
other => panic!("expected interleaved, got {other:?}"),
}
assert_eq!(t.stats.interleaved_responses, 1);
}
#[test]
fn interleaved_flag_tracks_use_not_capability() {
let mut t = table();
let rx1 = NtpTimestamp::from_unix(1_756_224_000, 0);
let _ = t.admit(&1, 0.0);
t.note_response(&1, rx1, rx1);
t.note_transmit(&1, rx1);
let _ = t.admit(&1, 8.0);
let _ = t.response_mode(&1, NtpTimestamp::ZERO);
assert!(
!t.get(&1).expect("record").interleaved_now,
"a client that never asked must not be reported as using interleaved"
);
let _ = t.admit(&1, 16.0);
let _ = t.response_mode(&1, rx1);
assert!(t.get(&1).expect("record").interleaved_now);
let _ = t.admit(&1, 24.0);
let _ = t.response_mode(&1, NtpTimestamp::ZERO);
assert!(!t.get(&1).expect("record").interleaved_now);
}
#[test]
fn a_client_echoing_the_wrong_value_gets_basic_mode() {
let mut t = table();
let _ = t.admit(&1, 0.0);
let rx1 = NtpTimestamp::from_unix(1_756_224_000, 0);
t.note_response(&1, rx1, rx1);
t.note_transmit(&1, NtpTimestamp::from_unix(1_756_224_000, 500));
let _ = t.admit(&1, 8.0);
assert_eq!(
t.response_mode(&1, NtpTimestamp(0xDEAD_BEEF)),
ResponseMode::Basic
);
assert_eq!(t.response_mode(&1, NtpTimestamp::ZERO), ResponseMode::Basic);
}
#[test]
fn eviction_downgrades_to_basic_rather_than_lying() {
let mut t = table(); let rx = NtpTimestamp::from_unix(1_756_224_000, 0);
let _ = t.admit(&1, 0.0);
t.note_response(&1, rx, rx);
t.note_transmit(&1, rx);
for c in 10..15u32 {
let _ = t.admit(&c, 100.0 + c as f64);
}
assert!(t.get(&1).is_none());
let _ = t.admit(&1, 200.0);
assert_eq!(t.response_mode(&1, rx), ResponseMode::Basic);
}
#[test]
fn heavy_client_churn_stays_tractable() {
let capacity = 4_096;
let mut t = ClientTable::<u32>::new(capacity, RateLimitConfig::default());
let churn = 200_000u32;
for client in 0..churn {
let _ = t.admit(&client, client as f64 * 0.001);
}
assert_eq!(t.len(), capacity, "table must sit exactly at capacity");
assert_eq!(
t.stats.evicted,
(churn as u64) - capacity as u64,
"every client past capacity must have cost exactly one eviction"
);
assert!(t.get(&(churn - 1)).is_some(), "newest client was evicted");
assert!(t.get(&0).is_none(), "oldest client survived");
}
#[test]
fn eviction_is_true_lru_not_insertion_order() {
let mut t = ClientTable::<u32>::new(3, RateLimitConfig::default());
let _ = t.admit(&1, 0.0);
let _ = t.admit(&2, 1.0);
let _ = t.admit(&3, 2.0);
let _ = t.admit(&1, 3.0);
let _ = t.admit(&4, 4.0);
assert!(t.get(&1).is_some(), "recently used client was evicted");
assert!(t.get(&2).is_none(), "stalest client should have gone");
assert!(t.get(&3).is_some());
assert!(t.get(&4).is_some());
}
#[test]
fn mru_report_is_ordered_and_bounded() {
let mut t = ClientTable::<u32>::new(16, RateLimitConfig::default());
for c in 0..10u32 {
let _ = t.admit(&c, c as f64);
}
let mru = t.most_recent(3);
assert_eq!(mru.len(), 3);
assert_eq!(mru[0].0, 9, "most recent first");
assert_eq!(mru[2].0, 7);
}
}