rusting_engine 2.7.0

//! Per-peer rate limits for what clients send the host, and bandwidth
//! budgets for what the host sends each client.

use super::{invalid, NetSession, PeerId};
use std::collections::BTreeMap;
use std::io;
use std::time::{Duration, Instant};

/// A token bucket per peer: each peer may send `burst` messages at once,
/// then `per_second` on average. Ask it before handling a client message
/// and drop (or count against the sender) the ones it refuses.
///
/// [`RateLimit::allow_at`] times the buckets by the fixed tick, so the same
/// messages on the same ticks get the same answers on every run;
/// [`RateLimit::allow`] times them by the wall clock. Use one of the two
/// per limiter.
#[derive(Clone, Debug)]
pub struct RateLimit {
    per_second: f64,
    burst: f64,
    /// Tokens left and the time they were counted at, in seconds.
    buckets: BTreeMap<PeerId, (f64, f64)>,
    /// What [`RateLimit::allow`] counts its seconds from: the first
    /// `Instant` it saw.
    epoch: Option<Instant>,
}

impl RateLimit {
    /// A limit of `per_second` messages a second per peer, `burst` at once.
    pub fn new(per_second: f64, burst: u32) -> Self {
        Self {
            per_second,
            burst: f64::from(burst),
            buckets: BTreeMap::new(),
            epoch: None,
        }
    }

    /// Takes one message from `peer` at `now`; `false` when it is over its
    /// rate.
    pub fn allow(&mut self, peer: PeerId, now: Instant) -> bool {
        let seconds = self.seconds(now);
        self.take(peer, 1.0, seconds)
    }

    /// Takes one message from `peer` on fixed tick `tick`, ticks being `dt`
    /// apart: pass `time.fixed_tick` and `time.fixed_delta` from
    /// [`FrameTime`](crate::runtime::FrameTime). `false` when it is over its rate.
    /// Unlike [`RateLimit::allow`] it reads no clock, so a server that
    /// floods it in a scenario refuses the same messages every run.
    pub fn allow_at(&mut self, peer: PeerId, tick: u64, dt: Duration) -> bool {
        self.take(peer, 1.0, tick as f64 * dt.as_secs_f64())
    }

    /// Seconds from the first `Instant` seen to `now`.
    fn seconds(&mut self, now: Instant) -> f64 {
        let epoch = *self.epoch.get_or_insert(now);
        now.saturating_duration_since(epoch).as_secs_f64()
    }

    fn take(&mut self, peer: PeerId, cost: f64, now: f64) -> bool {
        let (tokens, last) =
            self.buckets.entry(peer).or_insert((self.burst, now));
        let earned = (now - *last).max(0.0);
        *tokens = (*tokens + earned * self.per_second).min(self.burst);
        *last = now;
        if *tokens >= cost {
            *tokens -= cost;
            true
        } else {
            false
        }
    }

    /// Forgets `peer`, on `NetEvent::Disconnected`.
    pub fn forget(&mut self, peer: PeerId) {
        self.buckets.remove(&peer);
    }
}

/// What a [`Bandwidth`] measured for one client. Bytes are payload bytes.
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct Usage {
    pub bytes_sent: u64,
    pub messages_sent: u64,
    /// Messages refused because the client was over its budget.
    pub messages_dropped: u64,
    /// Bytes sent during the last whole second.
    pub bytes_per_second: u64,
    window: Option<(Instant, u64)>,
}

/// A bandwidth budget per client: the host may send each one `burst`
/// bytes at once, then `bytes_per_second` on average, and the budget
/// measures what each one got. Send per-tick state through it; a refused
/// message is simply not sent, so use it for messages that a later one
/// replaces, such as [`snapshot_for`](super::replicate::Replication::snapshot_for).
#[derive(Clone, Debug)]
pub struct Bandwidth {
    cap: RateLimit,
    usage: BTreeMap<PeerId, Usage>,
}

impl Bandwidth {
    /// A budget of `bytes_per_second` per client, `burst` bytes at once.
    pub fn new(bytes_per_second: f64, burst: u32) -> Self {
        Self {
            cap: RateLimit::new(bytes_per_second, burst),
            usage: BTreeMap::new(),
        }
    }

    /// Sends `bytes` to `to` unreliably if its budget allows at `now`.
    /// `Ok(false)` when it does not. Fails for a message bigger than the
    /// burst, which could never be sent.
    pub fn send_unreliable(
        &mut self,
        session: &NetSession,
        to: PeerId,
        bytes: &[u8],
        now: Instant,
    ) -> io::Result<bool> {
        let size = bytes.len() as f64;
        if size > self.cap.burst {
            return Err(invalid("message is bigger than the bandwidth burst"));
        }
        let seconds = self.cap.seconds(now);
        let allowed = self.cap.take(to, size, seconds);
        if allowed {
            session.send_unreliable(to, bytes)?;
        }
        let usage = self.usage.entry(to).or_default();
        let (start, bytes_in_window) = usage.window.get_or_insert((now, 0));
        if now.saturating_duration_since(*start) >= Duration::from_secs(1) {
            usage.bytes_per_second = *bytes_in_window;
            *start = now;
            *bytes_in_window = 0;
        }
        if allowed {
            *bytes_in_window += bytes.len() as u64;
            usage.bytes_sent += bytes.len() as u64;
            usage.messages_sent += 1;
        } else {
            usage.messages_dropped += 1;
        }
        Ok(allowed)
    }

    /// What `peer` was sent so far.
    pub fn usage(&self, peer: PeerId) -> Usage {
        self.usage.get(&peer).copied().unwrap_or_default()
    }

    /// Forgets `peer`, on `NetEvent::Disconnected`.
    pub fn forget(&mut self, peer: PeerId) {
        self.cap.forget(peer);
        self.usage.remove(&peer);
    }
}

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

    #[test]
    fn each_peer_gets_a_burst_then_its_rate() {
        let mut limit = RateLimit::new(10.0, 3);
        let start = Instant::now();
        let allowed = |limit: &mut RateLimit, peer, at| {
            (0..5).filter(|_| limit.allow(peer, at)).count()
        };
        assert_eq!(allowed(&mut limit, 1, start), 3);
        assert_eq!(allowed(&mut limit, 2, start), 3); // its own bucket
        let later = start + Duration::from_millis(200);
        assert_eq!(allowed(&mut limit, 1, later), 2);
        // A long pause refills only up to the burst.
        let much_later = start + Duration::from_secs(60);
        assert_eq!(allowed(&mut limit, 1, much_later), 3);
        limit.forget(1);
        assert_eq!(allowed(&mut limit, 1, much_later), 3);
    }

    #[test]
    fn a_tick_timed_limit_gives_the_same_answers_every_run() {
        let dt = Duration::from_secs_f64(1.0 / 60.0);
        // A client sending three messages every tick against 30 a second,
        // bursts of 10: the answers depend only on the ticks.
        let flood = || {
            let mut limit = RateLimit::new(30.0, 10);
            (1..=120u64)
                .flat_map(|tick| (0..3).map(move |_| tick))
                .map(|tick| limit.allow_at(7, tick, dt))
                .collect::<Vec<_>>()
        };
        let answers = flood();
        assert_eq!(answers, flood(), "the same run gives the same answers");
        // The burst, then one every other tick for the 119 ticks after.
        let allowed = answers.iter().filter(|ok| **ok).count();
        assert!((69..=70).contains(&allowed), "{allowed}");
        assert!(answers[..11].iter().all(|ok| *ok), "burst and refill");
        assert!(!answers[11]);

        let mut limit = RateLimit::new(10.0, 3);
        let allowed = |limit: &mut RateLimit, peer, tick| {
            (0..5).filter(|_| limit.allow_at(peer, tick, dt)).count()
        };
        assert_eq!(allowed(&mut limit, 1, 0), 3);
        assert_eq!(allowed(&mut limit, 2, 0), 3); // its own bucket
        assert_eq!(allowed(&mut limit, 1, 12), 2, "0.2 s later");
        assert_eq!(allowed(&mut limit, 1, 6000), 3, "refills to the burst");
    }

    #[test]
    fn bandwidth_caps_and_measures_each_client() {
        let (host, _clients) = NetSession::loopback(2);
        host.poll();
        let mut budget = Bandwidth::new(1000.0, 400);
        let start = Instant::now();
        let message = [0u8; 100];
        let sent = |budget: &mut Bandwidth, peer, at| {
            (0..10)
                .filter(|_| {
                    budget.send_unreliable(&host, peer, &message, at).unwrap()
                })
                .count()
        };
        assert_eq!(sent(&mut budget, 1, start), 4, "the burst");
        assert_eq!(sent(&mut budget, 2, start), 4, "its own budget");
        let later = start + Duration::from_secs(1);
        assert_eq!(sent(&mut budget, 1, later), 4, "refilled to the burst");
        let usage = budget.usage(1);
        assert_eq!((usage.bytes_sent, usage.messages_sent), (800, 8));
        assert_eq!(usage.messages_dropped, 12);
        assert_eq!(usage.bytes_per_second, 400, "the first second");
        assert_eq!(host.stats().bytes_sent, 1200 + 12 * 9);
        assert!(budget.send_unreliable(&host, 1, &[0; 401], later).is_err());
        budget.forget(1);
        assert_eq!(budget.usage(1), Usage::default());
    }
}