use std::cell::RefCell;
use std::collections::HashMap;
use std::hash::{BuildHasher, RandomState};
use std::sync::Mutex;
use pagesift_stat::Delta;
use pagesift_stat::rng::Rng;
pub const PRIOR_SUFFICIENT: f32 = 3.0;
pub const MAX_HOSTS: usize = 100_000;
#[derive(Default, Clone, Copy, Debug)]
struct Counts {
sufficient: f32,
insufficient: f32,
}
impl Counts {
fn observe(&mut self, sufficient: bool, keep: f32) {
self.sufficient *= keep;
self.insufficient *= keep;
if sufficient {
self.sufficient += 1.0;
} else {
self.insufficient += 1.0;
}
}
fn sample_sufficient(&self, rng: &mut Rng) -> bool {
if self.insufficient <= 0.0 {
return true;
}
let a = f64::from(PRIOR_SUFFICIENT + self.sufficient);
let b = f64::from(self.insufficient);
rng.beta(a, b) > 0.5
}
fn evidence(&self) -> f32 {
self.sufficient + self.insufficient
}
}
#[derive(Default)]
struct Shard {
hosts: HashMap<u64, Vec<Counts>>,
}
pub type EngineObserver = std::sync::Arc<dyn Fn(&str, usize, bool) + Send + Sync>;
pub struct HostLearner {
observer: std::sync::OnceLock<EngineObserver>,
hasher: RandomState,
shards: Box<[Mutex<Shard>]>,
per_shard: usize,
keep: f32,
}
thread_local! {
static RNG: RefCell<Rng> = RefCell::new(Rng::from_entropy());
}
impl Default for HostLearner {
fn default() -> Self {
Self::with_capacity(MAX_HOSTS, Delta::DEFAULT)
}
}
impl HostLearner {
fn with_capacity(capacity: usize, delta: Delta) -> Self {
let delta = delta.get();
let min_per_shard = (1.0 / delta).ceil() as usize;
let cores = std::thread::available_parallelism().map_or(1, |n| n.get());
let shards = cores
.next_power_of_two()
.min((capacity / min_per_shard).max(1));
HostLearner {
observer: std::sync::OnceLock::new(),
hasher: RandomState::new(),
shards: (0..shards).map(|_| Mutex::default()).collect(),
per_shard: capacity.div_ceil(shards),
keep: (1.0 - delta) as f32,
}
}
fn shard(&self, key: u64) -> std::sync::MutexGuard<'_, Shard> {
let i = (key as usize) % self.shards.len();
self.shards[i].lock().unwrap_or_else(|e| e.into_inner())
}
pub fn start_engine(&self, host: &str) -> usize {
let key = self.hasher.hash_one(host);
let shard = self.shard(key);
let Some(engines) = shard.hosts.get(&key) else {
return 0;
};
RNG.with_borrow_mut(|rng| {
engines
.iter()
.position(|c| c.sample_sufficient(rng))
.unwrap_or(engines.len())
})
}
pub fn record_engine(&self, host: &str, from: usize, used: usize) {
let key = self.hasher.hash_one(host);
let mut shard = self.shard(key);
if shard.hosts.len() >= self.per_shard && !shard.hosts.contains_key(&key) {
evict_low_evidence(&mut shard.hosts);
}
let engines = shard.hosts.entry(key).or_default();
if engines.len() <= used {
engines.resize(used + 1, Counts::default());
}
for (i, c) in engines.iter_mut().enumerate().take(used + 1).skip(from) {
c.observe(i == used, self.keep);
}
drop(shard);
if let Some(o) = self.observer.get() {
for i in from..=used {
o(host, i, i == used);
}
}
}
pub fn observe_with(&self, o: EngineObserver) -> bool {
self.observer.set(o).is_ok()
}
pub fn import(&self, host: &str, engine: usize, sufficient: f64, insufficient: f64) {
let key = self.hasher.hash_one(host);
let mut shard = self.shard(key);
if shard.hosts.len() >= self.per_shard && !shard.hosts.contains_key(&key) {
evict_low_evidence(&mut shard.hosts);
}
let engines = shard.hosts.entry(key).or_default();
if engines.len() <= engine {
engines.resize(engine + 1, Counts::default());
}
let c = &mut engines[engine];
if c.evidence() == 0.0 {
c.sufficient = sufficient as f32;
c.insufficient = insufficient as f32;
}
}
}
fn evict_low_evidence(hosts: &mut HashMap<u64, Vec<Counts>>) {
let evidence = |v: &Vec<Counts>| v.iter().map(Counts::evidence).sum::<f32>();
let mut keys: Vec<(f32, u64)> = hosts.iter().map(|(k, v)| (evidence(v), *k)).collect();
let half = keys.len() / 2;
keys.select_nth_unstable_by(half, |a, b| a.0.total_cmp(&b.0));
for (_, k) in &keys[..half] {
hosts.remove(k);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cold_start_always_starts_at_the_first_engine() {
let l = HostLearner::default();
for _ in 0..1000 {
assert_eq!(l.start_engine("new.test"), 0);
}
l.record_engine("ok.test", 0, 0);
for _ in 0..1000 {
assert_eq!(l.start_engine("ok.test"), 0);
}
}
#[test]
fn learns_the_starting_engine_and_keeps_exploring() {
let l = HostLearner::default();
for _ in 0..20 {
l.record_engine("spa.test", 0, 1);
}
let n = 2000;
let skipped = (0..n).filter(|_| l.start_engine("spa.test") == 1).count();
assert!(skipped > n * 9 / 10, "{skipped}");
assert!(skipped < n, "{skipped}");
}
#[test]
fn three_misses_put_the_posterior_mean_at_one_half() {
let mut c = Counts::default();
for _ in 0..3 {
c.observe(false, 1.0);
}
let mean = (PRIOR_SUFFICIENT + c.sufficient) / (PRIOR_SUFFICIENT + c.evidence());
assert_eq!(mean, 0.5);
}
#[test]
fn relearns_after_the_site_changes() {
let l = HostLearner::default();
for _ in 0..30 {
l.record_engine("spa.test", 0, 1);
}
for _ in 0..30 {
l.record_engine("spa.test", 0, 0);
}
let n = 2000;
let first = (0..n).filter(|_| l.start_engine("spa.test") == 0).count();
assert!(first > n * 9 / 10, "{first}");
}
#[test]
fn full_shard_evicts_the_less_informed_half() {
let l = HostLearner::with_capacity(40, Delta::DEFAULT);
assert!(l.shards.len() <= 2);
l.record_engine("busy.test", 0, 1);
l.record_engine("busy.test", 0, 1);
for i in 0..200 {
l.record_engine(&format!("h{i}.test"), 0, 0);
}
let total: usize = l.shards.iter().map(|s| s.lock().unwrap().hosts.len()).sum();
assert!(total <= 40, "{total}");
let busy = l.hasher.hash_one("busy.test");
assert!(l.shard(busy).hosts.contains_key(&busy));
}
}