zc2 0.0.32

//! Choosing a worker in O(1), whatever the size of the fleet.
//!
//! A broker with a million workers cannot rank them per task: sorting or
//! solving an assignment on every request is polynomial work on a path that
//! has to run billions of times. Two pieces make the decision constant-time:
//!
//! - **Walker's alias method** turns a weight per worker into a table that is
//!   built once per tick in O(n) and drawn from in O(1).
//! - **Power of `d` choices**: draw `d` candidates and take the best. Maximum
//!   load falls from Θ(log n / log log n) under a single random choice to
//!   Θ(log log n / log d), so `d = 2` buys nearly all of the benefit and
//!   sampling the whole fleet buys almost none of the rest.
//!
//! The RNG is seeded and lives here rather than being pulled from the OS, so
//! a run is reproducible: the same seed and the same weights make the same
//! decisions, which is what lets a routing regression be debugged at all.

/// A small deterministic PRNG (xorshift64*). Not cryptographic — this picks
/// workers, it does not protect anything.
#[derive(Debug, Clone)]
pub struct Rng(u64);

impl Rng {
    pub fn seeded(seed: u64) -> Self {
        // Zero is a fixed point of xorshift; steer away from it.
        Rng(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15).max(1))
    }

    fn next_u64(&mut self) -> u64 {
        let mut x = self.0;
        x ^= x >> 12;
        x ^= x << 25;
        x ^= x >> 27;
        self.0 = x;
        x.wrapping_mul(0x2545_F491_4F6C_DD1D)
    }

    /// Uniform in [0, 1).
    pub fn next_f64(&mut self) -> f64 {
        (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
    }

    /// Uniform in [0, n).
    pub fn below(&mut self, n: usize) -> usize {
        match n {
            0 => 0,
            n => (self.next_u64() % n as u64) as usize,
        }
    }
}

/// Weighted sampling in constant time (Walker's alias method).
#[derive(Debug, Clone)]
pub struct AliasTable {
    prob: Vec<f64>,
    alias: Vec<u32>,
}

impl AliasTable {
    /// Build from non-negative weights. `None` when there is nothing to draw:
    /// no entries, or every weight zero.
    pub fn build(weights: &[f64]) -> Option<AliasTable> {
        let n = weights.len();
        if n == 0 {
            return None;
        }
        let clean: Vec<f64> = weights
            .iter()
            .map(|w| if w.is_finite() && *w > 0.0 { *w } else { 0.0 })
            .collect();
        let total: f64 = clean.iter().sum();
        if total <= 0.0 {
            return None;
        }

        let mut prob: Vec<f64> = clean.iter().map(|w| w * n as f64 / total).collect();
        let mut alias = vec![0u32; n];
        let (mut small, mut large): (Vec<usize>, Vec<usize>) = (0..n).partition(|&i| prob[i] < 1.0);

        while let (Some(s), Some(l)) = (small.pop(), large.pop()) {
            alias[s] = l as u32;
            prob[l] = (prob[l] + prob[s]) - 1.0;
            match prob[l] < 1.0 {
                true => small.push(l),
                false => large.push(l),
            }
        }
        // Whatever is left is 1.0 up to rounding.
        for i in small.into_iter().chain(large) {
            prob[i] = 1.0;
            alias[i] = i as u32;
        }
        Some(AliasTable { prob, alias })
    }

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

    pub fn is_empty(&self) -> bool {
        self.prob.is_empty()
    }

    /// One draw: two array reads and a comparison, whatever `len()` is.
    pub fn draw(&self, rng: &mut Rng) -> usize {
        let i = rng.below(self.prob.len());
        match rng.next_f64() < self.prob[i] {
            true => i,
            false => self.alias[i] as usize,
        }
    }
}

/// Draw `d` candidates and keep the one `cost` likes best — lower is better.
///
/// This is the whole per-task decision: `d` constant-time draws and `d` cost
/// evaluations, independent of how many workers exist.
pub fn best_of_d<F>(table: &AliasTable, rng: &mut Rng, d: usize, cost: F) -> Option<usize>
where
    F: Fn(usize) -> f64,
{
    if table.is_empty() {
        return None;
    }
    let mut best: Option<(usize, f64)> = None;
    for _ in 0..d.max(1) {
        let i = table.draw(rng);
        let c = cost(i);
        if best.is_none_or(|(_, bc)| c < bc) {
            best = Some((i, c));
        }
    }
    best.map(|(i, _)| i)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn counts(weights: &[f64], draws: usize, seed: u64) -> Vec<usize> {
        let table = AliasTable::build(weights).expect("weights");
        let mut rng = Rng::seeded(seed);
        let mut c = vec![0usize; weights.len()];
        for _ in 0..draws {
            c[table.draw(&mut rng)] += 1;
        }
        c
    }

    /// The table has to actually sample in proportion to the weights, or every
    /// decision built on it is quietly biased.
    #[test]
    fn draws_land_in_proportion_to_the_weights() {
        let weights = [1.0, 3.0, 6.0];
        let draws = 200_000;
        let c = counts(&weights, draws, 42);
        for (i, w) in weights.iter().enumerate() {
            let want = w / 10.0;
            let got = c[i] as f64 / draws as f64;
            assert!(
                (got - want).abs() < 0.01,
                "index {i}: drawn {got:.3}, weight says {want:.3}"
            );
        }
    }

    /// A drained or unhealthy worker gets weight zero, and must then never be
    /// handed work — not rarely, never.
    #[test]
    fn a_zero_weight_is_never_drawn() {
        let c = counts(&[5.0, 0.0, 5.0], 50_000, 7);
        assert_eq!(c[1], 0, "zero weight drawn {} times", c[1]);
        assert!(c[0] > 0 && c[2] > 0);
    }

    #[test]
    fn nothing_to_draw_from_is_not_a_panic() {
        assert!(AliasTable::build(&[]).is_none());
        assert!(AliasTable::build(&[0.0, 0.0]).is_none(), "all drained");
        assert!(
            AliasTable::build(&[f64::NAN, -1.0, 2.0]).is_some(),
            "a bad weight is ignored, the good one still works"
        );
    }

    /// Same seed, same decisions — otherwise a routing regression cannot be
    /// reproduced, let alone debugged.
    #[test]
    fn the_same_seed_makes_the_same_choices() {
        let a = counts(&[1.0, 2.0, 3.0], 1_000, 99);
        let b = counts(&[1.0, 2.0, 3.0], 1_000, 99);
        let different = counts(&[1.0, 2.0, 3.0], 1_000, 100);
        assert_eq!(a, b);
        assert_ne!(a, different, "a different seed explores differently");
    }

    /// Power of d: with two draws it should usually avoid the worst worker,
    /// while a single draw follows the weights blindly.
    #[test]
    fn two_choices_beat_one() {
        let weights = vec![1.0; 10];
        let table = AliasTable::build(&weights).unwrap();
        // Worker 9 is terrible, the rest are fine.
        let cost = |i: usize| if i == 9 { 100.0 } else { 1.0 };

        let mut rng = Rng::seeded(5);
        let mut bad_with_one = 0;
        let mut bad_with_two = 0;
        for _ in 0..10_000 {
            if best_of_d(&table, &mut rng, 1, cost) == Some(9) {
                bad_with_one += 1;
            }
            if best_of_d(&table, &mut rng, 2, cost) == Some(9) {
                bad_with_two += 1;
            }
        }
        // One draw hits it about a tenth of the time; two only when both land
        // on it, about a hundredth.
        assert!(bad_with_one > 700, "one draw: {bad_with_one}");
        assert!(
            bad_with_two * 5 < bad_with_one,
            "two draws should be far better: {bad_with_two} vs {bad_with_one}"
        );
    }

    #[test]
    fn best_of_d_picks_the_best_it_saw() {
        let table = AliasTable::build(&[1.0, 1.0, 1.0]).unwrap();
        let mut rng = Rng::seeded(1);
        // With d large enough to see everything, it must find the minimum.
        let chosen = best_of_d(&table, &mut rng, 50, |i| i as f64).unwrap();
        assert_eq!(chosen, 0);
    }
}

#[cfg(test)]
mod scale_tests {
    use super::*;
    use std::time::Instant;

    /// The point of the alias table: a draw costs the same whether the broker
    /// knows a hundred workers or a million. It is not literally flat — a 16 MB
    /// table misses cache where a 1 KB one does not — so this asserts the cost
    /// stays within a small constant factor, and prints both so a regression
    /// that changes the shape is visible.
    #[test]
    fn a_draw_costs_the_same_at_a_hundred_workers_and_at_a_million() {
        let per_draw = |n: usize| -> f64 {
            let table = AliasTable::build(&vec![1.0; n]).unwrap();
            let mut rng = Rng::seeded(n as u64);
            let draws = 200_000;
            // Warm the path so the first-touch page faults aren't the measurement.
            let mut sink = 0usize;
            for _ in 0..10_000 {
                sink ^= table.draw(&mut rng);
            }
            let start = Instant::now();
            for _ in 0..draws {
                sink ^= table.draw(&mut rng);
            }
            let ns = start.elapsed().as_nanos() as f64 / draws as f64;
            assert!(sink < usize::MAX, "keep the work");
            ns
        };

        let small = per_draw(100);
        let large = per_draw(1_000_000);
        println!("  per draw: 100 workers {small:.1} ns, 1e6 workers {large:.1} ns");
        assert!(
            large < small * 20.0 + 100.0,
            "a million workers must not change the shape of the cost: \
             {small:.1} ns vs {large:.1} ns"
        );
    }

    /// Building is O(n) and happens once a tick, not once a task.
    #[test]
    fn building_a_million_entry_table_is_affordable_once_a_tick() {
        let weights = vec![1.0; 1_000_000];
        let start = Instant::now();
        let table = AliasTable::build(&weights).expect("built");
        let ms = start.elapsed().as_secs_f64() * 1000.0;
        println!("  built 1e6 entries in {ms:.0} ms");
        assert_eq!(table.len(), 1_000_000);
        assert!(ms < 2_000.0, "{ms:.0} ms to build");
    }
}