zc2 0.0.32

//! Which worker should run the next task.
//!
//! Two things have to be true at once: the choice must exploit what has been
//! measured, and it must keep learning, because a mesh changes underneath the
//! estimates. Thompson sampling does both without an explore/exploit schedule
//! to tune — draw each candidate's service time from its own posterior, then
//! take the best. A worker measured a hundred times has a narrow posterior and
//! is judged on its merits; one measured twice has a wide one and is sometimes
//! sampled favourably and tried. A worker measured *never* is neither instantly
//! the best (which is what made an unmeasured worker outrank a 5 ms one) nor
//! starved.
//!
//! Only the service time is sampled. Success and trust use their posterior
//! means: they move the cost far less than latency does, and sampling a Beta
//! needs machinery this does not otherwise need.

use super::estimator::Stats;
use super::sampler::{best_of_d, AliasTable, Rng};

/// What the policy knows about one candidate worker.
#[derive(Debug, Clone)]
pub struct Candidate {
    pub stats: Stats,
    /// Requests already handed to it and not yet finished.
    pub in_flight: u32,
    /// Credits per hour, for the cost term.
    pub price_per_hour: f64,
}

/// How the terms trade off. Defaults optimise for completion time, with price
/// and failure as secondary.
#[derive(Debug, Clone, Copy)]
pub struct Weights {
    pub latency: f64,
    pub price: f64,
    /// What a failure costs, as a multiple of the expected service time — a
    /// failed task is the whole wait plus doing it again somewhere else.
    pub retry: f64,
    /// Assumed service time, in ms, for a worker nothing is known about. It
    /// only has to be the right order of magnitude: the posterior takes over
    /// as soon as there are real observations.
    pub unmeasured_ms: f64,
}

impl Default for Weights {
    fn default() -> Self {
        Weights {
            latency: 1.0,
            price: 0.0,
            retry: 2.0,
            unmeasured_ms: 250.0,
        }
    }
}

/// A standard normal, by Box-Muller.
fn standard_normal(rng: &mut Rng) -> f64 {
    // u1 must not be 0, or ln(0) is -inf.
    let u1 = rng.next_f64().max(f64::MIN_POSITIVE);
    let u2 = rng.next_f64();
    (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
}

/// Draw a plausible service time for this worker from its posterior.
///
/// With `n` observations of log service time, the posterior over the mean log
/// is Normal(μ̂, σ̂²/n): the more that has been seen, the tighter the draw sits
/// to what was measured. With nothing measured, fall back to the assumed time
/// with a wide spread, so the worker is tried but not favoured.
pub fn sample_service_ms(stats: &Stats, w: &Weights, rng: &mut Rng) -> f64 {
    let n = stats.weight();
    match (stats.median_service_ms(), n >= 1.0) {
        (Some(median), true) => {
            let sd = stats.log_sd().unwrap_or(0.5).max(0.05);
            let spread = sd / n.sqrt();
            (median.ln() + standard_normal(rng) * spread).exp()
        }
        _ => (w.unmeasured_ms.ln() + standard_normal(rng) * 0.8).exp(),
    }
}

/// What handing this task to this worker is expected to cost.
///
/// Queue first: what a caller waits for is the work already in front of it
/// plus its own. Ranking on service time alone piles every task onto the same
/// "fastest" worker until it is deep in its own queue — measured on a live
/// mesh, that moved p90 from 2.07 s to 6.14 s.
pub fn cost(candidate: &Candidate, sampled_service_ms: f64, w: &Weights) -> f64 {
    let completion = sampled_service_ms * (1.0 + candidate.in_flight as f64);
    let failure = 1.0 - candidate.stats.success_rate() * candidate.stats.trust();
    let price = candidate.price_per_hour * sampled_service_ms / 3_600_000.0;
    w.latency * completion + w.price * price + w.retry * failure * completion
}

/// Pick a worker: `d` constant-time draws, each judged on a sample from its
/// own posterior. Independent of how many workers exist.
pub fn choose(
    candidates: &[Candidate],
    table: &AliasTable,
    rng: &mut Rng,
    d: usize,
    w: &Weights,
) -> Option<usize> {
    if candidates.is_empty() || table.is_empty() {
        return None;
    }
    // Sampled once per candidate considered, inside the cost closure, so a
    // worker drawn twice is judged on two independent samples.
    let sampled = std::cell::RefCell::new(rng.clone());
    let chosen = best_of_d(table, rng, d, |i| {
        let c = &candidates[i];
        let s = sample_service_ms(&c.stats, w, &mut sampled.borrow_mut());
        cost(c, s, w)
    });
    chosen.filter(|i| *i < candidates.len())
}

/// The weight a worker gets in the sampling table, rebuilt once a tick.
///
/// Inverse expected completion: a worker twice as fast is drawn twice as
/// often, before any of the per-task sampling happens. An unmeasured worker
/// gets the assumed time, so it is in the table from the start.
pub fn table_weight(candidate: &Candidate, w: &Weights) -> f64 {
    let service = candidate
        .stats
        .expected_service_ms()
        .unwrap_or(w.unmeasured_ms);
    let completion = service * (1.0 + candidate.in_flight as f64);
    if completion <= 0.0 {
        return 0.0;
    }
    // Squared, not linear. Linear weighting leaves too much traffic on the
    // slow end of a mixed fleet: with four workers at 254 ms and nine at
    // 1000 ms it sends 36% of draws to the slow ones, and even after
    // power-of-2 filtering that is ~13% of tasks and ~99 ms of avoidable wait
    // per task — which is exactly the steady-state regret this measured
    // before the change. Squaring cuts it to about 12 ms while leaving every
    // worker a real share, so a slow worker that becomes fast is still tried.
    let rate = 1.0 / completion;
    rate * rate
}

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

    fn measured(median_ms: f64, n: usize, ok: bool) -> Stats {
        let mut s = Stats::new();
        for _ in 0..n {
            s.observe(median_ms, ok);
        }
        s
    }

    fn candidate(median_ms: f64, n: usize, in_flight: u32) -> Candidate {
        Candidate {
            stats: measured(median_ms, n, true),
            in_flight,
            price_per_hour: 3.6,
        }
    }

    #[test]
    fn a_well_measured_worker_is_sampled_close_to_what_it_measured() {
        let w = Weights::default();
        let mut rng = Rng::seeded(1);
        let stats = measured(50.0, 200, true);
        let draws: Vec<f64> = (0..500)
            .map(|_| sample_service_ms(&stats, &w, &mut rng))
            .collect();
        let mean = draws.iter().sum::<f64>() / draws.len() as f64;
        assert!(
            (mean - 50.0).abs() < 10.0,
            "200 observations should pin it near 50 ms, got {mean:.1}"
        );
    }

    /// The exploration that makes this adaptive: a worker seen twice must be
    /// sampled over a much wider range than one seen two hundred times, so it
    /// still gets tried.
    #[test]
    fn a_barely_measured_worker_is_sampled_widely() {
        let w = Weights::default();
        let spread = |n: usize| {
            let stats = measured(50.0, n, true);
            let mut rng = Rng::seeded(7);
            let draws: Vec<f64> = (0..2000)
                .map(|_| sample_service_ms(&stats, &w, &mut rng))
                .collect();
            let mean = draws.iter().sum::<f64>() / draws.len() as f64;
            let var = draws.iter().map(|d| (d - mean).powi(2)).sum::<f64>() / draws.len() as f64;
            var.sqrt()
        };
        // Both have the same observed value; only the confidence differs.
        assert!(
            spread(2) > spread(200) * 3.0,
            "little evidence must sample widely: {:.1} vs {:.1}",
            spread(2),
            spread(200)
        );
    }

    #[test]
    fn queueing_counts_against_a_worker() {
        let w = Weights::default();
        let idle = candidate(30.0, 50, 0);
        let busy = candidate(10.0, 50, 4);
        assert!(
            cost(&busy, 10.0, &w) > cost(&idle, 30.0, &w),
            "four queued at 10 ms is worse than idle at 30 ms"
        );
    }

    #[test]
    fn failure_and_price_move_the_cost_the_way_they_should() {
        let w = Weights::default();
        let mut flaky = candidate(50.0, 0, 0);
        for _ in 0..20 {
            flaky.stats.observe(50.0, false);
        }
        let solid = candidate(50.0, 20, 0);
        assert!(
            cost(&flaky, 50.0, &w) > cost(&solid, 50.0, &w),
            "a worker that keeps failing costs more than one that does not"
        );

        let priced = Weights {
            price: 1e6,
            ..Weights::default()
        };
        let cheap = Candidate {
            price_per_hour: 1.0,
            ..candidate(50.0, 20, 0)
        };
        let dear = Candidate {
            price_per_hour: 100.0,
            ..candidate(50.0, 20, 0)
        };
        assert!(cost(&dear, 50.0, &priced) > cost(&cheap, 50.0, &priced));
    }

    #[test]
    fn the_table_weights_a_fast_idle_worker_above_a_slow_busy_one() {
        let w = Weights::default();
        assert!(
            table_weight(&candidate(10.0, 20, 0), &w) > table_weight(&candidate(10.0, 20, 4), &w)
        );
        assert!(
            table_weight(&candidate(10.0, 20, 0), &w) > table_weight(&candidate(100.0, 20, 0), &w)
        );
        // A worker nothing is known about is still in the table.
        assert!(table_weight(&candidate(0.0, 0, 0), &w) > 0.0);
    }
}

/// Does the policy actually beat what the broker does today, and does it
/// recover when the mesh changes underneath it? These simulate a fleet with
/// known true service times and compare cumulative waiting against an oracle
/// that knows them.
#[cfg(test)]
mod regret_tests {
    use super::*;

    /// A fleet whose true service times we know, so regret can be measured.
    struct Fleet {
        truth_ms: Vec<f64>,
        stats: Vec<Stats>,
        in_flight: Vec<u32>,
    }

    impl Fleet {
        fn new(truth_ms: Vec<f64>) -> Self {
            let n = truth_ms.len();
            Fleet {
                truth_ms,
                stats: vec![Stats::new(); n],
                in_flight: vec![0; n],
            }
        }

        fn candidates(&self) -> Vec<Candidate> {
            (0..self.truth_ms.len())
                .map(|i| Candidate {
                    stats: self.stats[i],
                    in_flight: self.in_flight[i],
                    price_per_hour: 3.6,
                })
                .collect()
        }

        /// Run one task on `i`, with a little noise, and learn from it.
        fn run(&mut self, i: usize, rng: &mut Rng) -> f64 {
            let noise = 0.75 + rng.next_f64() * 0.5;
            let took = self.truth_ms[i] * noise;
            self.stats[i].decay(0.98);
            self.stats[i].observe(took, true);
            took
        }
    }

    fn run_policy(truth: Vec<f64>, tasks: usize, seed: u64, d: usize) -> f64 {
        let mut fleet = Fleet::new(truth);
        let mut rng = Rng::seeded(seed);
        let w = Weights::default();
        let mut total = 0.0;
        for t in 0..tasks {
            let cands = fleet.candidates();
            let weights: Vec<f64> = cands.iter().map(|c| table_weight(c, &w)).collect();
            let table = AliasTable::build(&weights).expect("a fleet");
            let i = choose(&cands, &table, &mut rng, d, &w).expect("a choice");
            total += fleet.run(i, &mut rng);
            // Rebuild the table every tick, not every task.
            let _ = t;
        }
        total
    }

    fn run_round_robin(truth: Vec<f64>, tasks: usize, seed: u64) -> f64 {
        let mut fleet = Fleet::new(truth);
        let mut rng = Rng::seeded(seed);
        let n = fleet.truth_ms.len();
        (0..tasks).map(|t| fleet.run(t % n, &mut rng)).sum()
    }

    fn run_oracle(truth: Vec<f64>, tasks: usize, seed: u64) -> f64 {
        let best = truth
            .iter()
            .enumerate()
            .min_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, _)| i)
            .unwrap();
        let mut fleet = Fleet::new(truth);
        let mut rng = Rng::seeded(seed);
        (0..tasks).map(|_| fleet.run(best, &mut rng)).sum()
    }

    /// The mesh this was measured on: a few fast local workers and a lot of
    /// slow remote ones.
    fn mesh() -> Vec<f64> {
        let mut t = vec![254.0; 4];
        t.extend(std::iter::repeat_n(1000.0, 9));
        t
    }

    #[test]
    fn the_policy_beats_round_robin_by_a_wide_margin() {
        let tasks = 3_000;
        let policy = run_policy(mesh(), tasks, 11, 2);
        let rr = run_round_robin(mesh(), tasks, 11);
        let oracle = run_oracle(mesh(), tasks, 11);

        let policy_regret = policy - oracle;
        let rr_regret = rr - oracle;
        println!(
            "  per task: oracle {:.0} ms, policy {:.0} ms, round-robin {:.0} ms",
            oracle / tasks as f64,
            policy / tasks as f64,
            rr / tasks as f64
        );
        assert!(
            policy_regret * 3.0 < rr_regret,
            "policy regret {policy_regret:.0} ms vs round-robin {rr_regret:.0} ms"
        );
    }

    /// Regret must grow more slowly than the number of tasks: the cost of
    /// learning is paid once, not per task, or the policy never converges.
    #[test]
    fn regret_per_task_shrinks_as_it_learns() {
        let oracle_per_task = |tasks: usize| run_oracle(mesh(), tasks, 3) / tasks as f64;
        let policy_per_task = |tasks: usize| run_policy(mesh(), tasks, 3, 2) / tasks as f64;

        let early = policy_per_task(300) - oracle_per_task(300);
        let late = policy_per_task(6_000) - oracle_per_task(6_000);
        println!("  regret per task: early {early:.1} ms, late {late:.1} ms");
        assert!(
            late < early * 0.6,
            "should keep improving: early {early:.1} ms, late {late:.1} ms"
        );
    }

    /// The harder direction, and the risk that comes with weighting sharply:
    /// a worker the policy has learned to avoid is drawn rarely, so it must
    /// still notice when that worker becomes the best one on the mesh. What
    /// saves it is decay — evidence ages, the posterior widens, and a rare
    /// draw is then sampled optimistically enough to be tried again.
    #[test]
    fn it_notices_a_slow_worker_that_becomes_fast() {
        let mut fleet = Fleet::new(mesh());
        let mut rng = Rng::seeded(33);
        let w = Weights::default();

        let mut dispatch = |fleet: &mut Fleet, rng: &mut Rng| -> usize {
            let cands = fleet.candidates();
            let weights: Vec<f64> = cands.iter().map(|c| table_weight(c, &w)).collect();
            let table = AliasTable::build(&weights).expect("a fleet");
            let i = choose(&cands, &table, rng, 2, &w).expect("a choice");
            fleet.run(i, rng);
            i
        };

        // Learn that workers 4.. are the slow ones.
        for _ in 0..1_500 {
            dispatch(&mut fleet, &mut rng);
        }
        let before: usize = (0..300)
            .map(|_| usize::from(dispatch(&mut fleet, &mut rng) == 12))
            .sum();

        // Worker 12 becomes by far the best on the mesh.
        fleet.truth_ms[12] = 20.0;
        for _ in 0..4_000 {
            dispatch(&mut fleet, &mut rng);
        }
        let after: usize = (0..300)
            .map(|_| usize::from(dispatch(&mut fleet, &mut rng) == 12))
            .sum();

        println!("  share on the worker that got fast: {before}/300 → {after}/300");
        assert!(before < 30, "it had learned to avoid it: {before}/300");
        assert!(
            after > 150,
            "it must find a worker that improves, not write it off: {after}/300"
        );
    }

    /// The point of discounting: when the fast workers go slow, the policy has
    /// to notice and move, not keep sending work where it used to be good.
    #[test]
    fn it_follows_the_mesh_when_conditions_change() {
        let mut fleet = Fleet::new(mesh());
        let mut rng = Rng::seeded(21);
        let w = Weights::default();

        let mut dispatch = |fleet: &mut Fleet, rng: &mut Rng| -> usize {
            let cands = fleet.candidates();
            let weights: Vec<f64> = cands.iter().map(|c| table_weight(c, &w)).collect();
            let table = AliasTable::build(&weights).expect("a fleet");
            let i = choose(&cands, &table, rng, 2, &w).expect("a choice");
            fleet.run(i, rng);
            i
        };

        // Learn the fleet as it is.
        for _ in 0..1_500 {
            dispatch(&mut fleet, &mut rng);
        }
        let before: usize = (0..300)
            .map(|_| usize::from(dispatch(&mut fleet, &mut rng) < 4))
            .sum();

        // The local workers become the slow ones.
        for i in 0..4 {
            fleet.truth_ms[i] = 2_000.0;
        }
        for _ in 0..1_500 {
            dispatch(&mut fleet, &mut rng);
        }
        let after: usize = (0..300)
            .map(|_| usize::from(dispatch(&mut fleet, &mut rng) < 4))
            .sum();

        println!("  share on the formerly-fast workers: {before}/300 → {after}/300");
        assert!(before > 200, "it found the fast ones: {before}/300");
        assert!(
            after < 90,
            "it moved off them once they slowed: {after}/300"
        );
    }
}