use 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 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)))
}
fn glob_match(pat: &[u8], text: &[u8]) -> bool {
let (mut p, mut t, mut star) = (0, 0, None);
while t < text.len() {
if p < pat.len() && pat[p] == b'*' {
star = Some((p + 1, t));
p += 1;
} else if p < pat.len() && (pat[p] == b'?' || pat[p] == text[t]) {
p += 1;
t += 1;
} else if let Some((sp, st)) = &mut star {
p = *sp;
*st += 1;
t = *st;
} else {
return false;
}
}
pat[p..].iter().all(|&b| b == b'*')
}
#[derive(Debug, Clone)]
pub struct TierRule {
pub pattern: SmolStr,
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 {
/// an unknown override gives "default" without trying the host's rule
pub fn resolve_name<'a>(&'a self, host: &str, override_name: Option<&'a SmolStr>) -> &'a str {
let rule = || {
self.rules
.iter()
.find(|r| glob_match(r.pattern.as_bytes(), host.as_bytes()))
.map(|r| &r.tier_name)
};
override_name
.or_else(rule)
.filter(|name| self.tiers.contains_key(*name))
.map_or("default", SmolStr::as_str)
}
/// 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| glob_match(r.pattern.as_bytes(), host.as_bytes()));
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::*;
use crate::util::hex;
#[test]
fn tier_names_resolve_in_the_order_tiers_do() {
let tier = |name: &str| (SmolStr::new(name), crate::config::RateTier::trusted());
let policy = TierPolicy {
tiers: [tier("default"), tier("trusted"), tier("strict")].into(),
rules: vec![
TierRule {
pattern: "*.strict.example".into(),
tier_name: "strict".into(),
},
TierRule {
pattern: "gone.example".into(),
tier_name: "unknown".into(),
},
],
};
let trusted = SmolStr::new("trusted");
let nonesuch = SmolStr::new("nonesuch");
assert_eq!(policy.resolve_name("a.strict.example", None), "strict");
assert_eq!(
policy.resolve_name("a.strict.example", Some(&trusted)),
"trusted"
);
assert_eq!(policy.resolve_name("elsewhere.example", None), "default");
assert_eq!(policy.resolve_name("gone.example", None), "default");
assert_eq!(
policy.resolve_name("a.strict.example", Some(&nonesuch)),
"default"
);
}
#[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);
}
}
}