peerspeak/src/core/fetchbudget.rs

//! Byte/request budgets for AUTOMATIC chat-attachment fetches (Phase 3B).
//!
//! The four-permit semaphore bounds how many auto-fetch tasks run at once, but
//! not how much a peer can make us download over time: with permits released
//! after each transfer, an insider could stream distinct ≤4 MiB images
//! sequentially forever. This budget adds per-author and session (room-wide)
//! token buckets over both request COUNT and declared BYTES. Like the Phase 2
//! chat gate, time is passed in — never read from a clock — so every refill
//! boundary is unit-testable.
//!
//! Only the automatic path consults this; a user's explicit click (Save /
//! Download / Load image) is human-rate-limited and always allowed through to
//! the fetch (still subject to the transfer cap and cache/decoder budgets).

use iroh::EndpointId;
use std::collections::HashMap;

/// Per-author request burst: how many auto-fetches one author can trigger
/// back-to-back before refill pacing binds.
pub const AUTHOR_REQ_BURST: f64 = 8.0;
/// Per-author request refill: one recovered every 10 s.
pub const AUTHOR_REQ_REFILL_PER_MS: f64 = 1.0 / 10_000.0;
/// Per-author byte burst (declared sizes): a couple of full-size auto images
/// plus a normal working set.
pub const AUTHOR_BYTES_BURST: f64 = (16 * 1024 * 1024) as f64;
/// Per-author byte refill: 64 KiB/s (~one 4 MiB auto image per minute).
pub const AUTHOR_BYTES_REFILL_PER_MS: f64 = (64 * 1024) as f64 / 1000.0;

/// Session-wide request burst across all authors.
pub const SESSION_REQ_BURST: f64 = 16.0;
/// Session-wide request refill: one recovered every 5 s.
pub const SESSION_REQ_REFILL_PER_MS: f64 = 1.0 / 5_000.0;
/// Session-wide byte burst across all authors.
pub const SESSION_BYTES_BURST: f64 = (48 * 1024 * 1024) as f64;
/// Session-wide byte refill: 128 KiB/s.
pub const SESSION_BYTES_REFILL_PER_MS: f64 = (128 * 1024) as f64 / 1000.0;

/// Bound on the per-author bucket map. Authors are roster members (≤32 live),
/// so this tracks the roster plus recently departed; the least-recently-active
/// entry is pruned past the cap.
pub const AUTHOR_MAP_CAP: usize = 64;

/// A deterministic token bucket that can take a WEIGHTED cost (bytes), unlike
/// the unit-cost bucket in the gossip chat gate.
#[derive(Debug, Clone, Copy)]
struct WeightedBucket {
    tokens: f64,
    last_ms: u64,
}

impl WeightedBucket {
    fn full(burst: f64, now_ms: u64) -> Self {
        Self {
            tokens: burst,
            last_ms: now_ms,
        }
    }

    /// Refill for elapsed time (capped at `burst`) without consuming.
    fn refill(&mut self, burst: f64, refill_per_ms: f64, now_ms: u64) {
        let elapsed = now_ms.saturating_sub(self.last_ms) as f64;
        self.tokens = (self.tokens + elapsed * refill_per_ms).min(burst);
        self.last_ms = now_ms;
    }

    fn has(&self, cost: f64) -> bool {
        self.tokens >= cost
    }

    fn take(&mut self, cost: f64) {
        self.tokens -= cost;
    }
}

/// One author's pair of buckets plus last activity (for idle pruning).
#[derive(Debug)]
struct AuthorBudget {
    reqs: WeightedBucket,
    bytes: WeightedBucket,
    last_seen_ms: u64,
}

/// Admission budget for automatic attachment fetches. All four buckets are
/// checked BEFORE any is consumed, so a rejection never burns tokens (no
/// refund bookkeeping — the check-then-take is atomic within `admit`).
#[derive(Debug)]
pub struct AutoFetchBudget {
    session_reqs: WeightedBucket,
    session_bytes: WeightedBucket,
    authors: HashMap<EndpointId, AuthorBudget>,
}

impl AutoFetchBudget {
    pub fn new(now_ms: u64) -> Self {
        Self {
            session_reqs: WeightedBucket::full(SESSION_REQ_BURST, now_ms),
            session_bytes: WeightedBucket::full(SESSION_BYTES_BURST, now_ms),
            authors: HashMap::new(),
        }
    }

    /// Whether an auto-fetch of `size` declared bytes for `author` may start
    /// now. Consumes one request token and `size` byte tokens from BOTH the
    /// author's and the session's buckets — or nothing at all on rejection.
    pub fn admit(&mut self, author: EndpointId, size: u64, now_ms: u64) -> bool {
        self.prune(author, now_ms);
        let entry = self.authors.entry(author).or_insert_with(|| AuthorBudget {
            reqs: WeightedBucket::full(AUTHOR_REQ_BURST, now_ms),
            bytes: WeightedBucket::full(AUTHOR_BYTES_BURST, now_ms),
            last_seen_ms: now_ms,
        });
        entry.last_seen_ms = now_ms;
        entry
            .reqs
            .refill(AUTHOR_REQ_BURST, AUTHOR_REQ_REFILL_PER_MS, now_ms);
        entry
            .bytes
            .refill(AUTHOR_BYTES_BURST, AUTHOR_BYTES_REFILL_PER_MS, now_ms);
        self.session_reqs
            .refill(SESSION_REQ_BURST, SESSION_REQ_REFILL_PER_MS, now_ms);
        self.session_bytes
            .refill(SESSION_BYTES_BURST, SESSION_BYTES_REFILL_PER_MS, now_ms);

        let cost = size as f64;
        let ok = entry.reqs.has(1.0)
            && entry.bytes.has(cost)
            && self.session_reqs.has(1.0)
            && self.session_bytes.has(cost);
        if ok {
            let entry = self.authors.get_mut(&author).expect("just inserted");
            entry.reqs.take(1.0);
            entry.bytes.take(cost);
            self.session_reqs.take(1.0);
            self.session_bytes.take(cost);
        }
        ok
    }

    /// Keep the author map bounded: past the cap, drop the least-recently
    /// active entry that isn't the author being admitted. A pruned author
    /// returns with full buckets, but authors are roster-gated upstream, so
    /// the map can't be churned by strangers.
    fn prune(&mut self, keep: EndpointId, _now_ms: u64) {
        while self.authors.len() >= AUTHOR_MAP_CAP {
            let Some(victim) = self
                .authors
                .iter()
                .filter(|(id, _)| **id != keep)
                .min_by_key(|(_, b)| b.last_seen_ms)
                .map(|(id, _)| *id)
            else {
                break;
            };
            self.authors.remove(&victim);
        }
    }

    #[cfg(test)]
    fn author_count(&self) -> usize {
        self.authors.len()
    }
}

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

    const T0: u64 = 1_000_000;
    const MIB: u64 = 1024 * 1024;

    fn author() -> EndpointId {
        SecretKey::generate().public()
    }

    #[test]
    fn author_request_burst_then_refill_recovers() {
        let mut b = AutoFetchBudget::new(T0);
        let a = author();
        // Tiny sizes so only the REQUEST buckets can bind.
        for _ in 0..AUTHOR_REQ_BURST as usize {
            assert!(b.admit(a, 1, T0));
        }
        assert!(!b.admit(a, 1, T0), "author request burst exhausted");
        // One request refills after 10 s.
        assert!(b.admit(a, 1, T0 + 10_000));
        assert!(!b.admit(a, 1, T0 + 10_000));
    }

    #[test]
    fn author_byte_budget_binds_and_recovers() {
        let mut b = AutoFetchBudget::new(T0);
        let a = author();
        // 4 × 4 MiB = the full 16 MiB author byte burst (well under the
        // 8-request burst, so bytes are the binding constraint).
        for _ in 0..4 {
            assert!(b.admit(a, 4 * MIB, T0));
        }
        assert!(!b.admit(a, 4 * MIB, T0), "author byte burst exhausted");
        // 64 KiB/s → a 4 MiB image is affordable again after 64 s (which also
        // refills 6 request tokens, so bytes stay the binding constraint).
        assert!(!b.admit(a, 4 * MIB, T0 + 32_000));
        assert!(b.admit(a, 4 * MIB, T0 + 64_000));
    }

    #[test]
    fn session_budget_binds_across_authors_without_burning_author_tokens() {
        let mut b = AutoFetchBudget::new(T0);
        // Three authors × 16 MiB exhausts the 48 MiB session byte burst even
        // though each author is within their own budget.
        for _ in 0..3 {
            let a = author();
            for _ in 0..4 {
                assert!(b.admit(a, 4 * MIB, T0));
            }
        }
        let fresh = author();
        assert!(!b.admit(fresh, 4 * MIB, T0), "session bytes exhausted");
        // The rejection consumed NOTHING: once the session refills enough for
        // one image (4 MiB / 128 KiB/s = 32 s), the fresh author's own full
        // burst is intact and admits immediately.
        assert!(b.admit(fresh, 4 * MIB, T0 + 32_000));
    }

    #[test]
    fn session_request_bucket_binds_across_authors() {
        let mut b = AutoFetchBudget::new(T0);
        // 16 tiny requests from distinct authors exhaust the session request
        // burst while every author bucket stays nearly full.
        for _ in 0..SESSION_REQ_BURST as usize {
            assert!(b.admit(author(), 1, T0));
        }
        assert!(!b.admit(author(), 1, T0), "session requests exhausted");
        assert!(b.admit(author(), 1, T0 + 5_000), "one recovers after 5 s");
    }

    #[test]
    fn author_map_stays_bounded_pruning_least_recent() {
        let mut b = AutoFetchBudget::new(T0);
        // Session request refill would bind over a naive loop; space the
        // admissions out so only the map bound is under test.
        let mut t = T0;
        let first = author();
        assert!(b.admit(first, 1, t));
        for _ in 0..(AUTHOR_MAP_CAP + 10) {
            t += 10_000;
            assert!(b.admit(author(), 1, t));
            assert!(b.author_count() <= AUTHOR_MAP_CAP);
        }
        assert!(b.author_count() <= AUTHOR_MAP_CAP);
    }
}