diff --git a/README.md b/README.md index 7474f93..44cb401 100644 --- a/README.md +++ b/README.md @@ -171,7 +171,7 @@ directory, it will also be loaded automatically. | `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` format (e.g. `trusted:5000/10.0/18000000/432000000`). built-in tiers (`default`, `trusted`) are always present and can be overridden. | +| `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. | ## build features @@ -329,6 +329,14 @@ tiers can be defined via `RATE_TIERS`. the per-second limit scales with the number of active accounts on the PDS: `max(per_second_base, accounts × per_second_account_mul)`. +you can also define an optional `account_limit` for a rate tier. if a PDS +exceeds this number of active accounts, hydrant will reject any new account +creation events from it. + +the built-in tiers are defined as follows: +- `default`: `50` per sec (floor), `+0.5` per account. max `3_600_000`/hr, `86_400_000`/day. `100` account limit. +- `trusted`: `5000` per sec (floor), `+10.0` per account. max `18_000_000`/hr, `432_000_000`/day. `10_000_000` account limit. + - `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 } } }`. @@ -342,10 +350,10 @@ the per-second limit scales with the number of active accounts on the PDS: - 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. - - body: `{ "host": string }`. + - query parameter: `?host=` (e.g. `?host=pds.example.com`). - 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" }`. + - 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. diff --git a/flake.nix b/flake.nix index 103a3ec..98338f5 100644 --- a/flake.nix +++ b/flake.nix @@ -30,6 +30,7 @@ http-nu clang wild + psmisc ]; }; }; diff --git a/src/api/pds.rs b/src/api/pds.rs index 808ebc3..4f5b314 100644 --- a/src/api/pds.rs +++ b/src/api/pds.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use axum::{ Json, Router, - extract::State, + extract::{Query, State}, http::StatusCode, routing::{delete, get, put}, }; @@ -64,17 +64,17 @@ pub async fn set_tier( } #[derive(Deserialize)] -pub struct RemoveTierBody { +pub struct RemoveTierQuery { pub host: String, } pub async fn remove_tier( State(hydrant): State, - Json(body): Json, + Query(query): Query, ) -> Result { hydrant .pds - .remove_tier(body.host) + .remove_tier(query.host) .await .map(|_| StatusCode::OK) .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) diff --git a/src/api/xrpc/get_host_status.rs b/src/api/xrpc/get_host_status.rs index e754b39..f9590bd 100644 --- a/src/api/xrpc/get_host_status.rs +++ b/src/api/xrpc/get_host_status.rs @@ -1,8 +1,5 @@ -use jacquard_api::com_atproto::sync::{ - HostStatus, - get_host_status::{ - GetHostStatusError, GetHostStatusOutput, GetHostStatusRequest, GetHostStatusResponse, - }, +use jacquard_api::com_atproto::sync::get_host_status::{ + GetHostStatusError, GetHostStatusOutput, GetHostStatusRequest, GetHostStatusResponse, }; use jacquard_common::CowStr; @@ -26,7 +23,7 @@ pub async fn handle( account_count: Some(host.account_count as i64), hostname: CowStr::Owned(host.name), seq: Some(host.seq), - status: host.is_banned.then_some(HostStatus::Banned), + status: Some(host.status.into()), extra_data: None, })) } diff --git a/src/api/xrpc/list_hosts.rs b/src/api/xrpc/list_hosts.rs index d83e230..dbc978a 100644 --- a/src/api/xrpc/list_hosts.rs +++ b/src/api/xrpc/list_hosts.rs @@ -1,6 +1,5 @@ -use jacquard_api::com_atproto::sync::{ - HostStatus, - list_hosts::{Host, ListHostsOutput, ListHostsRequest, ListHostsResponse}, +use jacquard_api::com_atproto::sync::list_hosts::{ + Host, ListHostsOutput, ListHostsRequest, ListHostsResponse, }; use jacquard_common::CowStr; @@ -25,7 +24,7 @@ pub async fn handle( .map(|h| Host { hostname: CowStr::Owned(h.name), seq: Some(h.seq), - status: h.is_banned.then_some(HostStatus::Banned), + status: Some(h.status.into()), account_count: Some(h.account_count as i64), extra_data: None, }) diff --git a/src/config.rs b/src/config.rs index d5bbaa4..4204c6a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -22,6 +22,8 @@ pub struct RateTier { pub per_hour: u64, /// per-day limit. pub per_day: u64, + /// maximum active account limit for this host before dropping tracking of new accounts + pub account_limit: Option, } impl RateTier { @@ -32,6 +34,7 @@ impl RateTier { per_second_account_mul: 10.0, per_hour: 5000 * 3600, per_day: 5000 * 86400, + account_limit: Some(10_000_000), } } @@ -42,13 +45,14 @@ impl RateTier { per_second_account_mul: 0.5, per_hour: 1000 * 3600, per_day: 1000 * 86400, + account_limit: Some(100), } } - /// parse `base/mul/hourly/daily` format used by `HYDRANT_RATE_TIERS`. + /// parse `base/mul/hourly/daily[/account_limit]` format used by `HYDRANT_RATE_TIERS`. fn parse(s: &str) -> Option { let parts: Vec<&str> = s.split('/').collect(); - if parts.len() != 4 { + if parts.len() < 4 || parts.len() > 5 { return None; } Some(Self { @@ -56,6 +60,7 @@ impl RateTier { per_second_account_mul: parts[1].parse().ok()?, per_hour: parts[2].parse().ok()?, per_day: parts[3].parse().ok()?, + account_limit: parts.get(4).and_then(|p| p.parse().ok()), }) } } diff --git a/src/control/mod.rs b/src/control/mod.rs index e56c777..dff981b 100644 --- a/src/control/mod.rs +++ b/src/control/mod.rs @@ -56,6 +56,7 @@ use stream::event_stream_thread; #[cfg(feature = "relay")] use stream::relay_stream_thread; +#[derive(Debug, Clone)] /// infromation about a host hydrant is consuming from. pub struct Host { /// hostname of the host. @@ -64,8 +65,8 @@ pub struct Host { pub seq: i64, /// the amount of accounts hydrant has seen from this host. pub account_count: u64, - /// whether this host is banned or not. - pub is_banned: bool, + /// the status of this host in hydrant. + pub status: crate::pds_meta::HostStatus, } /// an event emitted by the hydrant event stream. @@ -767,25 +768,40 @@ impl Hydrant { tokio::task::spawn_blocking(move || { let key = keys::firehose_cursor_key(&hostname); - let Some(seq) = state.db.cursors.get(&key).into_diagnostic()? else { - return Ok(None); - }; - let seq = i64::from_be_bytes( - seq.as_ref() - .try_into() - .into_diagnostic() - .wrap_err("cursor value is not 8 bytes")?, - ); + + let mut seq = 0; + if let Some(cursor_bytes) = state.db.cursors.get(&key).into_diagnostic()? { + seq = i64::from_be_bytes(cursor_bytes.as_ref().try_into().into_diagnostic()?); + } else { + // if it has no cursor, check if it's explicitly tracked in hosts map + // or firehose tasks (recently added via API but no messages yet) + let meta = state.pds_meta.load(); + if !meta.hosts.contains_key(hostname.as_str()) { + // we should also allow it if it's an active firehose ingestor + let mut found_in_cursors = false; + state.firehose_cursors.iter_sync(|u, _| { + if u.host_str() == Some(hostname.as_str()) { + found_in_cursors = true; + } + !found_in_cursors // continue if not found + }); + + if !found_in_cursors { + return Ok(None); + } + } + } + let account_count = state .db .get_count_sync(&keys::pds_account_count_key(&hostname)); - let is_banned = state.pds_meta.load().is_banned(&hostname); + let status = state.pds_meta.load().status(&hostname); Ok(Some(Host { name: hostname.into(), seq, account_count, - is_banned, + status, })) }) .await @@ -835,12 +851,12 @@ impl Hydrant { let account_count = state .db .get_count_sync(&keys::pds_account_count_key(hostname)); - let is_banned = state.pds_meta.load().is_banned(&hostname); + let status = state.pds_meta.load().status(hostname); hosts.push(Host { name: hostname.into(), seq, account_count, - is_banned, + status, }); } diff --git a/src/control/pds.rs b/src/control/pds.rs index 9b7558f..41305b5 100644 --- a/src/control/pds.rs +++ b/src/control/pds.rs @@ -7,7 +7,7 @@ use smol_str::SmolStr; use crate::config::RateTier; use crate::db::pds_meta as db_pds; -use crate::pds_meta::PdsMeta; +use crate::pds_meta::{HostStatus, PdsMeta}; use crate::state::AppState; /// a single PDS-to-tier assignment. @@ -24,6 +24,7 @@ pub struct PdsTierDefinition { pub per_second_account_mul: f64, pub per_hour: u64, pub per_day: u64, + pub account_limit: Option, } impl From for PdsTierDefinition { @@ -33,6 +34,7 @@ impl From for PdsTierDefinition { per_second_account_mul: t.per_second_account_mul, per_hour: t.per_hour, per_day: t.per_day, + account_limit: t.account_limit, } } } @@ -64,13 +66,20 @@ 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. pub async fn list_tiers(&self) -> HashMap { let snapshot = self.0.pds_meta.load(); snapshot - .tiers + .hosts .iter() - .map(|(host, tier)| (host.clone(), tier.to_string())) + .filter_map(|(host, desc)| desc.tier.as_ref().map(|t| (host.clone(), t.to_string()))) .collect() } @@ -78,8 +87,9 @@ impl PdsControl { pub fn get_tier(&self, host: impl AsRef) -> String { let snapshot = self.0.pds_meta.load(); snapshot - .tiers + .hosts .get(host.as_ref()) + .and_then(|h| h.tier.as_ref()) .map(|t| t.to_string()) .unwrap_or_else(|| "default".to_string()) } @@ -92,7 +102,13 @@ impl PdsControl { /// list all currently banned PDS hosts. pub async fn list_banned(&self) -> Vec { let snapshot = self.0.pds_meta.load(); - snapshot.banned.iter().cloned().collect() + snapshot + .hosts + .iter() + .filter_map(|(host, desc)| { + matches!(desc.status, HostStatus::Banned).then(|| host.clone()) + }) + .collect() } /// list all configured rate tier definitions. @@ -117,10 +133,24 @@ impl PdsControl { let host = host.as_ref().to_string(); 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); + self.update( - move |batch, ks| db_pds::set_tier(batch, ks, &host_clone, &tier_clone), + move |batch, ks| { + let _ = db_pds::set_tier(batch, ks, &host_clone, &tier_clone); + if let Some(status) = maybe_status { + let _ = db_pds::set_status(batch, ks, &host_clone, status); + } + }, move |meta| { - meta.tiers.insert(host, SmolStr::new(&tier)); + meta.update_host_entry(&host, |entry| { + entry.tier = Some(SmolStr::new(&tier)); + if let Some(status) = maybe_status { + entry.status = status; + } + }); }, ) .await @@ -130,10 +160,28 @@ impl PdsControl { 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); + self.update( - move |batch, ks| db_pds::remove_tier(batch, ks, &host_clone), + move |batch, ks| { + let _ = db_pds::remove_tier(batch, ks, &host_clone); + if let Some(status) = maybe_status { + let _ = db_pds::set_status(batch, ks, &host_clone, status); + } + }, move |meta| { - meta.tiers.remove(&host); + meta.update_host_entry(&host, |desc| { + desc.tier = None; + if let Some(status) = maybe_status { + desc.status = status; + } + }); }, ) .await @@ -144,9 +192,13 @@ impl PdsControl { let host = host.as_ref().to_string(); let host_clone = host.clone(); self.update( - move |batch, ks| db_pds::set_banned(batch, ks, &host_clone), + move |batch, ks| { + let _ = db_pds::set_status(batch, ks, &host_clone, HostStatus::Banned); + }, move |meta| { - meta.banned.insert(host); + meta.update_host_entry(&host, |desc| { + desc.status = HostStatus::Banned; + }); }, ) .await @@ -157,9 +209,11 @@ impl PdsControl { let host = host.as_ref().to_string(); let host_clone = host.clone(); self.update( - move |batch, ks| db_pds::remove_banned(batch, ks, &host_clone), + move |batch, ks| db_pds::remove_status(batch, ks, &host_clone), move |meta| { - meta.banned.remove(&host); + meta.update_host_entry(&host, |desc| { + desc.status = HostStatus::Active; + }); }, ) .await diff --git a/src/crawler/list_repos.rs b/src/crawler/list_repos.rs index 62e5659..8c3fca8 100644 --- a/src/crawler/list_repos.rs +++ b/src/crawler/list_repos.rs @@ -234,8 +234,8 @@ impl SignalChecker { return (did, retry_state.into()); } if is_throttle_worthy(&e) { - if let Some(mins) = throttle.record_failure() { - warn!(url = %pds_url, mins, "throttling pds due to hard failure"); + if let Some(secs) = throttle.record_failure() { + warn!(url = %pds_url, secs, "throttling pds due to hard failure"); } let mut retry_state = throttle.to_retry_state(); retry_state.status = e.status(); diff --git a/src/db/migration/mod.rs b/src/db/migration/mod.rs index cbdd43a..a188b36 100644 --- a/src/db/migration/mod.rs +++ b/src/db/migration/mod.rs @@ -8,6 +8,7 @@ mod v1; mod v2; mod v3; mod v4; +mod v5; type MigrationFn = fn(&Db, &mut OwnedWriteBatch) -> Result<()>; @@ -17,6 +18,7 @@ const MIGRATIONS: &[(&str, MigrationFn)] = &[ ("repo_state_root_commit", v2::repo_state_root_commit), ("firehose_source_is_pds", v3::firehose_source_is_pds), ("repo_state_active", v4::repo_state_active), + ("pds_meta_layout", v5::pds_meta_layout), ]; fn read_version(db: &Db) -> Result { diff --git a/src/db/migration/v5.rs b/src/db/migration/v5.rs new file mode 100644 index 0000000..a506a91 --- /dev/null +++ b/src/db/migration/v5.rs @@ -0,0 +1,38 @@ +use crate::db::Db; +use fjall::OwnedWriteBatch; +use miette::{Context, IntoDiagnostic, Result}; + +pub mod v4 { + pub const PDS_TIER_PREFIX: &[u8] = b"pt|"; + pub const PDS_BANNED_PREFIX: &[u8] = b"pb|"; +} + +pub(crate) fn pds_meta_layout(db: &Db, batch: &mut OwnedWriteBatch) -> Result<()> { + for guard in db.filter.prefix(v4::PDS_TIER_PREFIX) { + let (k, v) = guard.into_inner().into_diagnostic()?; + let host = std::str::from_utf8(&k[v4::PDS_TIER_PREFIX.len()..]) + .into_diagnostic() + .wrap_err("failed to parse host as utf8")?; + let tier = std::str::from_utf8(&v) + .into_diagnostic() + .wrap_err("failed to parse tier as utf8")?; + crate::db::pds_meta::set_tier(batch, &db.filter, host, tier); + batch.remove(&db.filter, k); + } + + for guard in db.filter.prefix(v4::PDS_BANNED_PREFIX) { + let (k, _) = guard.into_inner().into_diagnostic()?; + let host = std::str::from_utf8(&k[v4::PDS_BANNED_PREFIX.len()..]) + .into_diagnostic() + .wrap_err("failed to parse host as utf8")?; + crate::db::pds_meta::set_status( + batch, + &db.filter, + host, + crate::pds_meta::HostStatus::Banned, + )?; + batch.remove(&db.filter, k); + } + + Ok(()) +} diff --git a/src/db/mod.rs b/src/db/mod.rs index 76b6976..0016186 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -85,8 +85,8 @@ macro_rules! update_gauge_diff_impl { // pending match ($old, $new) { (GaugeState::Pending, GaugeState::Pending) => {} - (GaugeState::Pending, _) => $self.$update_method("pending", -1) $(.$await)?, - (_, GaugeState::Pending) => $self.$update_method("pending", 1) $(.$await)?, + (GaugeState::Pending, _) => {$self.$update_method("pending", -1) $(.$await)?;}, + (_, GaugeState::Pending) => {$self.$update_method("pending", 1) $(.$await)?;}, _ => {} } @@ -94,8 +94,8 @@ macro_rules! update_gauge_diff_impl { let old_resync = $old.is_resync(); let new_resync = $new.is_resync(); match (old_resync, new_resync) { - (true, false) => $self.$update_method("resync", -1) $(.$await)?, - (false, true) => $self.$update_method("resync", 1) $(.$await)?, + (true, false) => {$self.$update_method("resync", -1) $(.$await)?;}, + (false, true) => {$self.$update_method("resync", 1) $(.$await)?;}, _ => {} } @@ -697,7 +697,7 @@ impl Db { .into_diagnostic()? } - pub fn update_count(&self, key: &str, delta: i64) { + pub fn update_count(&self, key: &str, delta: i64) -> u64 { let mut entry = self.counts_map.entry_sync(SmolStr::new(key)).or_insert(0); if delta >= 0 { *entry = entry.saturating_add(delta as u64); @@ -715,6 +715,7 @@ impl Db { *entry -= decrement; } } + *entry } pub async fn update_count_async(&self, key: &str, delta: i64) { diff --git a/src/db/pds_meta.rs b/src/db/pds_meta.rs index 7e9489c..4c2593b 100644 --- a/src/db/pds_meta.rs +++ b/src/db/pds_meta.rs @@ -1,62 +1,77 @@ +use crate::pds_meta::HostStatus; use fjall::{Keyspace, OwnedWriteBatch}; use miette::{IntoDiagnostic, Result}; use smol_str::SmolStr; -pub const PDS_TIER_PREFIX: &[u8] = b"pt|"; +pub mod v5 { + use super::*; -// `pt|{host}` -> tier name -pub fn pds_tier_key(host: &str) -> Vec { - let mut key = Vec::with_capacity(PDS_TIER_PREFIX.len() + host.len()); - key.extend_from_slice(PDS_TIER_PREFIX); - key.extend_from_slice(host.as_bytes()); - key -} + // `{host}|tier` -> tier name + pub fn pds_tier_key(host: &str) -> Vec { + let mut key = Vec::with_capacity(host.len() + 5); + key.extend_from_slice(host.as_bytes()); + key.extend_from_slice(b"|tier"); + key + } -/// load all PDS tier assignments from the filter keyspace -pub fn load_tiers(ks: &Keyspace) -> Result> { - let mut out = Vec::new(); - for guard in ks.prefix(PDS_TIER_PREFIX) { - let (k, v) = guard.into_inner().into_diagnostic()?; - let host = std::str::from_utf8(&k[PDS_TIER_PREFIX.len()..]).into_diagnostic()?; - let tier = std::str::from_utf8(&v).into_diagnostic()?; - out.push((SmolStr::new(host), SmolStr::new(tier))); + /// load all PDS tier assignments from the filter keyspace + pub fn load_tiers(ks: &Keyspace) -> Result> { + let mut out = Vec::new(); + for guard in ks.iter() { + let (k, v) = guard.into_inner().into_diagnostic()?; + if k.ends_with(b"|tier") { + let host = std::str::from_utf8(&k[..k.len() - 5]).into_diagnostic()?; + let tier = std::str::from_utf8(&v).into_diagnostic()?; + out.push((SmolStr::new(host), SmolStr::new(tier))); + } + } + Ok(out) } - Ok(out) -} -pub fn set_tier(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str, tier: &str) { - batch.insert(ks, pds_tier_key(host), tier.as_bytes()); -} + pub fn set_tier(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str, tier: &str) { + batch.insert(ks, pds_tier_key(host), tier.as_bytes()); + } -pub fn remove_tier(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str) { - batch.remove(ks, pds_tier_key(host)); -} + pub fn remove_tier(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str) { + batch.remove(ks, pds_tier_key(host)); + } -pub const PDS_BANNED_PREFIX: &[u8] = b"pb|"; + // `{host}|status` -> encoded HostStatus (msgpack) + pub fn pds_status_key(host: &str) -> Vec { + let mut key = Vec::with_capacity(host.len() + 7); + key.extend_from_slice(host.as_bytes()); + key.extend_from_slice(b"|status"); + key + } -// `pb|{host}` -> empty value -pub fn pds_banned_key(host: &str) -> Vec { - let mut key = Vec::with_capacity(PDS_BANNED_PREFIX.len() + host.len()); - key.extend_from_slice(PDS_BANNED_PREFIX); - key.extend_from_slice(host.as_bytes()); - key -} + /// load all host statuses from the filter keyspace + pub fn load_statuses(ks: &Keyspace) -> Result> { + let mut out = Vec::new(); + for guard in ks.iter() { + let (k, v) = guard.into_inner().into_diagnostic()?; + if k.ends_with(b"|status") { + let host = std::str::from_utf8(&k[..k.len() - 7]).into_diagnostic()?; + let status: HostStatus = rmp_serde::from_slice(&v).into_diagnostic()?; + out.push((SmolStr::new(host), status)); + } + } + Ok(out) + } -/// load all banned PDS hosts from the filter keyspace -pub fn load_banned(ks: &Keyspace) -> Result> { - let mut out = Vec::new(); - for guard in ks.prefix(PDS_BANNED_PREFIX) { - let (k, _) = guard.into_inner().into_diagnostic()?; - let host = std::str::from_utf8(&k[PDS_BANNED_PREFIX.len()..]).into_diagnostic()?; - out.push(SmolStr::new(host)); + pub fn set_status( + batch: &mut OwnedWriteBatch, + ks: &Keyspace, + host: &str, + status: HostStatus, + ) -> Result<()> { + let bytes = rmp_serde::to_vec(&status).into_diagnostic()?; + batch.insert(ks, pds_status_key(host), bytes); + Ok(()) } - Ok(out) -} -pub fn set_banned(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str) { - batch.insert(ks, pds_banned_key(host), &[]); + pub fn remove_status(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str) { + batch.remove(ks, pds_status_key(host)); + } } -pub fn remove_banned(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str) { - batch.remove(ks, pds_banned_key(host)); -} +pub use v5::*; diff --git a/src/ingest/firehose.rs b/src/ingest/firehose.rs index 4bdecf3..30e3ee6 100644 --- a/src/ingest/firehose.rs +++ b/src/ingest/firehose.rs @@ -1,6 +1,7 @@ use crate::filter::{FilterHandle, FilterMode}; use crate::ingest::stream::{FirehoseError, FirehoseStream, SubscribeReposMessage, decode_frame}; use crate::ingest::{BufferTx, IngestMessage}; +use crate::pds_meta::HostStatus; use crate::state::AppState; use crate::util::throttle::ThrottleHandle; use crate::util::{ @@ -142,8 +143,15 @@ impl FirehoseIngestor { || matches!(&e, FirehoseError::EmptyFrame); let timeout = if do_throttle { self.throttle.record_failure(); + if self.is_pds && self.throttle.consecutive_failures() >= 4 { + if let Err(e) = self.set_host_status(HostStatus::Offline) { + error!(err = %e, "failed to update host status to offline"); + } + } let until = self.throttle.throttled_until(); - Duration::from_secs((until - chrono::Utc::now().timestamp()) as u64) + Duration::from_secs( + 0.max((until - chrono::Utc::now().timestamp()) as i64) as u64 + ) } else { Duration::from_secs(10) }; @@ -157,7 +165,10 @@ impl FirehoseIngestor { self.throttle.record_success(); info!("firehose connected"); - let connected_at = tokio::time::Instant::now(); + let mut marked_active = false; + let active_sleep_secs = if cfg!(debug_assertions) { 1 } else { 60 }; + let mut active_sleep = + std::pin::pin!(tokio::time::sleep(Duration::from_secs(active_sleep_secs))); let res = loop { tokio::select! { @@ -188,7 +199,7 @@ impl FirehoseIngestor { } } } - self.handle_message(msg).await + self.handle_message(msg).await; }, Err(e) => match e { // dont disconnect on unknown op or type @@ -204,8 +215,26 @@ impl FirehoseIngestor { e => break Err(e), }, } - if connected_at.elapsed() > Duration::from_secs(60) { - backoff = Duration::from_secs(0); + } + _ = &mut active_sleep, if !marked_active => { + marked_active = true; + backoff = Duration::from_secs(0); + 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)) + }; + if current_status != HostStatus::Banned { + 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); + + if current_status != new_status { + if let Err(e) = self.set_host_status(new_status) { + error!(err = %e, "failed to update host status"); + } + } + } } } _ = self.enabled.changed() => { @@ -218,20 +247,34 @@ impl FirehoseIngestor { }; if let Err(e) = res { - if let FirehoseError::StreamClosed { code, reason } = &e - && *code == 1001 - { - debug!(reason = %reason, "host gone away"); - tokio::time::sleep(Duration::from_secs(1)).await; - continue; - } - if let FirehoseError::RelayError { error, message } = e { - let message = message.map_or(Cow::Borrowed(""), Cow::Owned); - error!(err = %error, "relay sent error: {message}"); - } else if backoff.as_secs() < 60 { - // stop logging errors after a minute of retries - // as to not spam logs, unlikely for error to change atp - error!(err = %e, "firehose stream error"); + match &e { + FirehoseError::StreamClosed { code: 1001, reason } => { + debug!(reason = %reason, "host gone away"); + tokio::time::sleep(Duration::from_secs(1)).await; + continue; + } + FirehoseError::FutureCursor => { + if self.is_pds + && let Err(e) = self.set_host_status(HostStatus::Idle) + { + error!(err = %e, "failed to update host status to idle"); + } + debug!("outdated cursor, backing off"); + tokio::time::sleep(Duration::from_secs(60)).await; + continue; + } + FirehoseError::RelayError { error, message } => { + let message = message + .as_deref() + .map_or(Cow::Borrowed(""), Cow::Borrowed); + error!(err = %error, "relay sent error: {message}"); + } + _ if backoff.as_secs() < 60 => { + // stop logging errors after a minute of retries + // as to not spam logs, unlikely for error to change atp + error!(err = %e, "firehose stream error"); + } + _ => {} } if backoff.is_zero() { backoff = Duration::from_secs(5); @@ -245,6 +288,22 @@ impl FirehoseIngestor { } } + fn set_host_status(&self, status: HostStatus) -> Result<()> { + let Some(host) = self.relay_host.host_str() else { + return Ok(()); + }; + + debug!(host = %host, status = ?status, "updating host status"); + + let mut batch = self.state.db.inner.batch(); + crate::db::pds_meta::set_status(&mut batch, &self.state.db.filter, host, status)?; + batch.commit().into_diagnostic()?; + + crate::pds_meta::PdsMeta::update_host(&self.state.pds_meta, host, |h| h.status = status); + + Ok(()) + } + async fn handle_message(&self, msg: SubscribeReposMessage<'_>) { let did = match &msg { SubscribeReposMessage::Commit(commit) => &commit.repo, diff --git a/src/ingest/relay.rs b/src/ingest/relay.rs index ffc33cb..7d2e928 100644 --- a/src/ingest/relay.rs +++ b/src/ingest/relay.rs @@ -524,10 +524,38 @@ impl RelayWorker { if is_pds { if let Some(host) = firehose.host_str() { let count_key = keys::pds_account_count_key(host); - if !was_active && repo_state.active { - ctx.state.db.update_count(&count_key, 1); + let changed = if !was_active && repo_state.active { + Some(ctx.state.db.update_count(&count_key, 1)) } else if was_active && !repo_state.active { - ctx.state.db.update_count(&count_key, -1); + Some(ctx.state.db.update_count(&count_key, -1)) + } else { + None + }; + + 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 + }); + } + } } } } @@ -848,6 +876,24 @@ impl WorkerContext<'_> { warn!(did = %did, got = ?pds_host, expected = ?msg.firehose.host_str(), "message rejected: wrong host for new account"); return Ok(None); } + + 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); + } + } + } } // try to get upstream status diff --git a/src/ingest/stream.rs b/src/ingest/stream.rs index 7491d86..c18db50 100644 --- a/src/ingest/stream.rs +++ b/src/ingest/stream.rs @@ -47,6 +47,8 @@ pub enum FirehoseError { StreamClosed { code: u16, reason: String }, #[error("tcp layer dropped")] TcpDropped, + #[error("future cursor")] + FutureCursor, } impl From> for FirehoseError { @@ -612,7 +614,13 @@ pub fn decode_frame<'i>(bytes: &'i [u8]) -> Result, Fi SubscribeReposMessage::Identity(Box::new(Deserialize::deserialize(&mut de)?)) } "#sync" => SubscribeReposMessage::Sync(Box::new(Deserialize::deserialize(&mut de)?)), - "#info" => SubscribeReposMessage::Info(Box::new(Deserialize::deserialize(&mut de)?)), + "#info" => { + let info: crate::ingest::stream::Info<'i> = Deserialize::deserialize(&mut de)?; + if info.name == InfoName::OutdatedCursor { + return Err(FirehoseError::FutureCursor); + } + SubscribeReposMessage::Info(Box::new(info)) + } other => return Err(FirehoseError::UnknownType(other.to_string())), }; diff --git a/src/pds_meta.rs b/src/pds_meta.rs index 44281b1..905dab5 100644 --- a/src/pds_meta.rs +++ b/src/pds_meta.rs @@ -1,13 +1,87 @@ use crate::config::RateTier; use arc_swap::ArcSwap; +use serde::{Deserialize, Serialize}; use smol_str::SmolStr; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum HostStatus { + Active, + Idle, + Offline, + Throttled, + Banned, +} + +impl From for jacquard_api::com_atproto::sync::HostStatus<'static> { + 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 { + /// returns the new status to apply if the count dynamically crossed limits. + pub fn check_limit_transition( + self, + current_count: u64, + account_limit: Option, + ) -> Option { + if self == Self::Banned { + return None; + } + 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, + } + } +} + +impl Default for HostStatus { + fn default() -> Self { + Self::Active + } +} + +#[derive(Debug, Default, Clone)] +pub struct HostDesc { + pub tier: Option, + pub status: HostStatus, +} + #[derive(Default, Clone)] pub(crate) struct PdsMeta { - pub tiers: HashMap, - pub banned: HashSet, + pub hosts: HashMap, +} + +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, + host: &str, + mut f: impl FnMut(&mut HostDesc), + ) { + cell.rcu(|meta| { + let mut next = (**meta).clone(); + next.update_host_entry(host, &mut f); + next + }); + } } impl PdsMeta { @@ -16,14 +90,22 @@ impl PdsMeta { .get("default") .copied() .unwrap_or_else(RateTier::default_tier); - self.tiers + 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 + .get(host) + .map(|h| h.status) + .unwrap_or(HostStatus::Active) + } + pub fn is_banned(&self, host: &str) -> bool { - self.banned.contains(host) + self.status(host) == HostStatus::Banned } } diff --git a/src/state.rs b/src/state.rs index 6f56dfb..75aff57 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::future::Future; use std::sync::atomic::AtomicI64; use std::time::Duration; @@ -55,27 +55,42 @@ impl AppState { let filter = new_filter_handle(filter_config); - // load persisted per-PDS tier assignments from the filter keyspace. - // trusted_hosts from config are merged in as defaults (not persisted here; they seed - // only if the host has no existing assignment in the DB). - let mut tiers: HashMap = crate::db::pds_meta::load_tiers(&db.filter) + let tiers: HashMap = crate::db::pds_meta::load_tiers(&db.filter) .unwrap_or_default() .into_iter() .map(|(host, tier)| (host.to_string(), tier)) .collect(); + + let statuses: HashMap = + crate::db::pds_meta::load_statuses(&db.filter) + .unwrap_or_default() + .into_iter() + .map(|(host, stat)| (host.to_string(), stat)) + .collect(); + + let mut hosts = HashMap::new(); + for (host, tier) in tiers { + hosts + .entry(host) + .or_insert_with(crate::pds_meta::HostDesc::default) + .tier = Some(tier); + } + for (host, stat) in statuses { + hosts + .entry(host) + .or_insert_with(crate::pds_meta::HostDesc::default) + .status = stat; + } for host in &config.trusted_hosts { - tiers + let entry = hosts .entry(host.clone()) - .or_insert_with(|| SmolStr::new("trusted")); + .or_insert_with(crate::pds_meta::HostDesc::default); + if entry.tier.is_none() { + entry.tier = Some(SmolStr::new("trusted")); + } } - let banned: HashSet = crate::db::pds_meta::load_banned(&db.filter) - .unwrap_or_default() - .into_iter() - .map(|host| host.to_string()) - .collect(); - - let pds_meta = new_pds_handle(PdsMeta { tiers, banned }); + let pds_meta = new_pds_handle(PdsMeta { hosts }); let relay_cursors = scc::HashIndex::new(); diff --git a/src/util/throttle.rs b/src/util/throttle.rs index 0bf64aa..b13455d 100644 --- a/src/util/throttle.rs +++ b/src/util/throttle.rs @@ -112,29 +112,31 @@ impl ThrottleHandle { } /// called on hard failures (timeout, TLS error, bad gateway, etc). - /// returns throttle duration in minutes if this is a *new* throttle, - /// and notifies all in-flight tasks to cancel immediately. + /// always increments `consecutive_failures`. only sets a new `throttled_until` + /// (and notifies waiters) if not already throttled. pub fn record_failure(&self) -> Option { - if self.is_throttled() { - return None; - } - let failures = self .state .consecutive_failures .fetch_add(1, Ordering::AcqRel) + 1; - // 30 min, 60 min, 120 min, ... capped at ~512 hours - let base_minutes = 30u64; + if self.is_throttled() { + return None; + } + + let base_secs = 15u64; let exponent = (failures as u32).saturating_sub(1); - let minutes = base_minutes * 2u64.pow(exponent.min(10)); - let until = chrono::Utc::now().timestamp() + (minutes * 60) as i64; + let secs = (base_secs * 2u64.pow(exponent.min(10))).min(300); + #[cfg(debug_assertions)] + let secs = secs.min(1); + + let until = chrono::Utc::now().timestamp() + secs as i64; self.state.throttled_until.store(until, Ordering::Release); self.state.failure_notify.notify_waiters(); - Some(minutes) + Some(secs) } /// returns current timeout duration — 3s, 6s, or 12s depending on prior timeouts. @@ -143,6 +145,10 @@ impl ThrottleHandle { Duration::from_secs(3 * 2u64.pow(n.min(2) as u32)) } + pub fn consecutive_failures(&self) -> usize { + self.state.consecutive_failures.load(Ordering::Acquire) as usize + } + /// returns whether the timeout attempts are exhausted pub fn record_timeout(&self) -> bool { let timeouts = self diff --git a/tests/api.nu b/tests/api.nu index 07865d9..f90ff0c 100644 --- a/tests/api.nu +++ b/tests/api.nu @@ -272,7 +272,7 @@ def test-pds-tiers [url: string, pid: int] { let rate_tiers = (http get $"($url)/pds/rate-tiers") for tier_name in ["default", "trusted"] { let tier = ($rate_tiers | get $tier_name) - for field in ["per_second_base", "per_second_account_mul", "per_hour", "per_day"] { + for field in ["per_second_base", "per_second_account_mul", "per_hour", "per_day", "account_limit"] { if not ($field in $tier) { fail $"($tier_name) tier missing field ($field)" $pid } @@ -346,9 +346,7 @@ def test-pds-tiers [url: string, pid: int] { # remove the first host print " DELETE /pds/tiers (first host)..." - http delete -f -e -t application/json $"($url)/pds/tiers" --data { - host: "pds.example.com" - } | assert-status 200 "DELETE /pds/tiers" $pid + http delete -f -e $"($url)/pds/tiers?host=pds.example.com" | assert-status 200 "DELETE /pds/tiers" $pid let after_del = (http get $"($url)/pds/tiers") if ($after_del.assignments | columns | length) != 1 { fail $"expected 1 assignment after delete, got ($after_del.assignments | columns | length)" $pid @@ -359,15 +357,11 @@ def test-pds-tiers [url: string, pid: int] { print " ok: correct host removed, other assignment intact" # remove the second host - http delete -f -e -t application/json $"($url)/pds/tiers" --data { - host: "other.example.com" - } | assert-status 200 "DELETE /pds/tiers second" $pid + http delete -f -e $"($url)/pds/tiers?host=other.example.com" | assert-status 200 "DELETE /pds/tiers second" $pid # deleting a non-existent host is idempotent (returns 200, not an error) print " DELETE /pds/tiers (non-existent, expect 200)..." - http delete -f -e -t application/json $"($url)/pds/tiers" --data { - host: "pds.example.com" - } | assert-status 200 "DELETE /pds/tiers non-existent" $pid + http delete -f -e $"($url)/pds/tiers?host=pds.example.com" | assert-status 200 "DELETE /pds/tiers non-existent" $pid let after_idempotent = (http get $"($url)/pds/tiers") if ($after_idempotent.assignments | columns | length) != 0 { fail "expected empty assignments after cleanup" $pid diff --git a/tests/mock_pds.nu b/tests/mock_pds.nu new file mode 100644 index 0000000..7b8929a --- /dev/null +++ b/tests/mock_pds.nu @@ -0,0 +1,12 @@ +export def start-mock-pds [port: int] { + # kill any stale process from a previous failed run holding this port + try { bash -c $"fuser -k ($port)/tcp" } catch {} + sleep 100ms + let log_file = (mktemp) + let pid = (bash -c $"websocat -s ($port) >($log_file) 2>&1 & echo $!" | str trim | into int) + { pid: $pid, log: $log_file } +} + +export def stop-mock-pds [handle: record] { + try { kill $handle.pid } +} diff --git a/tests/pds_status.nu b/tests/pds_status.nu new file mode 100644 index 0000000..ff76b79 --- /dev/null +++ b/tests/pds_status.nu @@ -0,0 +1,139 @@ +source common.nu + +source mock_pds.nu + +def main [] { + let port = resolve-test-port 3033 + let url = $"http://localhost:($port)" + let binary = build-hydrant + let db = (mktemp -d -t hydrant_test.XXXXXX) + + let instance = (with-env { + HYDRANT_RELAY_HOSTS: "", + HYDRANT_CRAWLER_URLS: "", + HYDRANT_RATE_TIERS: "custom:1/1/1/1/0" + } { + start-hydrant $binary $db $port + }) + if not (wait-for-api $url) { + fail "hydrant did not start" $instance.pid + } + + let mock_port = resolve-test-mock-port 9999 + let mock_host = "127.0.0.1" + + # kill any stale listener on the mock port from a previous failed run + try { bash -c $"fuser -k ($mock_port)/tcp" } catch {} + sleep 100ms + + print "adding offline mock pds via firehose sources..." + http post -t application/json $"($url)/firehose/sources" { + url: $"ws://($mock_host):($mock_port)/", + is_pds: true + } + + print "checking status transitions to Offline..." + mut offline = false + + # the throttle backoff will cap at 1 second in debug builds. + # it takes 4 consecutive failures to mark as offline. + # therefore, 4 * 1 = ~4 seconds maximum for transition. + for i in 1..20 { + let res = (http get -fe $"($url)/xrpc/com.atproto.sync.getHostStatus?hostname=($mock_host)") + if $res.status == 200 { + if $res.body.status == "offline" { + $offline = true + break + } + if $res.body.status == "active" { + print $" ... currently ($res.body.status), waiting for offline" + } + } else { + print $" ... could not get status, waiting: ($res.status)" + } + sleep 2sec + } + + if not $offline { + fail "host did not transition to offline within time limit" $instance.pid + } + print "ok: host transitioned to offline successfully." + + print "starting mock pds websocket server..." + let mock_pds_handle = (start-mock-pds $mock_port) + + print "checking status transitions back to Active..." + mut active = false + + # now wait for it to successfully reconnect and the active_sleep of 1s to pass. + for i in 1..20 { + let res = (http get -fe $"($url)/xrpc/com.atproto.sync.getHostStatus?hostname=($mock_host)") + if $res.status == 200 { + if $res.body.status == "active" { + $active = true + break + } + if $res.body.status == "offline" { + print $" ... currently ($res.body.status), waiting for active" + } + } else { + print $" ... could not get status, waiting: ($res.status)" + } + sleep 2sec + } + + if $active { + print "ok: host transitioned to active successfully." + } else { + stop-mock-pds $mock_pds_handle + try { kill $instance.pid } + fail "host did not transition to active within time limit" + } + + print "checking status transitions to Throttled..." + let put_res = (http put -fe -t application/json $"($url)/pds/tiers" { + host: $mock_host, + tier: "custom" + }) + if $put_res.status != 200 { + print $"PUT /pds/tiers failed with status ($put_res.status)" + print $put_res.body + stop-mock-pds $mock_pds_handle + try { kill $instance.pid } + fail "failed to change tier" + } + + # since we updated the tier via API, the status should change immediately + mut throttled = false + let res = (http get -fe $"($url)/xrpc/com.atproto.sync.getHostStatus?hostname=($mock_host)") + if $res.status == 200 and $res.body.status == "throttled" { + $throttled = true + } + + if not $throttled { + stop-mock-pds $mock_pds_handle + try { kill $instance.pid } + fail "host did not transition to throttled after tier update" + } + print "ok: host transitioned to throttled successfully." + + print "checking status transitions back to Active when limits loosen..." + http delete -fe $"($url)/pds/tiers?host=($mock_host)" + + # should change back immediately + mut re_active = false + let res = (http get -fe $"($url)/xrpc/com.atproto.sync.getHostStatus?hostname=($mock_host)") + if $res.status == 200 and $res.body.status == "active" { + $re_active = true + } + + stop-mock-pds $mock_pds_handle + try { kill $instance.pid } + + 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" + } +} diff --git a/tests/run_all.nu b/tests/run_all.nu index b318f0a..c8ee9ca 100644 --- a/tests/run_all.nu +++ b/tests/run_all.nu @@ -46,7 +46,7 @@ def main [--only: list = [], --skip-creds] { print "" # discover all test scripts, excluding infrastructure files - mut excluded = ["common", "mock_relay", "run_all"] + mut excluded = ["common", "mock_relay", "mock_pds", "run_all"] if $skip_creds { $excluded = ($excluded | append ["authenticated_stream", "repo_sync_integrity"]) }