crates/admit/src/share.rs
//! [`UsageShare`] and its [`Settings`].
use std::fmt;
use std::hash::{Hash, Hasher};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use crate::estimate::Estimator;
use crate::table::{default_shards, Table};
use crate::{
sentence, Admit, Class, Consumer, Cost, DayToken, EstimateRow, Guard, PassThrough, Plan,
Refusal, Usage, Work, DAY_SECS, DEFAULT_ANONYMOUS_PER_DAY, DEFAULT_LIMIT_BYTES,
DEFAULT_WINDOW_DAYS,
};
/// Consumers a table holds by default: 16 384 named, 16 384 anonymous tokens.
/// About 4 MB preallocated for a 30-day window.
pub const DEFAULT_MAX_CONSUMERS: usize = 16_384;
/// **Concurrent-slot shares per class**, for a product that has no gate of its
/// own. `Reduced` and `Floor` draw on ONE pool — "over budget" is the
/// anonymous share, by design, not a third tier. `Queued` takes no slot.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize))]
#[cfg_attr(feature = "serde", serde(deny_unknown_fields))]
pub struct Slots {
/// Most `Full` work in flight at once.
pub full: usize,
/// Most `Reduced` + `Floor` work in flight at once.
pub anonymous: usize,
}
/// **The `[admit]` table of an estate config.**
///
/// `limit = "none"` is PassThrough ([`Settings::is_pass_through`]); sizes are
/// SI (`"100 GB"` = 10^11 bytes) or IEC (`"5 GiB"`), or a bare integer of bytes.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Settings {
/// Bytes per window a named consumer moves at `Full`; `None` = unlimited.
pub limit: Option<u64>,
/// The rolling window, in UTC days.
pub window_days: u32,
/// Bytes one anonymous day token may move per UTC day; `None` = unlimited.
pub anonymous_per_day: Option<u64>,
/// Class shares, when the product has no gate of its own.
pub slots: Option<Slots>,
/// Consumers each table preallocates room for.
pub max_consumers: usize,
}
impl Default for Settings {
fn default() -> Settings {
Settings {
limit: Some(DEFAULT_LIMIT_BYTES),
window_days: DEFAULT_WINDOW_DAYS,
anonymous_per_day: Some(DEFAULT_ANONYMOUS_PER_DAY),
slots: None,
max_consumers: DEFAULT_MAX_CONSUMERS,
}
}
}
impl Settings {
/// `limit = "none"`: the product runs [`PassThrough`].
pub fn pass_through() -> Settings {
Settings { limit: None, ..Settings::default() }
}
pub fn is_pass_through(&self) -> bool {
self.limit.is_none()
}
/// **The policy these settings select**: PassThrough for `limit = none`,
/// otherwise a UsageShare.
pub fn build(self) -> Arc<dyn Admit> {
self.build_seeded(None)
}
/// [`Settings::build`], with a durable ledger to seed each consumer from on
/// first sight (see [`Seed`]).
pub fn build_seeded(self, seed: Option<Arc<dyn Seed>>) -> Arc<dyn Admit> {
if self.is_pass_through() {
Arc::new(PassThrough::new())
} else {
let share = UsageShare::new(self);
Arc::new(match seed {
Some(seed) => share.with_seed(seed),
None => share,
})
}
}
}
/// A setting that did not parse.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SettingsError(pub String);
impl fmt::Display for SettingsError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
impl std::error::Error for SettingsError {}
/// **`"none"` → `None`, `"100 GB"` → `Some(10^11)`**, `"5 GiB"`, `"123"`.
pub fn parse_bytes(s: &str) -> Result<Option<u64>, SettingsError> {
let t = s.trim();
if t.eq_ignore_ascii_case("none") || t.eq_ignore_ascii_case("unlimited") {
return Ok(None);
}
let split = t.find(|c: char| !(c.is_ascii_digit() || c == '.')).unwrap_or(t.len());
let (num, unit) = (t[..split].trim(), t[split..].trim());
let n: f64 = num
.parse()
.map_err(|_| SettingsError(format!("`{s}` is not a size (e.g. \"100 GB\" or \"none\")")))?;
let mul: f64 = match unit.to_ascii_lowercase().as_str() {
"" | "b" => 1.0,
"kb" => 1e3,
"mb" => 1e6,
"gb" => 1e9,
"tb" => 1e12,
"kib" => 1024.0,
"mib" => 1024.0 * 1024.0,
"gib" => 1024.0 * 1024.0 * 1024.0,
"tib" => 1024.0 * 1024.0 * 1024.0 * 1024.0,
other => return Err(SettingsError(format!("`{other}` is not a size unit in `{s}`"))),
};
Ok(Some((n * mul).round() as u64))
}
#[cfg(feature = "serde")]
mod de {
use super::*;
#[derive(serde::Deserialize)]
#[serde(untagged)]
enum Size {
Bytes(u64),
Text(String),
}
fn size<E: serde::de::Error>(v: Option<Size>, default: Option<u64>) -> Result<Option<u64>, E> {
match v {
None => Ok(default),
Some(Size::Bytes(n)) => Ok(Some(n)),
Some(Size::Text(t)) => parse_bytes(&t).map_err(E::custom),
}
}
#[derive(serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct Raw {
limit: Option<Size>,
window_days: Option<u32>,
anonymous_per_day: Option<Size>,
slots: Option<Slots>,
max_consumers: Option<usize>,
}
impl<'de> serde::Deserialize<'de> for Settings {
fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Settings, D::Error> {
let raw = Raw::deserialize(d)?;
let def = Settings::default();
Ok(Settings {
limit: size(raw.limit, def.limit)?,
window_days: raw.window_days.unwrap_or(def.window_days).max(1),
anonymous_per_day: size(raw.anonymous_per_day, def.anonymous_per_day)?,
slots: raw.slots,
max_consumers: raw.max_consumers.unwrap_or(def.max_consumers),
})
}
}
}
/// **A durable ledger to seed a consumer from on first sight.** gunnar's is
/// its flow ledger (bytes per namespace per UTC day, kept 400 days): the
/// table here is then a CACHE of that ledger, filled once per consumer per
/// process — off the hot path, on the insert that first sees it — and kept
/// current by `record`. Returns `(day, bytes)` for the days it has, from
/// `first_day` to `today` inclusive.
pub trait Seed: Send + Sync {
fn seed(&self, who: &Consumer, first_day: u64, today: u64) -> Vec<(u64, u64)>;
}
impl<F> Seed for F
where
F: Fn(&Consumer, u64, u64) -> Vec<(u64, u64)> + Send + Sync,
{
fn seed(&self, who: &Consumer, first_day: u64, today: u64) -> Vec<(u64, u64)> {
self(who, first_day, today)
}
}
/// Class slot counters, one cache line each.
#[repr(align(64))]
#[derive(Debug, Default)]
struct Counter(AtomicUsize);
#[derive(Debug, Default)]
struct SlotTable {
/// `Full`'s cell and the shared `Reduced`/`Floor` cell.
in_flight: [Counter; 2],
}
/// A slot held by a [`Guard`]; released on drop.
#[derive(Debug)]
pub(crate) struct SlotHold {
table: Arc<SlotTable>,
cell: usize,
}
impl Drop for SlotHold {
fn drop(&mut self) {
self.table.in_flight[self.cell].0.fetch_sub(1, Ordering::AcqRel);
}
}
/// The table identity of a named consumer: a keyed hash of its kind and
/// name, with no string built for it.
fn named_key(table: &Table, who: &Consumer) -> Option<u64> {
struct Named<'a>(&'a Consumer);
impl Hash for Named<'_> {
fn hash<H: Hasher>(&self, h: &mut H) {
match self.0 {
Consumer::Principal(p) => (0u8, p.as_str()).hash(h),
Consumer::Namespace(n) => (1u8, n.as_str()).hash(h),
Consumer::Peer(p) => (2u8, p.as_str()).hash(h),
Consumer::Anonymous(_) => unreachable!("anonymous has no named key"),
}
}
}
(!who.is_anonymous()).then(|| table.key(&Named(who)))
}
/// A day token's identity: its first 8 bytes are already a keyed hash.
fn token_key(t: &DayToken) -> u64 {
Table::fix(u64::from_le_bytes(t.token[..8].try_into().expect("8 bytes")))
}
/// **Rolling-window bytes per consumer decide the class.**
///
/// * named consumer, window usage below `limit` → [`Class::Full`];
/// * at or above it → [`Class::Reduced`] (the anonymous share);
/// * anonymous → [`Class::Floor`], refused by [`Admit::admit`] once its day
/// token has moved `anonymous_per_day`.
///
/// **Multi-core, no global lock** (see `table` and `estimate`): usage is a
/// sharded preallocated table of atomic day buckets, the estimate a striped
/// set of atomic EWMAs, and a class slot one `fetch_add`. A new consumer takes
/// its SHARD's insert lock once; nothing else locks.
pub struct UsageShare {
settings: Settings,
estimator: Estimator,
named: Table,
anonymous: Table,
slots: Option<Arc<SlotTable>>,
seed: Option<Arc<dyn Seed>>,
clock: fn() -> u64,
}
impl fmt::Debug for UsageShare {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("UsageShare").field("settings", &self.settings).finish_non_exhaustive()
}
}
impl UsageShare {
pub fn new(settings: Settings) -> UsageShare {
let shards = default_shards();
UsageShare {
settings,
estimator: Estimator::with_stripes(shards),
named: Table::new(settings.max_consumers, settings.window_days, shards),
anonymous: Table::new(settings.max_consumers, 1, shards),
slots: settings.slots.map(|_| Arc::new(SlotTable::default())),
seed: None,
clock: crate::visitors::now_unix,
}
}
/// The same share on another clock (Unix seconds) — a test seam.
pub fn with_clock(mut self, clock: fn() -> u64) -> UsageShare {
self.clock = clock;
self
}
/// Seed each named consumer from a durable ledger on first sight.
pub fn with_seed(mut self, seed: Arc<dyn Seed>) -> UsageShare {
self.seed = Some(seed);
self
}
pub fn settings(&self) -> &Settings {
&self.settings
}
fn today(&self) -> u64 {
(self.clock)() / DAY_SECS
}
/// The named table's slot for `who`, inserting (and seeding) it on first sight.
fn held(&self, h: u64, who: &Consumer, today: u64) -> bool {
if self.named.contains(h) {
return true;
}
let first = today.saturating_sub(u64::from(self.settings.window_days) - 1);
let seed = self.seed.as_ref().map(|s| s.seed(who, first, today)).unwrap_or_default();
self.named.insert(h, today, &seed)
}
fn anonymous_used(&self, t: &DayToken) -> u64 {
self.anonymous.on(token_key(t), t.day).unwrap_or(0)
}
fn cell(class: Class) -> Option<usize> {
match class {
Class::Full => Some(0),
Class::Reduced | Class::Floor => Some(1),
Class::Queued => None,
}
}
fn share(&self, class: Class) -> usize {
match (self.settings.slots, class) {
(Some(s), Class::Full) => s.full,
(Some(s), _) => s.anonymous,
(None, _) => usize::MAX,
}
}
}
impl Admit for UsageShare {
fn plan(&self, work: &Work, who: &Consumer, usage: &Usage) -> Plan {
let estimate = self.estimator.estimate(work);
match who {
Consumer::Anonymous(t) => {
let used = Usage::new(self.anonymous_used(t), 1);
let limit = self.settings.anonymous_per_day.map(|l| Usage::new(l, 1));
Plan {
class: Class::Floor,
estimate,
why: sentence(limit, &used, work, &estimate, Class::Floor),
budget_left: limit
.map(|l| Usage::new(l.bytes.saturating_sub(used.bytes), 1))
.unwrap_or(Usage::unlimited(1)),
work: *work,
who: who.clone(),
}
}
named => {
let window = usage.window_days;
let (class, left, limit) = match self.settings.limit {
Some(l) => (
if usage.bytes >= l { Class::Reduced } else { Class::Full },
Usage::new(l.saturating_sub(usage.bytes), window),
Some(Usage::new(l, self.settings.window_days)),
),
None => (Class::Full, Usage::unlimited(window), None),
};
Plan {
class,
estimate,
why: sentence(limit, usage, work, &estimate, class),
budget_left: left,
work: *work,
who: named.clone(),
}
}
}
}
fn admit(&self, plan: &Plan) -> Result<Guard, Refusal> {
if let (Consumer::Anonymous(t), Some(limit)) = (&plan.who, self.settings.anonymous_per_day) {
let used = self.anonymous_used(t);
if used >= limit {
return Err(Refusal::DayTokenSpent { used, limit });
}
}
let slot = match (&self.slots, Self::cell(plan.class)) {
(Some(table), Some(cell)) => {
let share = self.share(plan.class);
// fetch_add, compare, and give it back on a refusal — the
// gunnar-ssh admission pattern: one atomic op on the way in.
let was = table.in_flight[cell].0.fetch_add(1, Ordering::AcqRel);
if was >= share {
table.in_flight[cell].0.fetch_sub(1, Ordering::AcqRel);
return Err(Refusal::ClassFull { class: plan.class, in_flight: was, share });
}
Some(SlotHold { table: Arc::clone(table), cell })
}
_ => None,
};
Ok(Guard { class: plan.class, work: plan.work, who: plan.who.clone(), slot })
}
fn record(&self, guard: Guard, actual: &Cost) {
self.estimator.learn(&guard.work, actual);
let bytes = actual.usage_bytes();
match &guard.who {
Consumer::Anonymous(t) => {
let h = token_key(t);
// Yesterday's token charges yesterday's bucket, which today's
// reads never look at — and the table drops a day older than
// the one its bucket holds.
if !self.anonymous.add(h, t.day, bytes) && self.anonymous.insert(h, t.day, &[]) {
self.anonymous.add(h, t.day, bytes);
}
}
named => {
if let Some(h) = named_key(&self.named, named) {
let today = self.today();
if !self.named.add(h, today, bytes) && self.held(h, named, today) {
self.named.add(h, today, bytes);
}
}
}
}
drop(guard);
}
fn usage(&self, who: &Consumer) -> Usage {
match who {
Consumer::Anonymous(t) => Usage::new(self.anonymous_used(t), 1),
named => {
let today = self.today();
let window = self.settings.window_days;
let Some(h) = named_key(&self.named, named) else {
return Usage::new(0, window);
};
let bytes = match self.named.sum(h, today) {
Some(b) => b,
None if self.seed.is_some() && self.held(h, named, today) => {
self.named.sum(h, today).unwrap_or(0)
}
None => 0,
};
Usage::new(bytes, window)
}
}
}
fn estimates(&self) -> Vec<EstimateRow> {
self.estimator.rows()
}
fn policy(&self) -> &'static str {
"usage-share"
}
fn limit(&self) -> Option<Usage> {
self.settings.limit.map(|l| Usage::new(l, self.settings.window_days))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::daykey::DayKey;
use crate::visitors::Client;
use crate::WorkKind;
use std::sync::atomic::AtomicU64;
const DAY0: u64 = 20_000 * DAY_SECS;
static NOW: AtomicU64 = AtomicU64::new(DAY0 + 10);
fn clock() -> u64 {
NOW.load(Ordering::Relaxed)
}
fn ns(n: &str) -> Consumer {
Consumer::Namespace(n.into())
}
fn clone(h: Option<u64>) -> Work {
Work::new(WorkKind::Clone, h)
}
const GB: u64 = 1_000_000_000;
#[test]
fn below_the_limit_is_full_and_at_or_above_is_reduced() {
let s = UsageShare::new(Settings::default());
let w = clone(Some(1_200_000_000));
let p = s.plan(&w, &ns("acme"), &Usage::new(63 * GB, 30));
assert_eq!(p.class, Class::Full);
assert_eq!(p.why, "100 GB/30d used 63 GB; this clone ~1.2 GB; Full");
assert_eq!(p.budget_left, Usage::new(37 * GB, 30));
let p = s.plan(&w, &ns("acme"), &Usage::new(100 * GB, 30));
assert_eq!(p.class, Class::Reduced);
assert!(p.why.ends_with("; Reduced"), "{}", p.why);
assert_eq!(p.budget_left.bytes, 0);
}
#[test]
fn plan_is_pure() {
let s = UsageShare::new(Settings { slots: Some(Slots { full: 1, anonymous: 1 }), ..Settings::default() });
for _ in 0..10 {
let _ = s.plan(&clone(None), &ns("a"), &Usage::none());
}
assert!(s.admit(&s.plan(&clone(None), &ns("a"), &Usage::none())).is_ok());
assert!(s.named.is_empty(), "a plan inserted a consumer");
}
/// ★ The classes are shares of slots, refused fast with a reason; Reduced
/// and Floor draw on ONE pool; a released guard frees its slot.
#[test]
fn classes_are_shares_and_a_full_share_refuses_with_a_reason() {
let s = UsageShare::new(Settings { slots: Some(Slots { full: 2, anonymous: 1 }), ..Settings::default() });
let full = s.plan(&clone(None), &ns("a"), &Usage::none());
let over = s.plan(&clone(None), &ns("b"), &Usage::new(200 * GB, 30));
assert_eq!(over.class, Class::Reduced);
let anon = Consumer::Anonymous(DayKey::new(DAY0).token(Client::Ip("203.0.113.7".parse().unwrap()), DAY0));
let floor = s.plan(&clone(None), &anon, &Usage::none());
assert_eq!(floor.class, Class::Floor);
let g1 = s.admit(&full).unwrap();
let _g2 = s.admit(&full).unwrap();
assert_eq!(
s.admit(&full).unwrap_err(),
Refusal::ClassFull { class: Class::Full, in_flight: 2, share: 2 }
);
let r = s.admit(&over).unwrap();
assert!(matches!(s.admit(&floor), Err(Refusal::ClassFull { class: Class::Floor, in_flight: 1, share: 1 })));
drop(r);
assert!(s.admit(&floor).is_ok(), "the Reduced slot came back and Floor may take it");
s.record(g1, &Cost::default());
assert!(s.admit(&full).is_ok(), "record releases the slot");
}
#[test]
fn a_refusal_gives_its_slot_back() {
let s = UsageShare::new(Settings { slots: Some(Slots { full: 1, anonymous: 1 }), ..Settings::default() });
let p = s.plan(&clone(None), &ns("a"), &Usage::none());
let held = s.admit(&p).unwrap();
for _ in 0..100 {
assert!(s.admit(&p).is_err());
}
drop(held);
assert!(s.admit(&p).is_ok(), "a hundred refusals leaked a slot");
}
#[test]
fn queued_takes_no_slot() {
let s = UsageShare::new(Settings { slots: Some(Slots { full: 0, anonymous: 0 }), ..Settings::default() });
let mut p = s.plan(&clone(None), &ns("a"), &Usage::none());
p.class = Class::Queued;
assert!(s.admit(&p).is_ok());
}
/// ★ An anonymous token has a per-day budget; a new day is a new token and
/// a fresh budget, and yesterday's token cannot be charged into today.
#[test]
fn an_anonymous_day_token_has_a_daily_budget() {
let s = UsageShare::new(Settings::default());
let key = DayKey::new(DAY0);
let c = Client::Ip("203.0.113.7".parse().unwrap());
let who = Consumer::Anonymous(key.token(c, DAY0 + 10));
let w = Work::new(WorkKind::ArtifactDownload, None);
let g = s.admit(&s.plan(&w, &who, &Usage::none())).unwrap();
s.record(g, &Cost { net_out: 4 * GB, ..Cost::default() });
let p = s.plan(&w, &who, &Usage::none());
assert_eq!(p.budget_left, Usage::new(GB, 1));
assert!(p.why.starts_with("5 GB/day used 4 GB; this download ~4 GB; Floor"), "{}", p.why);
let g = s.admit(&p).unwrap();
s.record(g, &Cost { net_out: 2 * GB, ..Cost::default() });
let spent = s.admit(&s.plan(&w, &who, &Usage::none())).unwrap_err();
assert_eq!(spent, Refusal::DayTokenSpent { used: 6 * GB, limit: 5 * GB });
let other = Consumer::Anonymous(key.token(Client::Ip("198.51.100.9".parse().unwrap()), DAY0 + 20));
assert!(s.admit(&s.plan(&w, &other, &Usage::none())).is_ok());
// Midnight: the same address is a new token with a whole budget…
let tomorrow = Consumer::Anonymous(key.token(c, DAY0 + DAY_SECS + 1));
assert_ne!(tomorrow, who);
let g = s.admit(&s.plan(&w, &tomorrow, &Usage::none())).unwrap();
s.record(g, &Cost { net_out: 1, ..Cost::default() });
assert_eq!(s.usage(&tomorrow).bytes, 1);
// …and a guard still holding yesterday's token charges nothing to today.
let late = Guard { class: Class::Floor, work: w, who: who.clone(), slot: None };
s.record(late, &Cost { net_out: 100 * GB, ..Cost::default() });
assert_eq!(s.usage(&tomorrow).bytes, 1);
}
/// ★ Without a ledger of its own a caller reads `usage`, fed by `record`
/// over the rolling window; the estimate learns from the same record.
#[test]
fn record_feeds_the_rolling_usage_and_the_estimate() {
NOW.store(DAY0 + 10, Ordering::Relaxed);
let s = UsageShare::new(Settings { limit: Some(10 * GB), ..Settings::default() }).with_clock(clock);
let who = ns("acme");
let w = clone(Some(GB));
let g = s.admit(&s.plan(&w, &who, &s.usage(&who))).unwrap();
s.record(g, &Cost { net_out: 11 * GB, ..Cost::default() });
assert_eq!(s.usage(&who), Usage::new(11 * GB, 30));
let p = s.plan(&w, &who, &s.usage(&who));
assert_eq!(p.class, Class::Reduced);
assert_eq!(p.estimate.net_out, 11 * GB, "the estimate learnt from the record");
assert_eq!(s.estimates()[0].samples, 1);
NOW.store(DAY0 + 30 * DAY_SECS + 10, Ordering::Relaxed);
assert_eq!(s.usage(&who).bytes, 0, "a window later the day has left");
}
/// ★ A durable ledger seeds a consumer ONCE, on first sight; after that the
/// table is the cache and `record` keeps it current.
#[test]
fn a_seed_is_read_once_per_consumer() {
static READS: AtomicU64 = AtomicU64::new(0);
NOW.store(DAY0 + 10, Ordering::Relaxed);
let seed = |who: &Consumer, first: u64, today: u64| {
READS.fetch_add(1, Ordering::Relaxed);
assert_eq!(today - first, 29);
match who {
Consumer::Namespace(n) if n == "acme" => vec![(today - 3, 70 * GB), (today, 50 * GB)],
_ => vec![],
}
};
let s = UsageShare::new(Settings::default()).with_clock(clock).with_seed(Arc::new(seed));
assert_eq!(s.usage(&ns("acme")).bytes, 120 * GB);
assert_eq!(s.usage(&ns("acme")).bytes, 120 * GB);
assert_eq!(READS.load(Ordering::Relaxed), 1, "the durable ledger is read once");
let g = s.admit(&s.plan(&clone(None), &ns("acme"), &s.usage(&ns("acme")))).unwrap();
s.record(g, &Cost { net_in: 5, ..Cost::default() });
assert_eq!(s.usage(&ns("acme")).bytes, 120 * GB + 5);
assert_eq!(READS.load(Ordering::Relaxed), 1);
}
#[test]
fn limit_none_builds_pass_through() {
assert_eq!(Settings::pass_through().build().policy(), "pass-through");
assert_eq!(Settings::default().build().policy(), "usage-share");
}
#[test]
fn sizes_parse() {
assert_eq!(parse_bytes("none"), Ok(None));
assert_eq!(parse_bytes("100 GB"), Ok(Some(100 * GB)));
assert_eq!(parse_bytes("5GB"), Ok(Some(5 * GB)));
assert_eq!(parse_bytes("1.5 GiB"), Ok(Some(1_610_612_736)));
assert_eq!(parse_bytes("123"), Ok(Some(123)));
assert!(parse_bytes("5 parsecs").is_err());
}
}