didbot_serve/

write_meter.rs

//! Per-account write metering: one token bucket on write operations and one
//! on request bytes, both keyed by the authenticated account.
//!
//! Every account shares one write-ahead log, and `[capacity]
//! log_budget_bytes` bounds it for all of them together. The server trims the
//! log at `[capacity] checkpoint_log_bytes`, so what the budget has to hold is
//! what all the accounts write between two checkpoints. The meter bounds how
//! fast one account can grow it, so an account in a loop is one share of that
//! rather than all of it. It bounds a rate, not a total: every account has its
//! own buckets.
//!
//! # What is charged
//!
//! `createRecord`, `putRecord` and `deleteRecord` cost one operation;
//! `applyWrites` costs one per write in its batch. Each also costs its
//! request body's length in bytes. `uploadBlob` costs one operation and no
//! bytes, because `[blobs] account_quota_bytes` already bounds what an
//! account's blobs hold. A request is charged once write policy has allowed
//! it, just before it commits, so neither a token that may not write nor a
//! client that policy refuses can spend the account's share. A write the
//! store refuses after that, such as a failed swap, has still spent its
//! tokens. A write refused by the byte bucket has already spent its
//! operations.
//!
//! The buckets are in memory and start full when the server starts.
//!
//! # The defaults, worked through
//!
//! A write's cost to the log does not depend on the record's size: the body
//! goes to the record heap, and the log gets the record's slot and the
//! commit naming it. Measured, that is about 1 KiB for a single
//! `createRecord`, `putRecord` or `deleteRecord`, and about 240 bytes per
//! write inside a 200-write `applyWrites`. So the operation bucket is what
//! bounds the log:
//!
//! - At most [`DEFAULT_OPS_BURST`] + [`DEFAULT_OPS_PER_HOUR`] = 4,600
//!   operations in any hour.
//! - 4,600 × 1 KiB ≈ 4.5 MiB of log an hour, about 0.2% of the 2 GiB budget
//!   a default 20 GiB volume gets.
//! - One account writing flat out reaches the default 1 GiB watermark, where
//!   the server checkpoints and trims, in about 227 hours. A repository deep
//!   enough to double the cost of a write halves that.
//!
//! The byte bucket bounds the record heap, which shares the volume:
//!
//! - At most [`DEFAULT_BYTES_BURST`] + [`DEFAULT_BYTES_PER_HOUR`] = 72 MiB of
//!   request bodies in any hour, about 0.35% of a 20 GiB volume.
//! - A loop of full 2 MiB requests gets four back to back, then one every
//!   112 seconds.
//!
//! An agent writing a 1 KiB record every five seconds spends 720 operations
//! and under 1 MiB an hour: 20% of the operation rate and about 1% of the
//! byte rate.

use std::num::NonZeroU32;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use governor::clock::{Clock, DefaultClock};
use governor::middleware::NoOpMiddleware;
use governor::state::keyed::HashMapStateStore;
use governor::{Quota, RateLimiter};

/// [`WriteLimits::ops_per_hour`]'s default: one operation a second.
pub const DEFAULT_OPS_PER_HOUR: u32 = 3_600;

/// [`WriteLimits::ops_burst`]'s default.
pub const DEFAULT_OPS_BURST: u32 = 1_000;

/// [`WriteLimits::bytes_per_hour`]'s default: 64 MiB.
pub const DEFAULT_BYTES_PER_HOUR: u64 = 64 * 1024 * 1024;

/// [`WriteLimits::bytes_burst`]'s default: 8 MiB, four full request bodies.
pub const DEFAULT_BYTES_BURST: u64 = 8 * 1024 * 1024;

/// How fast one account may write, and how much it may write at once.
///
/// `[writes]` in the configuration file; each field's range is documented
/// there, on `didbot_config::sections::WritesSection`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WriteLimits {
    /// Operations an account's bucket regains an hour.
    pub ops_per_hour: u32,
    /// Operations an account's bucket holds when full.
    pub ops_burst: u32,
    /// Request bytes an account's bucket regains an hour.
    pub bytes_per_hour: u64,
    /// Request bytes an account's bucket holds when full.
    pub bytes_burst: u64,
}

impl Default for WriteLimits {
    fn default() -> Self {
        Self {
            ops_per_hour: DEFAULT_OPS_PER_HOUR,
            ops_burst: DEFAULT_OPS_BURST,
            bytes_per_hour: DEFAULT_BYTES_PER_HOUR,
            bytes_burst: DEFAULT_BYTES_BURST,
        }
    }
}

/// A `[writes]` key written outside the range it is documented with.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("{field} must be {low} to {high}, not {value}")]
pub struct WriteLimitError {
    /// The key, spelled as `[writes]` spells it.
    pub field: &'static str,
    /// The value that was written.
    pub value: u64,
    /// The lowest value the key accepts.
    pub low: u64,
    /// The highest value the key accepts.
    pub high: u64,
}

impl WriteLimits {
    /// Checks each key against the range `WritesSection` documents.
    ///
    /// The two burst floors are hard: a smaller `ops_burst` refuses every
    /// full `applyWrites`, and a smaller `bytes_burst` refuses every request
    /// body near the size limit, however long the caller waits.
    pub fn check(&self) -> Result<(), WriteLimitError> {
        for (field, value, low, high) in [
            ("ops_per_hour", u64::from(self.ops_per_hour), 600, 360_000),
            (
                "ops_burst",
                u64::from(self.ops_burst),
                crate::wire::MAX_APPLY_WRITES as u64,
                100_000,
            ),
            ("bytes_per_hour", self.bytes_per_hour, 8 << 20, 4 << 30),
            (
                "bytes_burst",
                self.bytes_burst,
                crate::routes::MAX_REQUEST_BODY as u64,
                1 << 30,
            ),
        ] {
            if value < low || value > high {
                return Err(WriteLimitError {
                    field,
                    value,
                    low,
                    high,
                });
            }
        }
        Ok(())
    }
}

/// Which bucket refused a write.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Bucket {
    /// [`WriteLimits::ops_per_hour`] and [`WriteLimits::ops_burst`].
    Ops,
    /// [`WriteLimits::bytes_per_hour`] and [`WriteLimits::bytes_burst`].
    Bytes,
}

/// A write the meter refused.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Refusal {
    /// The bucket that was short.
    pub bucket: Bucket,
    /// What the write asked that bucket for.
    pub wanted: u64,
    /// How long until the bucket holds `wanted`. `None` when it never can,
    /// because `wanted` is more than the bucket holds when full.
    pub retry_after: Option<Duration>,
}

/// How often [`WriteMeter::charge`] drops buckets that have refilled.
const SWEEP_EVERY: Duration = Duration::from_secs(60);

type Keyed<C> =
    RateLimiter<String, HashMapStateStore<String>, C, NoOpMiddleware<<C as Clock>::Instant>>;

/// The two buckets for every account that has written lately. A bucket
/// that has refilled is dropped, since a missing one starts full anyway.
pub struct WriteMeter<C: Clock = DefaultClock> {
    limits: WriteLimits,
    ops: Keyed<C>,
    bytes: Keyed<C>,
    clock: C,
    last_swept: Mutex<Instant>,
}

impl WriteMeter {
    /// A meter on the system's monotonic clock.
    pub fn new(limits: WriteLimits) -> Self {
        Self::with_clock(limits, DefaultClock::default())
    }
}

impl Default for WriteMeter {
    fn default() -> Self {
        Self::new(WriteLimits::default())
    }
}

impl<C: Clock + Clone> WriteMeter<C> {
    /// A meter on `clock`, which a test can advance by hand.
    ///
    /// # Panics
    ///
    /// If a rate or burst in `limits` is zero, a burst is past `u32::MAX`, or
    /// a rate refills a token in under a nanosecond. Overflows if a slow rate
    /// meets a large burst. [`WriteLimits::check`] refuses each of these.
    pub fn with_clock(limits: WriteLimits, clock: C) -> Self {
        let hour = Duration::from_secs(3600);
        let ops = quota(hour / limits.ops_per_hour, u64::from(limits.ops_burst));
        let bytes = quota(
            Duration::from_nanos(
                u64::try_from(hour.as_nanos() / u128::from(limits.bytes_per_hour))
                    .unwrap_or(u64::MAX),
            ),
            limits.bytes_burst,
        );
        Self {
            limits,
            ops: RateLimiter::hashmap_with_clock(ops, clock.clone()),
            bytes: RateLimiter::hashmap_with_clock(bytes, clock.clone()),
            clock,
            last_swept: Mutex::new(Instant::now()),
        }
    }

    /// The limits this meter enforces.
    pub fn limits(&self) -> WriteLimits {
        self.limits
    }

    /// Spends `ops` operations and `bytes` bytes of `did`'s buckets, or
    /// refuses and spends nothing more.
    ///
    /// Operations are spent first. A write the byte bucket refuses keeps
    /// them spent.
    pub fn charge(&self, did: &str, ops: u32, bytes: u64) -> Result<(), Refusal> {
        self.sweep();
        let key = did.to_owned();
        self.spend(&self.ops, &key, Bucket::Ops, u64::from(ops))?;
        self.spend(&self.bytes, &key, Bucket::Bytes, bytes)
    }

    /// Drops refilled buckets, at most once per [`SWEEP_EVERY`], so the maps
    /// hold the accounts writing now rather than every account that has.
    fn sweep(&self) {
        let mut last = self
            .last_swept
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        if last.elapsed() < SWEEP_EVERY {
            return;
        }
        *last = Instant::now();
        for limiter in [&self.ops, &self.bytes] {
            limiter.retain_recent();
            limiter.shrink_to_fit();
        }
    }

    fn spend(
        &self,
        limiter: &Keyed<C>,
        key: &String,
        bucket: Bucket,
        wanted: u64,
    ) -> Result<(), Refusal> {
        let Some(n) = u32::try_from(wanted).ok().and_then(NonZeroU32::new) else {
            if wanted == 0 {
                return Ok(());
            }
            return Err(Refusal {
                bucket,
                wanted,
                retry_after: None,
            });
        };
        match limiter.check_key_n(key, n) {
            Ok(Ok(())) => Ok(()),
            Ok(Err(not_until)) => Err(Refusal {
                bucket,
                wanted,
                retry_after: Some(not_until.wait_time_from(self.clock.now())),
            }),
            Err(_insufficient) => Err(Refusal {
                bucket,
                wanted,
                retry_after: None,
            }),
        }
    }
}

fn quota(per_token: Duration, burst: u64) -> Quota {
    Quota::with_period(per_token)
        .expect("a write rate refills each token in at least a nanosecond")
        .allow_burst(
            u32::try_from(burst)
                .ok()
                .and_then(NonZeroU32::new)
                .expect("a write burst is 1 to u32::MAX"),
        )
}

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

    const ALICE: &str = "did:web:alice.agents.localhost";

    #[test]
    fn the_defaults_are_inside_their_documented_ranges() {
        assert_eq!(WriteLimits::default().check(), Ok(()));
    }

    #[test]
    fn a_burst_below_one_full_request_is_refused() {
        let limits = WriteLimits {
            bytes_burst: 1024 * 1024,
            ..WriteLimits::default()
        };
        assert_eq!(limits.check().unwrap_err().field, "bytes_burst");
        let limits = WriteLimits {
            ops_burst: 100,
            ..WriteLimits::default()
        };
        assert_eq!(limits.check().unwrap_err().field, "ops_burst");
    }

    /// A day of an agent at work: a 1 KiB record every five seconds, a
    /// 50-write `applyWrites` of 200 KiB every ten minutes, and a blob every
    /// hour. None of it is refused at the defaults.
    #[test]
    fn a_busy_agent_stays_under_the_defaults_for_a_day() {
        let clock = FakeRelativeClock::default();
        let meter = WriteMeter::with_clock(WriteLimits::default(), clock.clone());
        for tick in 0..(24 * 3600 / 5) {
            meter
                .charge(ALICE, 1, 1024)
                .unwrap_or_else(|refusal| panic!("tick {tick}: {refusal:?}"));
            if tick % 120 == 0 {
                meter
                    .charge(ALICE, 50, 200 * 1024)
                    .unwrap_or_else(|refusal| panic!("batch at tick {tick}: {refusal:?}"));
            }
            if tick % 720 == 0 {
                meter
                    .charge(ALICE, 1, 0)
                    .unwrap_or_else(|refusal| panic!("blob at tick {tick}: {refusal:?}"));
            }
            clock.advance(Duration::from_secs(5));
        }
    }

    /// A loop of full-size requests that waits out each refusal gets its
    /// burst, then the byte rate, and no more than the two in an hour.
    #[test]
    fn a_loop_of_full_size_requests_is_held_to_burst_plus_rate() {
        let clock = FakeRelativeClock::default();
        let limits = WriteLimits::default();
        let meter = WriteMeter::with_clock(limits, clock.clone());
        let body = crate::routes::MAX_REQUEST_BODY as u64;
        let hour = Duration::from_secs(3600);
        let (mut elapsed, mut written) = (Duration::ZERO, 0u64);
        while elapsed < hour {
            match meter.charge(ALICE, 1, body) {
                Ok(()) => written += body,
                Err(refusal) => {
                    assert_eq!(refusal.bucket, Bucket::Bytes);
                    let wait = refusal.retry_after.expect("a full body fits the bucket");
                    clock.advance(wait);
                    elapsed += wait;
                }
            }
        }
        let ceiling = limits.bytes_burst + limits.bytes_per_hour;
        assert!(written <= ceiling, "{written} bytes past {ceiling}");
        assert!(written >= ceiling - body, "only {written} bytes admitted");
    }

    #[test]
    fn more_than_a_full_bucket_can_never_be_admitted() {
        let meter = WriteMeter::new(WriteLimits::default());
        let refusal = meter.charge(ALICE, DEFAULT_OPS_BURST + 1, 0).unwrap_err();
        assert_eq!(refusal.bucket, Bucket::Ops);
        assert_eq!(refusal.retry_after, None);
    }
}