crates/admit/src/estimate.rs
//! **The estimate learns**: an EWMA per [`WorkKind`], with size as a feature —
//! and with no lock, so a record on one core never waits for one on another.
//!
//! Two EWMAs per kind, because they answer two different questions:
//!
//! * **per hinted byte** — when the caller had a size hint, each field of the
//! actual [`Cost`] divided by the hint. A clone of a 1.2 GB repository and a
//! clone of a 12 MB one are the same *kind* of work at different *sizes*, and
//! the ratio is what carries across them. Used whenever the new work also
//! has a hint.
//! * **absolute** — the plain mean cost of this kind, for work with no hint.
//!
//! Before anything has been learnt, a hinted estimate falls back to a PRIOR
//! that says which directions the hint's bytes move in (a clone reads the disk
//! and writes the network; a push the reverse; a pull-through fetches
//! upstream, writes the disk and serves). With no hint and no history the
//! estimate is zero, and the plan's sentence says "no estimate yet".
//!
//! `α = 0.2`; the first sample of either EWMA is taken as-is.
//!
//! # Multi-core: striped, not shared
//!
//! The table is **striped**: one cache-line-aligned stripe per core, and each
//! thread records into its own stripe (assigned round-robin on the thread's
//! first record and kept). A stripe's fields are `f64` bits in `AtomicU64`s,
//! updated by CAS, so even two threads that happen to share a stripe never
//! block — they retry. `estimate` combines the stripes as a sample-weighted
//! mean of their EWMAs: each stripe is an EWMA over its own subsequence of
//! the kind's samples, and on one thread (every unit test here) there is one
//! stripe, which is exactly the classic EWMA.
use std::cell::Cell;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use crate::{Cost, Work, WorkKind};
/// The smoothing factor.
pub const ALPHA: f64 = 0.2;
const KINDS: usize = WorkKind::ALL.len();
#[derive(Default)]
struct Model {
samples: AtomicU64,
abs: [AtomicU64; 6],
hinted: AtomicU64,
per_byte: [AtomicU64; 6],
}
#[repr(align(64))]
#[derive(Default)]
struct Stripe {
kinds: [Model; KINDS],
}
fn load(a: &AtomicU64) -> f64 {
f64::from_bits(a.load(Ordering::Relaxed))
}
/// `mean += α (sample − mean)`, or `mean = sample` for the first; one CAS loop.
fn ewma(mean: &AtomicU64, sample: f64, first: bool) {
let mut cur = mean.load(Ordering::Relaxed);
loop {
let m = f64::from_bits(cur);
let next = if first { sample } else { m + ALPHA * (sample - m) };
match mean.compare_exchange_weak(cur, next.to_bits(), Ordering::Relaxed, Ordering::Relaxed) {
Ok(_) => return,
Err(now) => cur = now,
}
}
}
fn round(f: [f64; 6]) -> Cost {
Cost::from_fields(f.map(|v| if v <= 0.0 || v.is_nan() { 0 } else { v.round() as u64 }))
}
/// What one kind has learnt, for the admin view.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct EstimateRow {
pub kind: WorkKind,
/// Samples recorded for this kind.
pub samples: u64,
/// Of those, how many carried a size hint.
pub hinted: u64,
/// The absolute mean cost (for unhinted work).
pub mean: Cost,
/// The cost of one GB (10^9 bytes) of hinted size, or `None` before any
/// hinted sample.
pub per_gb: Option<Cost>,
}
/// The per-kind EWMA table, striped per core. See the module docs.
pub struct Estimator {
stripes: Box<[Stripe]>,
next: AtomicUsize,
}
impl std::fmt::Debug for Estimator {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Estimator").field("stripes", &self.stripes.len()).finish()
}
}
impl Default for Estimator {
fn default() -> Estimator {
Estimator::with_stripes(crate::table::default_shards())
}
}
thread_local! {
/// This thread's stripe ticket, drawn once.
static TICKET: Cell<Option<usize>> = const { Cell::new(None) };
}
impl Estimator {
pub fn new() -> Estimator {
Estimator::default()
}
/// An estimator with exactly `n` stripes (at least one).
pub fn with_stripes(n: usize) -> Estimator {
Estimator {
stripes: (0..n.max(1)).map(|_| Stripe::default()).collect(),
next: AtomicUsize::new(0),
}
}
fn mine(&self) -> &Stripe {
let ticket = TICKET.with(|t| match t.get() {
Some(i) => i,
None => {
let i = self.next.fetch_add(1, Ordering::Relaxed);
t.set(Some(i));
i
}
});
&self.stripes[ticket % self.stripes.len()]
}
/// The sample-weighted mean over stripes of one kind's `(abs, per_byte)`.
fn combined(&self, kind: WorkKind) -> (u64, [f64; 6], u64, [f64; 6]) {
let (mut n, mut abs, mut nh, mut per) = (0u64, [0f64; 6], 0u64, [0f64; 6]);
for s in self.stripes.iter() {
let m = &s.kinds[kind.index()];
let k = m.samples.load(Ordering::Relaxed);
if k > 0 {
n += k;
for (acc, a) in abs.iter_mut().zip(&m.abs) {
*acc += load(a) * k as f64;
}
}
let kh = m.hinted.load(Ordering::Relaxed);
if kh > 0 {
nh += kh;
for (acc, a) in per.iter_mut().zip(&m.per_byte) {
*acc += load(a) * kh as f64;
}
}
}
if n > 0 {
abs = abs.map(|v| v / n as f64);
}
if nh > 0 {
per = per.map(|v| v / nh as f64);
}
(n, abs, nh, per)
}
/// **What `work` is expected to cost.** Loads only.
pub fn estimate(&self, work: &Work) -> Cost {
let (n, abs, nh, per) = self.combined(work.kind);
match work.size_hint.filter(|h| *h > 0) {
Some(h) if nh > 0 => round(per.map(|r| r * h as f64)),
Some(h) => prior(work.kind, h),
None if n > 0 => round(abs),
None => Cost::default(),
}
}
/// **Learn from what `work` actually cost.** CAS on this thread's stripe.
pub fn learn(&self, work: &Work, actual: &Cost) {
let f = actual.fields().map(|v| v as f64);
let m = &self.mine().kinds[work.kind.index()];
let first = m.samples.load(Ordering::Relaxed) == 0;
for (mean, v) in m.abs.iter().zip(f) {
ewma(mean, v, first);
}
m.samples.fetch_add(1, Ordering::Relaxed);
if let Some(h) = work.size_hint.filter(|h| *h > 0) {
let first = m.hinted.load(Ordering::Relaxed) == 0;
for (mean, v) in m.per_byte.iter().zip(f) {
ewma(mean, v / h as f64, first);
}
m.hinted.fetch_add(1, Ordering::Relaxed);
}
}
/// One row per kind that has learnt anything, in [`WorkKind::ALL`] order.
pub fn rows(&self) -> Vec<EstimateRow> {
WorkKind::ALL
.iter()
.filter_map(|kind| {
let (n, abs, nh, per) = self.combined(*kind);
(n > 0).then(|| EstimateRow {
kind: *kind,
samples: n,
hinted: nh,
mean: round(abs),
per_gb: (nh > 0).then(|| round(per.map(|r| r * 1e9))),
})
})
.collect()
}
}
/// **Which directions a hint's bytes move in**, before anything is learnt.
pub fn prior(kind: WorkKind, h: u64) -> Cost {
let z = Cost::default();
match kind {
WorkKind::Clone | WorkKind::Fetch | WorkKind::ArtifactDownload | WorkKind::Other => {
Cost { disk_read: h, net_out: h, ..z }
}
WorkKind::Push => Cost { net_in: h, disk_write: h, ..z },
WorkKind::PullThrough => Cost { upstream_fetch: h, disk_write: h, net_out: h, ..z },
WorkKind::MirrorAsset => Cost { upstream_fetch: h, disk_write: h, ..z },
WorkKind::Promote => Cost { disk_read: h, disk_write: h, ..z },
}
}
#[cfg(test)]
mod tests {
use super::*;
fn clone(h: Option<u64>) -> Work {
Work::new(WorkKind::Clone, h)
}
#[test]
fn nothing_learnt_is_the_prior_with_a_hint_and_zero_without() {
let e = Estimator::new();
assert_eq!(
e.estimate(&clone(Some(1000))),
Cost { disk_read: 1000, net_out: 1000, ..Cost::default() }
);
assert_eq!(e.estimate(&clone(None)), Cost::default());
assert!(e.rows().is_empty());
}
#[test]
fn the_first_sample_is_taken_as_is_then_it_moves_by_alpha() {
let e = Estimator::new();
e.learn(&clone(None), &Cost { net_out: 1000, ..Cost::default() });
assert_eq!(e.estimate(&clone(None)).net_out, 1000);
e.learn(&clone(None), &Cost { net_out: 2000, ..Cost::default() });
// 1000 + 0.2 * (2000 - 1000)
assert_eq!(e.estimate(&clone(None)).net_out, 1200);
e.learn(&clone(None), &Cost { net_out: 2000, ..Cost::default() });
// 1200 + 0.2 * 800
assert_eq!(e.estimate(&clone(None)).net_out, 1360);
}
/// ★ Size is a feature: what is learnt from a 1 GB clone scales to a 10 GB one.
#[test]
fn a_hinted_estimate_scales_with_the_hint() {
let e = Estimator::new();
e.learn(
&clone(Some(1_000_000_000)),
&Cost { net_out: 1_500_000_000, cpu_ms: 2000, ..Cost::default() },
);
let big = e.estimate(&clone(Some(10_000_000_000)));
assert_eq!(big.net_out, 15_000_000_000);
assert_eq!(big.cpu_ms, 20_000);
assert_eq!(big.disk_read, 0, "learnt, not prior: this clone read nothing from disk");
let row = e.rows()[0];
assert_eq!((row.kind, row.samples, row.hinted), (WorkKind::Clone, 1, 1));
assert_eq!(row.per_gb.unwrap().net_out, 1_500_000_000);
}
#[test]
fn kinds_learn_separately() {
let e = Estimator::new();
e.learn(&Work::new(WorkKind::Push, None), &Cost { net_in: 7, ..Cost::default() });
assert_eq!(e.estimate(&clone(None)), Cost::default());
assert_eq!(e.estimate(&Work::new(WorkKind::Push, None)).net_in, 7);
}
#[test]
fn an_estimate_converges_on_a_steady_cost() {
let e = Estimator::new();
e.learn(&clone(Some(100)), &Cost { net_out: 1000, ..Cost::default() });
for _ in 0..60 {
e.learn(&clone(Some(100)), &Cost { net_out: 200, ..Cost::default() });
}
assert_eq!(e.estimate(&clone(Some(100))).net_out, 200);
}
/// ★ Many threads learning at once: every sample is counted, and a
/// steady cost is what the combined estimate converges on.
#[test]
fn concurrent_learning_counts_every_sample() {
let e = std::sync::Arc::new(Estimator::with_stripes(4));
let threads: Vec<_> = (0..8)
.map(|_| {
let e = std::sync::Arc::clone(&e);
std::thread::spawn(move || {
for _ in 0..1000 {
e.learn(&clone(None), &Cost { net_out: 500, ..Cost::default() });
}
})
})
.collect();
for t in threads {
t.join().unwrap();
}
let row = e.rows()[0];
assert_eq!(row.samples, 8000);
assert_eq!(row.mean.net_out, 500);
}
}