didbot_serve/

subscribers.rs

//! How many streams this server holds open at once, and how many of them one
//! caller may hold.
//!
//! `com.atproto.sync.subscribeRepos` is the route whose answer does not end. Every other route's cost is bounded by
//! [`crate::MAX_CONCURRENT_REQUESTS`], which releases its permit once the
//! response head is sent — which for a stream is immediately, and for the
//! next hour of that stream's life is nothing at all. What each one keeps
//! after that is a task and a frame in flight per subscriber, and the
//! routes are public, so without a bound of their own a caller can open them
//! until the process runs out of memory.
//!
//! Two bounds, for the reason `crate::rate_limit`'s own doc gives for its
//! two: a total, which is the resource actually shared, and a per-address
//! one, so a single caller cannot spend the whole total on its own. A slot is
//! held for the life of the stream rather than the life of the request, which
//! is what makes it a bound on concurrency rather than on arrival rate.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

/// How many streams this server holds open across every caller.
///
/// Well past what an honest deployment reaches — a relay, an operator's
/// dashboard, and one daemon per agent host — and far enough under
/// [`crate::MAX_CONCURRENT_REQUESTS`] that a flood of subscriptions cannot
/// crowd out ordinary requests.
pub const DEFAULT_MAX_SUBSCRIBERS: usize = 128;

/// How many of those one caller address may hold.
///
/// An agent host runs one daemon, and a relay crawls with one connection. A
/// browser holding a dashboard open is the case that legitimately wants more
/// than one, which is what this leaves room for.
pub const DEFAULT_MAX_SUBSCRIBERS_PER_CALLER: usize = 8;

/// The two bounds, and what is currently held against them.
#[derive(Debug)]
pub struct SubscriberLimit {
    total: usize,
    per_caller: usize,
    held: Mutex<Counts>,
}

#[derive(Debug, Default)]
struct Counts {
    /// How many streams are open in total.
    total: usize,
    /// How many each caller holds. A caller with none is not in the map, so
    /// this is bounded by `total` rather than by how many addresses have ever
    /// subscribed.
    per_caller: HashMap<String, usize>,
}

impl Default for SubscriberLimit {
    fn default() -> Self {
        Self::new(DEFAULT_MAX_SUBSCRIBERS, DEFAULT_MAX_SUBSCRIBERS_PER_CALLER)
    }
}

impl SubscriberLimit {
    /// A limiter with bounds of a deployment's own choosing.
    pub fn new(total: usize, per_caller: usize) -> Self {
        Self {
            total,
            per_caller,
            held: Mutex::new(Counts::default()),
        }
    }

    /// Takes a slot for `caller`, or `None` if either bound is full.
    ///
    /// The slot is released when what it is handed back to drops, so a caller
    /// that disconnects gives its slot back without anything having to notice
    /// the disconnection.
    pub fn take(self: &Arc<Self>, caller: &str) -> Option<SubscriberSlot> {
        let mut held = self.held.lock().unwrap_or_else(|p| p.into_inner());
        if held.total >= self.total {
            return None;
        }
        let mine = held.per_caller.entry(caller.to_owned()).or_insert(0);
        if *mine >= self.per_caller {
            // Left at zero rather than removed: the next `take` for this
            // caller re-inserts it, and an entry at zero is the same size as
            // the absence of one for as long as this lock is held.
            if *mine == 0 {
                held.per_caller.remove(caller);
            }
            return None;
        }
        *mine += 1;
        held.total += 1;
        Some(SubscriberSlot {
            limit: Arc::clone(self),
            caller: caller.to_owned(),
        })
    }

    /// How many streams are open right now. For a test, and for anything
    /// that later wants to log it.
    pub fn held(&self) -> usize {
        self.held.lock().unwrap_or_else(|p| p.into_inner()).total
    }

    fn release(&self, caller: &str) {
        let mut held = self.held.lock().unwrap_or_else(|p| p.into_inner());
        held.total = held.total.saturating_sub(1);
        if let Some(mine) = held.per_caller.get_mut(caller) {
            *mine -= 1;
            if *mine == 0 {
                held.per_caller.remove(caller);
            }
        }
    }
}

/// One open stream's claim on [`SubscriberLimit`], released when it drops.
///
/// Carried by the stream itself rather than by the handler, because the
/// handler returns as soon as the response head is written and the stream
/// outlives it by as long as the subscriber stays connected.
#[derive(Debug)]
pub struct SubscriberSlot {
    limit: Arc<SubscriberLimit>,
    caller: String,
}

impl Drop for SubscriberSlot {
    fn drop(&mut self) {
        self.limit.release(&self.caller);
    }
}

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

    #[test]
    fn a_caller_may_not_hold_more_than_its_share() {
        let limit = Arc::new(SubscriberLimit::new(10, 2));
        let first = limit.take("1.2.3.4").expect("the first fits");
        let _second = limit.take("1.2.3.4").expect("the second fits");
        assert!(limit.take("1.2.3.4").is_none(), "a third is past the share");

        // Another caller is unaffected, and the first gets its slot back the
        // moment one is dropped.
        assert!(limit.take("5.6.7.8").is_some());
        drop(first);
        assert!(limit.take("1.2.3.4").is_some());
    }

    #[test]
    fn the_total_refuses_whoever_asks_past_it() {
        let limit = Arc::new(SubscriberLimit::new(2, 10));
        let held: Vec<_> = ["a", "b"]
            .iter()
            .map(|caller| limit.take(caller).expect("both fit"))
            .collect();
        assert!(limit.take("c").is_none());
        assert_eq!(limit.held(), 2);
        drop(held);
        assert_eq!(limit.held(), 0);
        assert!(limit.take("c").is_some());
    }

    /// A caller that has given every slot back leaves nothing behind, so the
    /// map is a function of what is open rather than of how many addresses
    /// have ever subscribed — which, on a public route, is unbounded.
    #[test]
    fn a_caller_holding_nothing_is_not_remembered() {
        let limit = Arc::new(SubscriberLimit::new(100, 1));
        for n in 0..1000 {
            let slot = limit.take(&format!("10.0.0.{n}")).expect("each fits");
            drop(slot);
        }
        assert_eq!(limit.held(), 0);
        assert!(limit.held.lock().unwrap().per_caller.is_empty());
    }
}