From b918a3963815ec0f316965bd78757f5a0df24cd3 Mon Sep 17 00:00:00 2001 From: dawn <90008@gaze.systems> Date: Thu, 16 Apr 2026 05:48:58 +0300 Subject: [PATCH] [ingest,lib,api] add glob rules for pds tiers, TRUSTED_HOSTS -> TIER_RULES --- README.md | 25 ++--- src/api/xrpc/com_atproto_describe_repo.rs | 6 +- src/config.rs | 81 +++++++++----- src/control/firehose.rs | 6 ++ src/control/pds.rs | 78 +++++++++----- src/ingest/firehose.rs | 22 ++-- src/ingest/relay.rs | 50 +++------ src/pds_meta.rs | 68 +++++++++--- src/state.rs | 62 ++++++++--- tests/common.nu | 19 +++- tests/pds_status.nu | 124 +++++++++++++++++++++- 11 files changed, 395 insertions(+), 146 deletions(-) diff --git a/README.md b/README.md index bb97f09..0b284c4 100644 --- a/README.md +++ b/README.md @@ -237,8 +237,8 @@ directory, it will also be loaded automatically. | `ENABLE_CRAWLER` | `true` if full network or crawler sources are configured, `false` otherwise | whether to actively query the network for unknown repositories. | | `CRAWLER_MAX_PENDING_REPOS` | `2000` | max pending repos for crawler. | | `CRAWLER_RESUME_PENDING_REPOS` | `1000` | resume threshold for crawler pending repos. | -| `TRUSTED_HOSTS` | | comma-separated list of PDS hostnames to pre-assign to the `trusted` rate tier at startup. hosts not listed here use the `default` tier unless assigned via the API. | | `RATE_TIERS` | | comma-separated list of named rate tier definitions in `name:base/mul/hourly/daily[/account_limit]` format (e.g. `trusted:5000/10.0/18000000/432000000/10000000`). the optional account limit prevents new accounts from being created on this PDS once reached. built-in tiers (`default`, `trusted`) are always present and can be overridden. | +| `TIER_RULES` | | comma-separated ordered list of glob rules in `pattern:tier_name` format (e.g. `*.bsky.network:trusted`). rules are evaluated in order; first match wins. explicit API assignments via `PUT /pds/tiers` take precedence over rules; the `default` tier is the final fallback. uses standard glob wildcards (`*`, `?`) matched against the PDS hostname. | ## build features @@ -407,30 +407,27 @@ the built-in tiers are defined as follows: - `GET /pds/tiers`: list all current tier assignments alongside the available tier definitions. - returns `{ "assignments": [{ "host": string, "tier": string }], "rate_tiers": { : { "per_second_base": int, "per_second_account_mul": float, "per_hour": int, "per_day": int } } }`. - - `assignments` only contains PDSes with an explicit assignment; any PDS not - listed uses the `default` tier. + - `assignments` only contains PDSes with an explicit API assignment. hosts without one resolve via glob rules or fall back to `default`. - `PUT /pds/tiers`: assign a PDS to a named rate tier. - body: `{ "host": string, "tier": string }`. - `host` is the PDS hostname (e.g. `pds.example.com`). - `tier` must be one of the configured tier names. returns `400` if unknown. - assignments are persisted to the database and survive restarts. - re-assigning the same host updates the tier in place without creating a duplicate. -- `DELETE /pds/tiers`: remove an explicit tier assignment for a PDS, reverting - it to the `default` tier. +- `DELETE /pds/tiers`: remove an explicit tier assignment for a PDS. - query parameter: `?host=` (e.g. `?host=pds.example.com`). + - reverts the host to glob-rule resolution (not necessarily `default`, a matching `TIER_RULES` pattern still applies). - returns `200` even if no assignment existed. - `GET /pds/rate-tiers`: list the available rate tier definitions. - returns a map of tier name to `{ "per_second_base", "per_second_account_mul", "per_hour", "per_day", "account_limit" }`. -hosts listed in `TRUSTED_HOSTS` are seeded as `trusted` at startup, but only -when no database assignment already exists for that host — DB entries always win. -the seed is not written to the database, so it is re-applied on every restart. -consequences: if you remove a host from `TRUSTED_HOSTS` and it has no DB entry, -it will revert to `default` on the next restart. if you explicitly assign a host -via the API (which writes to the DB), that assignment persists regardless of -`TRUSTED_HOSTS`. if you delete a host's DB assignment via the API while it is -still listed in `TRUSTED_HOSTS`, it will be re-seeded as `trusted` on the next -restart. +tiers are resolved in this order: + +1. **explicit API assignment**, set via `PUT /pds/tiers`, stored in the database, survives restarts. +2. **glob rules**, from `TIER_RULES`, evaluated in order; first match wins. +3. **`default` tier**, applied if no rule or explicit assignment matches. + +deleting an API assignment reverts the host to glob-rule resolution, not necessarily back to `default`. if a rule like `*.bsky.network:trusted` matches the host, it will become trusted again without any further action. ### repository management diff --git a/src/api/xrpc/com_atproto_describe_repo.rs b/src/api/xrpc/com_atproto_describe_repo.rs index 094b1ab..b5d2f06 100644 --- a/src/api/xrpc/com_atproto_describe_repo.rs +++ b/src/api/xrpc/com_atproto_describe_repo.rs @@ -1,5 +1,7 @@ use futures::TryFutureExt; -use jacquard_api::com_atproto::repo::describe_repo::{DescribeRepoOutput, DescribeRepoRequest}; +use jacquard_api::com_atproto::repo::describe_repo::{ + DescribeRepoOutput, DescribeRepoRequest, DescribeRepoResponse, +}; use crate::util::invalid_handle; @@ -9,7 +11,7 @@ pub async fn handle( State(hydrant): State, ExtractXrpc(req): ExtractXrpc, ) -> XrpcResult>> { - let nsid = "com.atproto.repo.describeRepo"; + let nsid = DescribeRepoResponse::NSID; let resolver = &hydrant.state.resolver; let did = resolver diff --git a/src/config.rs b/src/config.rs index 4204c6a..47e1841 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,7 @@ +use crate::pds_meta::{TierPolicy, TierRule}; use miette::Result; use serde::{Deserialize, Serialize}; -use smol_str::ToSmolStr; +use smol_str::{SmolStr, ToSmolStr}; use std::collections::HashMap; use std::fmt; use std::path::PathBuf; @@ -375,16 +376,21 @@ pub struct Config { /// /// set via `HYDRANT_SEED_HOSTS` as a comma-separated list of base URLs. pub seed_hosts: Vec, - /// list of trusted PDS/relay hosts to pre-assign to the "trusted" rate tier at startup. - /// set via `HYDRANT_TRUSTED_HOSTS` as a comma-separated list of hostnames. - /// hosts not present in this list use the "default" tier unless assigned via the API. - pub trusted_hosts: Vec, /// named rate tier definitions for PDS rate limiting. /// /// built-in tiers ("default" and "trusted") are always present and may be overridden. /// set via `HYDRANT_RATE_TIERS` as a comma-separated list of `name:base/mul/hourly/daily` entries, /// e.g. `trusted:5000/10.0/18000000/432000000,custom:100/1.0/7200000/172800000`. - pub rate_tiers: HashMap, + /// + /// built from `HYDRANT_TIER_RULES` and `HYDRANT_RATE_TIERS` at startup. + pub tier_policy: TierPolicy, + + /// glob rules mapping host patterns to named rate tiers. + /// + /// set via `HYDRANT_TIER_RULES` as a comma-separated list of `pattern:tiername` entries, + /// e.g. `*.bsky.network:trusted,pds.example.com:custom`. rules are evaluated in order; + /// api-assigned per-host overrides always take priority over these rules. + pub tier_rules: Vec<(String, String)>, /// db internals, tune only if you know what you're doing. /// @@ -478,12 +484,15 @@ impl Default for Config { filter_collections: None, filter_excludes: None, enable_backlinks: false, - trusted_hosts: vec![], - rate_tiers: { - let mut m = HashMap::new(); - m.insert("default".to_string(), RateTier::default_tier()); - m.insert("trusted".to_string(), RateTier::trusted()); - m + tier_rules: vec![], + tier_policy: { + let mut tiers = HashMap::new(); + tiers.insert(SmolStr::new("default"), RateTier::default_tier()); + tiers.insert(SmolStr::new("trusted"), RateTier::trusted()); + TierPolicy { + tiers, + rules: vec![], + } }, cache_size: 256, data_compression: Compression::Zstd, @@ -651,16 +660,16 @@ impl Config { let enable_backlinks: bool = cfg!("ENABLE_BACKLINKS", defaults.enable_backlinks); - // start with built-in tiers, then layer in any env-defined overrides. + // start with built-in tier definitions, then layer in any env-defined overrides. // format: HYDRANT_RATE_TIERS=name:base/mul/hourly/daily,... - let mut rate_tiers = defaults.rate_tiers.clone(); + let mut tiers = defaults.tier_policy.tiers.clone(); if let Ok(s) = std::env::var("HYDRANT_RATE_TIERS") { for entry in s.split(',') { let entry = entry.trim(); if let Some((name, spec)) = entry.split_once(':') { match RateTier::parse(spec) { Some(tier) => { - rate_tiers.insert(name.trim().to_string(), tier); + tiers.insert(SmolStr::new(name.trim()), tier); } None => tracing::warn!( "ignoring invalid rate tier '{name}': expected base/mul/hourly/daily format" @@ -688,15 +697,35 @@ impl Config { }) .unwrap_or_else(|| defaults.seed_hosts.clone()); - let trusted_hosts = std::env::var("HYDRANT_TRUSTED_HOSTS") - .ok() - .map(|s| { - s.split(',') - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()) - .collect() - }) - .unwrap_or_else(|| defaults.trusted_hosts.clone()); + // build ordered glob rules from HYDRANT_TIER_RULES + let mut rules: Vec = vec![]; + let mut tier_rules: Vec<(String, String)> = vec![]; + if let Ok(s) = std::env::var("HYDRANT_TIER_RULES") { + for entry in s.split(',') { + let entry = entry.trim(); + if entry.is_empty() { + continue; + } + if let Some((pattern_str, tier_name)) = entry.split_once(':') { + let pattern_str = pattern_str.trim(); + let tier_name = tier_name.trim(); + match glob::Pattern::new(pattern_str) { + Ok(pattern) => { + rules.push(TierRule { + pattern, + tier_name: SmolStr::new(tier_name), + }); + tier_rules.push((pattern_str.to_string(), tier_name.to_string())); + } + Err(e) => tracing::warn!( + "ignoring invalid tier rule pattern '{pattern_str}': {e}" + ), + } + } + } + } + + let tier_policy = TierPolicy { tiers, rules }; let default_mode = CrawlerMode::default_for(full_network); let crawler_sources = match std::env::var("HYDRANT_CRAWLER_URLS") { @@ -743,8 +772,8 @@ impl Config { filter_collections, filter_excludes, enable_backlinks, - trusted_hosts, - rate_tiers, + tier_policy, + tier_rules, cache_size, data_compression, journal_compression, diff --git a/src/control/firehose.rs b/src/control/firehose.rs index 894f37b..8fd8da9 100644 --- a/src/control/firehose.rs +++ b/src/control/firehose.rs @@ -199,6 +199,12 @@ impl FirehoseHandle { let _ = self.persisted.insert_async(url.clone()).await; + // reset failure state so the fresh task gets a clean slate. + // if the previous task exited after max failures, the failure counter + // would otherwise cause the new task to exit immediately. + let throttle = self.state.throttler.get_handle(&url).await; + throttle.record_success(); + self.spawn_firehose_ingestor(&FirehoseSource { url, is_pds }, shared, false) .await?; diff --git a/src/control/pds.rs b/src/control/pds.rs index 41305b5..6813b9a 100644 --- a/src/control/pds.rs +++ b/src/control/pds.rs @@ -4,10 +4,12 @@ use std::sync::Arc; use miette::{IntoDiagnostic, Result}; use serde::Serialize; use smol_str::SmolStr; +use tracing::debug; use crate::config::RateTier; +use crate::db::keys::pds_account_count_key; use crate::db::pds_meta as db_pds; -use crate::pds_meta::{HostStatus, PdsMeta}; +use crate::pds_meta::{HostDesc, HostStatus, PdsMeta}; use crate::state::AppState; /// a single PDS-to-tier assignment. @@ -50,7 +52,7 @@ impl PdsControl { G: FnOnce(&mut PdsMeta), { let state = self.0.clone(); - tokio::task::spawn_blocking(move || { + tokio::task::spawn_blocking(move || -> Result<()> { let mut batch = state.db.inner.batch(); db_op(&mut batch, &state.db.filter); batch.commit().into_diagnostic()?; @@ -66,24 +68,21 @@ impl PdsControl { Ok(()) } - fn check_limit_transition(&self, host: &str, account_limit: Option) -> Option { - let count_key = crate::db::keys::pds_account_count_key(host); - let count = self.0.db.get_count_sync(&count_key); - let current_status = self.0.pds_meta.load().status(host); - current_status.check_limit_transition(count, account_limit) - } - - /// list all current per-PDS tier assignments. + /// list all current per-PDS tier assignments (explicit api-assigned overrides only). pub async fn list_tiers(&self) -> HashMap { let snapshot = self.0.pds_meta.load(); snapshot .hosts .iter() - .filter_map(|(host, desc)| desc.tier.as_ref().map(|t| (host.clone(), t.to_string()))) + .filter_map(|(host, desc): (&String, &HostDesc)| { + desc.tier + .as_ref() + .map(|t: &smol_str::SmolStr| (host.clone(), t.to_string())) + }) .collect() } - /// returns the assigned tier for `host`, or "default" if none is assigned. + /// returns the assigned tier for `host`, or \"default\" if none is assigned. pub fn get_tier(&self, host: impl AsRef) -> String { let snapshot = self.0.pds_meta.load(); snapshot @@ -105,7 +104,7 @@ impl PdsControl { snapshot .hosts .iter() - .filter_map(|(host, desc)| { + .filter_map(|(host, desc): (&String, &crate::pds_meta::HostDesc)| { matches!(desc.status, HostStatus::Banned).then(|| host.clone()) }) .collect() @@ -114,19 +113,22 @@ impl PdsControl { /// list all configured rate tier definitions. pub fn list_rate_tiers(&self) -> HashMap { self.0 - .rate_tiers + .tier_policy + .tiers .iter() - .map(|(name, tier)| (name.clone(), PdsTierDefinition::from(*tier))) + .map(|(name, tier): (&smol_str::SmolStr, &RateTier)| { + (name.to_string(), PdsTierDefinition::from(*tier)) + }) .collect() } - /// assign `host` to `tier`, persisting the change to the database. + /// assign `host` to `tier`. /// returns an error if `tier` is not a known tier name. pub async fn set_tier(&self, host: impl AsRef, tier: String) -> Result<()> { - if !self.0.rate_tiers.contains_key(&tier) { + if !self.0.tier_policy.tiers.contains_key(tier.as_str()) { miette::bail!( "unknown tier '{tier}'; known tiers: {:?}", - self.0.rate_tiers.keys().collect::>() + self.0.tier_policy.tiers.keys().collect::>() ); } @@ -134,8 +136,18 @@ impl PdsControl { let host_clone = host.clone(); let tier_clone = tier.clone(); - let new_tier_limit = self.0.rate_tiers.get(&tier).unwrap().account_limit; - let maybe_status = self.check_limit_transition(&host, new_tier_limit); + // read the new tier's account limit and check for a status transition, + // now that the override is about to change. + let new_tier_limit = self + .0 + .tier_policy + .tiers + .get(tier.as_str()) + .unwrap() + .account_limit; + let count = self.0.db.get_count_sync(&pds_account_count_key(&host)); + let current_status = self.0.pds_meta.load().status(&host); + let maybe_status = current_status.check_limit_transition(count, new_tier_limit); self.update( move |batch, ks| { @@ -156,17 +168,25 @@ impl PdsControl { .await } - /// remove any explicit tier assignment for `host`, reverting it to the default tier. + /// remove any explicit tier assignment for `host`, reverting it to the matched rule or default. pub async fn remove_tier(&self, host: impl AsRef) -> Result<()> { let host = host.as_ref().to_string(); let host_clone = host.clone(); - let default_tier_limit = self - .0 - .rate_tiers - .get("default") - .and_then(|t| t.account_limit); - let maybe_status = self.check_limit_transition(&host, default_tier_limit); + // after removing the override, the effective tier is determined by glob rules. + // resolve it without the override to get the correct limit. + let effective_limit = self.0.tier_policy.resolve(&host, None).account_limit; + let count = self.0.db.get_count_sync(&pds_account_count_key(&host)); + let current_status = self.0.pds_meta.load().status(&host); + let maybe_status = current_status.check_limit_transition(count, effective_limit); + debug!( + host, + ?current_status, + ?effective_limit, + count, + ?maybe_status, + "remove_tier: computed status transition" + ); self.update( move |batch, ks| { @@ -187,7 +207,7 @@ impl PdsControl { .await } - /// ban `host`, persisting the change to the database. + /// ban `host` pub async fn ban(&self, host: impl AsRef) -> Result<()> { let host = host.as_ref().to_string(); let host_clone = host.clone(); @@ -204,7 +224,7 @@ impl PdsControl { .await } - /// unban `host`, removing it from the database. + /// unban `host` pub async fn unban(&self, host: impl AsRef) -> Result<()> { let host = host.as_ref().to_string(); let host_clone = host.clone(); diff --git a/src/ingest/firehose.rs b/src/ingest/firehose.rs index 304ed0b..36a4304 100644 --- a/src/ingest/firehose.rs +++ b/src/ingest/firehose.rs @@ -130,7 +130,8 @@ impl FirehoseIngestor { if banned { break Ok(()); } - meta.tier_for(host, &self.state.rate_tiers) + let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref()); + self.state.tier_policy.resolve(host, override_name) }; let accounts = self.state.db.get_count(&count_key).await; tokio::select! { @@ -162,21 +163,30 @@ impl FirehoseIngestor { } _ = &mut active_sleep, if !marked_active => { marked_active = true; - // only reset failure state once the stream has been healthy for - // a full window — prevents hosts that connect but immediately - // send garbage from resetting their backoff on every attempt + // only reset failure state once the stream has been healthy for a bit + // so we dont get in a "connects successfully, sends garbage" situation self.throttle.record_success(); if self.is_pds { let (current_status, tier) = { let meta = self.state.pds_meta.load(); - (meta.status(host), meta.tier_for(host, &self.state.rate_tiers)) + let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref()); + (meta.status(host), self.state.tier_policy.resolve(host, override_name)) }; if current_status == HostStatus::Banned { break Ok(()); } let count = self.state.db.get_count_sync(&count_key); let new_status = tier.account_limit.is_some_and(|l| count >= l) - .then_some(HostStatus::Throttled).unwrap_or(HostStatus::Active); + .then_some(HostStatus::Throttled) + .unwrap_or(HostStatus::Active); + debug!( + host, + ?current_status, + account_limit = ?tier.account_limit, + count, + ?new_status, + "active_sleep: computed status transition" + ); if current_status != new_status { if let Err(e) = self.set_host_status(new_status) { diff --git a/src/ingest/relay.rs b/src/ingest/relay.rs index 7d2e928..8b4c5ab 100644 --- a/src/ingest/relay.rs +++ b/src/ingest/relay.rs @@ -18,6 +18,7 @@ use tokio::sync::mpsc; use tracing::{debug, error, info, info_span, trace, warn}; use url::Url; +use crate::db::keys::pds_account_count_key; use crate::db::{self, keys}; use crate::ingest::stream::AccountStatus; #[cfg(feature = "relay")] @@ -523,7 +524,7 @@ impl RelayWorker { // update per-PDS active account count on transitions if is_pds { if let Some(host) = firehose.host_str() { - let count_key = keys::pds_account_count_key(host); + let count_key = pds_account_count_key(host); let changed = if !was_active && repo_state.active { Some(ctx.state.db.update_count(&count_key, 1)) } else if was_active && !repo_state.active { @@ -533,28 +534,11 @@ impl RelayWorker { }; if let Some(count) = changed { - let (current_status, limit) = { - let meta = ctx.state.pds_meta.load(); - ( - meta.status(host), - meta.tier_for(host, &ctx.state.rate_tiers).account_limit, - ) - }; - - if let Some(status) = current_status.check_limit_transition(count, limit) { - debug!(%host, count, ?limit, ?status, "account count crossed limit, shifting status"); - if let Err(e) = crate::db::pds_meta::set_status( - &mut ctx.batch, - &ctx.state.db.filter, - host, - status, - ) { - error!(err = %e, "failed to write host status"); - } else { - crate::pds_meta::PdsMeta::update_host(&ctx.state.pds_meta, host, |h| { - h.status = status - }); - } + let mut batch_for_status = ctx.state.db.inner.batch(); + ctx.state + .apply_host_limit_status(&mut batch_for_status, host, count); + if let Err(e) = batch_for_status.commit() { + error!(%host, err = %e, "failed to commit host status update"); } } } @@ -878,20 +862,10 @@ impl WorkerContext<'_> { } if let Some(host) = msg.firehose.host_str() { - let tier = self - .state - .pds_meta - .load() - .tier_for(host, &self.state.rate_tiers); - if let Some(limit) = tier.account_limit { - let count = self - .state - .db - .get_count_sync(&crate::db::keys::pds_account_count_key(host)); - if count >= limit { - warn!(did = %did, host, count, limit, "account limit reached for host, dropping new account"); - return Ok(None); - } + let count = self.state.db.get_count_sync(&pds_account_count_key(host)); + if self.state.is_over_account_limit(host, count) { + warn!(did = %did, host, count, "account limit reached for host, dropping new account"); + return Ok(None); } } } @@ -924,7 +898,7 @@ impl WorkerContext<'_> { // track initial active state for per-PDS rate limiting if msg.is_pds && repo_state.active { if let Some(host) = msg.firehose.host_str() { - db.update_count(&keys::pds_account_count_key(host), 1); + db.update_count(&pds_account_count_key(host), 1); } } diff --git a/src/pds_meta.rs b/src/pds_meta.rs index 905dab5..1e99ad7 100644 --- a/src/pds_meta.rs +++ b/src/pds_meta.rs @@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize}; use smol_str::SmolStr; use std::collections::HashMap; use std::sync::Arc; +use tracing::debug; #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] pub enum HostStatus { @@ -82,20 +83,6 @@ impl PdsMeta { next }); } -} - -impl PdsMeta { - pub fn tier_for(&self, host: &str, rate_tiers: &HashMap) -> RateTier { - let default = rate_tiers - .get("default") - .copied() - .unwrap_or_else(RateTier::default_tier); - self.hosts - .get(host) - .and_then(|h| h.tier.as_ref()) - .and_then(|name| rate_tiers.get(name.as_str()).copied()) - .unwrap_or(default) - } pub fn status(&self, host: &str) -> HostStatus { self.hosts @@ -114,3 +101,56 @@ pub(crate) type PdsMetaHandle = Arc>; 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, + /// ordered glob rules, first match wins (for unassigned hosts). + pub rules: Vec, +} + +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_else(RateTier::default_tier); + + 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 + } +} diff --git a/src/state.rs b/src/state.rs index 75aff57..5c07e9e 100644 --- a/src/state.rs +++ b/src/state.rs @@ -11,10 +11,10 @@ use tokio::sync::watch; use url::Url; use crate::{ - config::{Config, RateTier}, + config::Config, db::Db, filter::{FilterHandle, new_handle as new_filter_handle}, - pds_meta::{PdsMeta, PdsMetaHandle, new_handle as new_pds_handle}, + pds_meta::{PdsMeta, PdsMetaHandle, TierPolicy, new_handle as new_pds_handle}, resolver::Resolver, util::throttle::Throttler, }; @@ -24,7 +24,7 @@ pub struct AppState { pub resolver: Resolver, pub(crate) filter: FilterHandle, pub(crate) pds_meta: PdsMetaHandle, - pub(crate) rate_tiers: HashMap, + pub(crate) tier_policy: TierPolicy, pub firehose_cursors: scc::HashIndex, #[cfg(feature = "indexer")] pub backfill_notify: Notify, @@ -81,14 +81,6 @@ impl AppState { .or_insert_with(crate::pds_meta::HostDesc::default) .status = stat; } - for host in &config.trusted_hosts { - let entry = hosts - .entry(host.clone()) - .or_insert_with(crate::pds_meta::HostDesc::default); - if entry.tier.is_none() { - entry.tier = Some(SmolStr::new("trusted")); - } - } let pds_meta = new_pds_handle(PdsMeta { hosts }); @@ -105,7 +97,7 @@ impl AppState { resolver, filter, pds_meta, - rate_tiers: config.rate_tiers.clone(), + tier_policy: config.tier_policy.clone(), firehose_cursors: relay_cursors, #[cfg(feature = "indexer")] backfill_notify: Notify::new(), @@ -153,4 +145,50 @@ impl AppState { result } + + /// applies an account limit status transition for `host`, writing to `batch` and updating + /// in-memory state. call this after any event that changes the active account count for a PDS. + pub(crate) fn apply_host_limit_status( + &self, + batch: &mut fjall::OwnedWriteBatch, + host: &str, + count: u64, + ) { + use crate::db::pds_meta as db_pds; + use crate::pds_meta::PdsMeta; + use tracing::{debug, error}; + + let (current_status, limit) = { + let meta = self.pds_meta.load(); + let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref()); + let tier = self.tier_policy.resolve(host, override_name); + (meta.status(host), tier.account_limit) + }; + + debug!(%host, ?current_status, ?limit, count, "apply_host_limit_status"); + + let Some(new_status) = current_status.check_limit_transition(count, limit) else { + return; + }; + + debug!(%host, count, ?limit, ?new_status, "account count crossed limit, shifting status"); + + if let Err(e) = db_pds::set_status(batch, &self.db.filter, host, new_status) { + error!(%host, err = %e, "failed to write host status"); + return; + } + + PdsMeta::update_host(&self.pds_meta, host, |h| h.status = new_status); + } + + /// checks whether `host` is at or over its account limit at the given count. + /// does not modify any state. + pub(crate) fn is_over_account_limit(&self, host: &str, count: u64) -> bool { + let meta = self.pds_meta.load(); + let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref()); + self.tier_policy + .resolve(host, override_name) + .account_limit + .is_some_and(|l| count >= l) + } } diff --git a/tests/common.nu b/tests/common.nu index 194db15..fbd6fe5 100644 --- a/tests/common.nu +++ b/tests/common.nu @@ -91,14 +91,25 @@ export def activate-account [pds_url: string, jwt: string] { curl -X POST -H "Content-Type: application/json" -H $"Authorization: Bearer ($jwt)" $"($pds_url)/xrpc/com.atproto.server.activateAccount" } +# extract the hydrant executable path from cargo's json build output +def parse-hydrant-executable [output: string] { + $output + | lines + | each { |line| try { $line | from json } catch { null } } + | compact + | where { |r| $r.reason? == "compiler-artifact" and $r.executable? != null and ($r.target?.name? == "hydrant") } + | last + | get executable +} + # build the hydrant binary export def build-hydrant [] { if ($env | get --optional HYDRANT_BINARY | is-not-empty) { return $env.HYDRANT_BINARY } print "building hydrant..." - cargo build - "target/debug/hydrant" + let out = (^cargo build --message-format json err> /dev/null | complete) + parse-hydrant-executable $out.stdout } # build the hydrant binary with extra cargo features (space-separated string) @@ -107,8 +118,8 @@ export def build-hydrant-features [features: string] { return $env.HYDRANT_BINARY } print $"building hydrant with features: ($features)..." - cargo build --features $features - "target/debug/hydrant" + let out = (^cargo build --features $features --message-format json err> /dev/null | complete) + parse-hydrant-executable $out.stdout } # start hydrant in the background diff --git a/tests/pds_status.nu b/tests/pds_status.nu index ff76b79..0c357cd 100644 --- a/tests/pds_status.nu +++ b/tests/pds_status.nu @@ -62,6 +62,12 @@ def main [] { print "starting mock pds websocket server..." let mock_pds_handle = (start-mock-pds $mock_port) + sleep 500ms + http post -t application/json $"($url)/firehose/sources" { + url: $"ws://($mock_host):($mock_port)/", + is_pds: true + } + print "checking status transitions back to Active..." mut active = false @@ -132,8 +138,124 @@ def main [] { if $re_active { print "ok: host transitioned back to active successfully." - exit 0 } else { fail "host did not transition back to active after tier removed" } + + # verify that, + # 1. the glob rule resolves the tier for unassigned hosts (no explicit PUT /pds/tiers needed) + # 2. after removing an explicit tier override, the glob rule still applies -> host stays throttled + + print "starting hydrant instance with a glob tier rule..." + let port2 = ($port + 100) + let url2 = $"http://localhost:($port2)" + let db2 = (mktemp -d -t hydrant_test.XXXXXX) + + let instance2 = (with-env { + HYDRANT_RELAY_HOSTS: "", + HYDRANT_CRAWLER_URLS: "", + HYDRANT_RATE_TIERS: "custom:1/1/1/1/0", + HYDRANT_TIER_RULES: $"127.0.0.*:custom" + } { + start-hydrant $binary $db2 $port2 + }) + if not (wait-for-api $url2) { + try { kill $instance2.pid } + fail "hydrant instance did not start" + } + + # kill any stale listener from first round + try { bash -c $"fuser -k ($mock_port)/tcp" } catch {} + sleep 100ms + + # connect mock pds and wait for offline -> active cycle (same as above) + http post -t application/json $"($url2)/firehose/sources" { + url: $"ws://($mock_host):($mock_port)/", + is_pds: true + } + + # wait for offline + print "waiting for offline..." + mut offline2 = false + for i in 1..20 { + let res = (http get -fe $"($url2)/xrpc/com.atproto.sync.getHostStatus?hostname=($mock_host)") + if $res.status == 200 and $res.body.status == "offline" { + $offline2 = true + break + } + sleep 2sec + } + if not $offline2 { + try { kill $instance2.pid } + fail "glob test: host did not go offline" + } + + print "starting mock pds for glob test..." + let mock_pds2 = (start-mock-pds $mock_port) + sleep 500ms + http post -t application/json $"($url2)/firehose/sources" { + url: $"ws://($mock_host):($mock_port)/", + is_pds: true + } + + # with account_limit=0 and the glob rule active, the host goes straight to throttled + # on the first successful connection — no explicit set_tier call needed. + print "waiting for connected (expect throttled, not active)..." + mut connected2 = false + for i in 1..20 { + let res = (http get -fe $"($url2)/xrpc/com.atproto.sync.getHostStatus?hostname=($mock_host)") + if $res.status == 200 and $res.body.status != "offline" { + $connected2 = true + print $" connected with status: ($res.body.status)" + break + } + sleep 2sec + } + if not $connected2 { + stop-mock-pds $mock_pds2 + try { kill $instance2.pid } + fail "glob test: host did not reconnect" + } + + # verify the glob rule throttled the host automatically (no set_tier was called) + print "checking glob test: glob rule throttles host without explicit tier assignment..." + let res = (http get -fe $"($url2)/xrpc/com.atproto.sync.getHostStatus?hostname=($mock_host)") + print $" status \(no set_tier\): ($res.body.status?)" + if $res.status != 200 or $res.body.status != "throttled" { + stop-mock-pds $mock_pds2 + try { kill $instance2.pid } + fail $"glob test: expected throttled without set_tier \(glob rule should apply\), got ($res.body.status?)" + } + print "ok: host throttled by glob rule without explicit tier assignment." + + # set explicit tier -> still throttled (sanity check) + print "checking glob test: explicit tier assignment also throttles..." + http put -fe -t application/json $"($url2)/pds/tiers" { host: $mock_host, tier: "custom" } + let res = (http get -fe $"($url2)/xrpc/com.atproto.sync.getHostStatus?hostname=($mock_host)") + print $" status after set_tier: ($res.body.status?)" + if $res.status != 200 or $res.body.status != "throttled" { + stop-mock-pds $mock_pds2 + try { kill $instance2.pid } + fail $"glob test: expected throttled after set_tier, got ($res.body.status?)" + } + print "ok: host throttled via explicit tier." + + # remove explicit override -> glob rule still applies -> still throttled + print "checking glob test: remove explicit tier keeps host throttled via glob rule..." + http delete -fe $"($url2)/pds/tiers?host=($mock_host)" + + sleep 500ms + let res = (http get -fe $"($url2)/xrpc/com.atproto.sync.getHostStatus?hostname=($mock_host)") + print $" status after remove_tier: ($res.body.status?)" + let still_throttled = ($res.status == 200 and $res.body.status == "throttled") + + stop-mock-pds $mock_pds2 + try { kill $instance2.pid } + + if $still_throttled { + print "ok: host remains throttled after tier override removed (glob rule applies)." + exit 0 + } else { + fail $"glob test: expected throttled after remove_tier \(glob rule should apply\), got ($res.body.status?)" + } } -- 2.51.2