Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199use crate::config::RateTier;use arc_swap::ArcSwap;use serde::{Deserialize, Serialize};use smol_str::SmolStr;use std::collections::HashMap;use std::sync::Arc;use tracing::debug;
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]pub enum HostStatus { #[default] Active, Idle, Offline, Throttled, Banned,}
impl From<HostStatus> for jacquard_api::com_atproto::sync::HostStatus { fn from(status: HostStatus) -> Self { match status { HostStatus::Active => Self::Active, HostStatus::Idle => Self::Idle, HostStatus::Offline => Self::Offline, HostStatus::Throttled => Self::Throttled, HostStatus::Banned => Self::Banned, } }}
impl HostStatus { pub const fn as_str(self) -> &'static str { match self { Self::Active => "active", Self::Idle => "idle", Self::Offline => "offline", Self::Throttled => "throttled", Self::Banned => "banned", } }
/// returns the new status to apply if the count dynamically crossed limits. pub fn check_limit_transition( self, current_count: u64, account_limit: Option<u64>, ) -> Option<Self> { match account_limit { Some(limit) if current_count >= limit && self != Self::Throttled => { Some(Self::Throttled) } Some(limit) if current_count < limit && self == Self::Throttled => Some(Self::Active), None if self == Self::Throttled => Some(Self::Active), _ => None, } }}
#[derive(Debug, Default, Clone)]pub struct HostDesc { pub tier: Option<SmolStr>, pub status: HostStatus, pub operator_ban: bool,}
#[derive(Default, Clone)]pub(crate) struct PdsMeta { pub hosts: HashMap<String, HostDesc>,}
impl PdsMeta { /// update (or insert) the `HostDesc` for `host` by applying `f` to it. pub fn update_host_entry(&mut self, host: &str, f: impl FnOnce(&mut HostDesc)) { f(self.hosts.entry(host.to_string()).or_default()); }
/// atomically update (or insert) the `HostDesc` for `host` by applying `f` to it via RCU. pub(crate) fn update_host( cell: &arc_swap::ArcSwap<Self>, host: &str, mut f: impl FnMut(&mut HostDesc), ) { cell.rcu(|meta| { let mut next = (**meta).clone(); next.update_host_entry(host, &mut f); next }); }
pub fn status(&self, host: &str) -> HostStatus { self.hosts .get(host) .map(|h| h.status) .unwrap_or(HostStatus::Active) }
pub fn is_banned(&self, host: &str) -> bool { self.hosts .get(host) .is_some_and(|desc| desc.operator_ban || desc.status == HostStatus::Banned) }}
pub(crate) type PdsMetaHandle = Arc<ArcSwap<PdsMeta>>;
pub(crate) fn new_handle(meta: PdsMeta) -> PdsMetaHandle { Arc::new(ArcSwap::new(Arc::new(meta)))}
#[derive(Debug, Clone)]pub struct TierRule { pub pattern: glob::Pattern, pub tier_name: SmolStr,}
/// policy for resolving rate tiers for PDS hosts.////// resolution order:/// 1. explicit api-assigned override for the host (stored in `HostDesc.tier`)/// 2. first matching rule in `rules` (config glob patterns, in order)/// 3. built-in `"default"` tier#[derive(Debug, Clone)]pub struct TierPolicy { /// named rate tier definitions. pub tiers: HashMap<SmolStr, RateTier>, /// ordered glob rules, first match wins (for unassigned hosts). pub rules: Vec<TierRule>,}
impl TierPolicy { /// resolves the effective `RateTier` for `host`. /// /// `override_name` is the api-assigned tier name from `HostDesc.tier`, if any. pub fn resolve(&self, host: &str, override_name: Option<&SmolStr>) -> RateTier { let default = self .tiers .get("default") .copied() .unwrap_or(crate::config::DEFAULT_RATETIER);
if let Some(name) = override_name { let tier = self.tiers.get(name).copied().unwrap_or(default); debug!(host, override = %name, account_limit = ?tier.account_limit, "tier resolved via explicit override"); return tier; }
let matched = self.rules.iter().find(|r| r.pattern.matches(host));
let tier = matched .and_then(|r| self.tiers.get(&r.tier_name).copied()) .unwrap_or(default);
debug!( host, matched_rule = matched.map(|r| format!("{}:{}", r.pattern, r.tier_name)).as_deref(), account_limit = ?tier.account_limit, "tier resolved via glob rules" ); tier }}
#[cfg(test)]mod tests { use super::*;
#[test] fn host_status_persisted_golden_fixtures() { for (status, msgpack, json) in [ (HostStatus::Active, "a6416374697665", "\"Active\""), (HostStatus::Idle, "a449646c65", "\"Idle\""), (HostStatus::Offline, "a74f66666c696e65", "\"Offline\""), ( HostStatus::Throttled, "a95468726f74746c6564", "\"Throttled\"", ), (HostStatus::Banned, "a642616e6e6564", "\"Banned\""), ] { let encoded = rmp_serde::to_vec(&status).unwrap(); assert_eq!(hex::encode(&encoded), msgpack); assert_eq!( rmp_serde::from_slice::<HostStatus>(&encoded).unwrap(), status ); assert_eq!(serde_json::to_string(&status).unwrap(), json); let external: jacquard_api::com_atproto::sync::HostStatus = status.into(); let expected_external = match status { HostStatus::Active => jacquard_api::com_atproto::sync::HostStatus::Active, HostStatus::Idle => jacquard_api::com_atproto::sync::HostStatus::Idle, HostStatus::Offline => jacquard_api::com_atproto::sync::HostStatus::Offline, HostStatus::Throttled => jacquard_api::com_atproto::sync::HostStatus::Throttled, HostStatus::Banned => jacquard_api::com_atproto::sync::HostStatus::Banned, }; assert_eq!(external, expected_external); } }}