From d73d9a15a09ce07eb14b911e5590fe85290c5363 Mon Sep 17 00:00:00 2001 From: dawn <90008@gaze.systems> Date: Tue, 31 Mar 2026 10:33:15 +0300 Subject: [PATCH] [relay] ratelimits, ratelimit tiers and management --- AGENTS.md | 2 +- Cargo.lock | 1 + Cargo.toml | 1 + README.md | 24 ++- src/api/filter.rs | 3 +- src/api/mod.rs | 2 + src/api/pds.rs | 74 ++++++++ src/api/xrpc/mod.rs | 71 +++++--- src/api/xrpc/request_crawl.rs | 24 +++ src/config.rs | 102 +++++++++++ src/control/crawler.rs | 35 ++-- src/control/filter.rs | 5 +- src/control/firehose.rs | 162 ++++++++-------- src/control/mod.rs | 66 +++---- src/control/pds.rs | 113 ++++++++++++ src/control/repos.rs | 18 +- src/crawler/list_repos.rs | 6 +- src/crawler/mod.rs | 3 +- src/crawler/worker.rs | 31 ++-- src/db/filter.rs | 21 +-- src/db/keys/mod.rs | 4 + src/db/mod.rs | 2 +- src/db/pds_tiers.rs | 33 ++++ src/filter.rs | 16 +- src/ingest/firehose.rs | 19 +- src/ingest/relay.rs | 47 +++-- src/lib.rs | 1 + src/patch.rs | 11 ++ src/state.rs | 46 ++++- src/{util.rs => util/mod.rs} | 2 + src/{crawler => util}/throttle.rs | 119 +++++++++++- tests/api.nu | 294 ++++++++++++++++++++++++++++++ 32 files changed, 1126 insertions(+), 232 deletions(-) create mode 100644 src/api/pds.rs create mode 100644 src/api/xrpc/request_crawl.rs create mode 100644 src/control/pds.rs create mode 100644 src/db/pds_tiers.rs create mode 100644 src/patch.rs rename src/{util.rs => util/mod.rs} (99%) rename src/{crawler => util}/throttle.rs (61%) diff --git a/AGENTS.md b/AGENTS.md index 8b7069a..1c919ee 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -143,7 +143,7 @@ Examples: # GitNexus — Code Intelligence -This project is indexed by GitNexus as **hydrant** (655 symbols, 1810 relationships, 55 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely. +This project is indexed by GitNexus as **hydrant** (1339 symbols, 3645 relationships, 113 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely. > If any GitNexus tool warns the index is stale, run `npx gitnexus analyze` in terminal first. diff --git a/Cargo.lock b/Cargo.lock index ad75b5e..6e0c282 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1491,6 +1491,7 @@ dependencies = [ "miette", "mimalloc", "multibase", + "parking_lot", "rand 0.10.0", "reqwest", "rmp-serde", diff --git a/Cargo.toml b/Cargo.toml index 61f9c0e..5df713a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,6 +58,7 @@ rustls = { version = "0.23", features = ["aws-lc-rs"] } tokio-tungstenite = { version = "0.28.0", features = ["rustls-tls-native-roots"] } multibase = "0.9.2" sha2 = "0.10.9" +parking_lot = "0.12.5" [dev-dependencies] tempfile = "3.26.0" diff --git a/README.md b/README.md index 8f5e4e6..28bd11b 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ -> [hydrant](#hydrant)
-> [vs tap](#vs-tap) | [stream](#stream-behavior) | [multi-relay](#multiple-relay-support) | [crawler sources](#crawler-sources)
--> [configuration](#configuration)
+-> [configuration](#configuration) | [build features](#build-features)
-> [rest api](#rest-api) | [filter](#filter-management) | [ingestion](#ingestion-control) | [crawler](#crawler-management) | [firehose](#firehose-management) | [repos](#repository-management)
-> [xrpc api](#data-access-xrpc) | [backlinks](#bluemicrocosmlinks) | [identity](#bluemicrocosmidentity) | [atproto](#comatproto) | [custom](#systemsgazehydrant) @@ -142,6 +142,18 @@ 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. | +## build features + +[<- back to toc](#table-of-contents) + +`hydrant` has several optional compile-time features: + +| feature | default | description | +| :--- | :--- | :--- | +| `indexer` | yes | enables the indexing logic. | +| `relay` | no | enables relay functionality. | +| `backlinks` | no | enables the backlinks indexer and XRPC endpoints (`blue.microcosm.links.*`). | + ## REST api [<- back to toc](#table-of-contents) @@ -329,6 +341,8 @@ the following are implemented currently: - `com.atproto.sync.getRepoStatus` - `com.atproto.sync.listRepos` - `com.atproto.sync.getLatestCommit` +- `com.atproto.sync.requestCrawl` (adds the host to firehose sources in relay mode) +- `com.atproto.sync.subscribeRepos` (WebSocket firehose stream, requires `relay` feature) ### systems.gaze.hydrant.* @@ -397,11 +411,3 @@ return the number of records that link to a given subject. | `source` | no | filter by source collection (same format as `getBacklinks`). | returns `{ count }`. - -### blue.microcosm.identity.* - -[<- back to toc](#table-of-contents) - -#### blue.microcosm.identity.resolveMiniDoc - -see [here](https://slingshot.microcosm.blue/#tag/slingshot-specific-queries/GET/xrpc/blue.microcosm.identity.resolveMiniDoc) for this XRPC's documentation. diff --git a/src/api/filter.rs b/src/api/filter.rs index 5bf9312..00503c0 100644 --- a/src/api/filter.rs +++ b/src/api/filter.rs @@ -1,5 +1,6 @@ use crate::control::{FilterPatch, Hydrant}; -use crate::filter::{FilterMode, SetUpdate}; +use crate::filter::FilterMode; +use crate::patch::SetUpdate; use axum::{ Json, Router, extract::State, diff --git a/src/api/mod.rs b/src/api/mod.rs index f4299bb..a08cfac 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -11,6 +11,7 @@ mod debug; mod filter; mod firehose; mod ingestion; +mod pds; mod repos; mod stats; #[cfg(feature = "indexer")] @@ -27,6 +28,7 @@ pub async fn serve(hydrant: Hydrant, port: u16) -> miette::Result<()> { let app = app .merge(xrpc::router()) .merge(filter::router()) + .merge(pds::router()) .merge(repos::router()) .merge(ingestion::router()) .merge(crawler::router()) diff --git a/src/api/pds.rs b/src/api/pds.rs new file mode 100644 index 0000000..ce6baae --- /dev/null +++ b/src/api/pds.rs @@ -0,0 +1,74 @@ +use std::collections::HashMap; + +use axum::{ + Json, Router, + extract::State, + http::StatusCode, + routing::{delete, get, put}, +}; +use serde::{Deserialize, Serialize}; + +use crate::control::{Hydrant, PdsTierAssignment, PdsTierDefinition}; + +pub fn router() -> Router { + Router::new() + .route("/pds/tiers", get(list_tiers)) + .route("/pds/tiers", put(set_tier)) + .route("/pds/tiers", delete(remove_tier)) + .route("/pds/rate-tiers", get(list_rate_tiers)) +} + +/// combined response: tier assignments + available tier definitions. +#[derive(Serialize)] +pub struct TiersResponse { + pub assignments: Vec, + pub rate_tiers: HashMap, +} + +pub async fn list_tiers(State(hydrant): State) -> Json { + Json(TiersResponse { + assignments: hydrant.pds.list_assignments().await, + rate_tiers: hydrant.pds.list_rate_tiers(), + }) +} + +pub async fn list_rate_tiers( + State(hydrant): State, +) -> Json> { + Json(hydrant.pds.list_rate_tiers()) +} + +#[derive(Deserialize)] +pub struct SetTierBody { + pub host: String, + pub tier: String, +} + +pub async fn set_tier( + State(hydrant): State, + Json(body): Json, +) -> Result { + hydrant + .pds + .set_tier(body.host, body.tier) + .await + .map(|_| StatusCode::OK) + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string())) +} + +#[derive(Deserialize)] +pub struct RemoveTierBody { + pub host: String, +} + +pub async fn remove_tier( + State(hydrant): State, + Json(body): Json, +) -> Result { + hydrant + .pds + .remove_tier(body.host) + .await + .map(|_| StatusCode::OK) + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) +} diff --git a/src/api/xrpc/mod.rs b/src/api/xrpc/mod.rs index b098376..a717555 100644 --- a/src/api/xrpc/mod.rs +++ b/src/api/xrpc/mod.rs @@ -1,18 +1,10 @@ -use crate::api::xrpc::count_records::CountRecords; -use crate::api::xrpc::describe_repo::DescribeRepo; use crate::control::Hydrant; use axum::extract::FromRequest; use axum::response::IntoResponse; use axum::routing::get; use axum::{Json, Router, extract::State, http::StatusCode}; -use jacquard_api::com_atproto::repo::{ - describe_repo::DescribeRepoRequest as AtprotoDescribeRepoRequest, - get_record::{GetRecordError, GetRecordOutput, GetRecordRequest}, - list_records::{ListRecordsOutput, ListRecordsRequest, Record as RepoRecord}, -}; use jacquard_api::com_atproto::sync::get_host_status::GetHostStatusRequest; use jacquard_api::com_atproto::sync::get_latest_commit::GetLatestCommitRequest; -use jacquard_api::com_atproto::sync::get_repo::GetRepoRequest; use jacquard_api::com_atproto::sync::get_repo_status::GetRepoStatusRequest; use jacquard_api::com_atproto::sync::list_hosts::ListHostsRequest; use jacquard_api::com_atproto::sync::list_repos::ListReposRequest; @@ -27,48 +19,79 @@ use jacquard_common::{ use serde::{Deserialize, Serialize}; use smol_str::ToSmolStr; use std::fmt::Display; +use std::result::Result; +#[cfg(feature = "indexer")] +use { + crate::api::xrpc::count_records::CountRecords, + crate::api::xrpc::describe_repo::DescribeRepo, + jacquard_api::com_atproto::repo::{ + describe_repo::DescribeRepoRequest as AtprotoDescribeRepoRequest, + get_record::{GetRecordError, GetRecordOutput, GetRecordRequest}, + list_records::{ListRecordsOutput, ListRecordsRequest, Record as RepoRecord}, + }, + jacquard_api::com_atproto::sync::get_repo::GetRepoRequest, +}; #[cfg(feature = "relay")] use { + jacquard_api::com_atproto::sync::request_crawl::RequestCrawlRequest, jacquard_api::com_atproto::sync::subscribe_repos::SubscribeReposEndpoint, jacquard_common::xrpc::SubscriptionEndpoint, }; +mod get_host_status; +mod get_latest_commit; +mod get_repo_status; +mod list_hosts; +mod list_repos; + +#[cfg(feature = "indexer")] mod com_atproto_describe_repo; +#[cfg(feature = "indexer")] mod count_records; +#[cfg(feature = "indexer")] mod describe_repo; -mod get_host_status; -mod get_latest_commit; +#[cfg(feature = "indexer")] mod get_record; +#[cfg(feature = "indexer")] mod get_repo; -mod get_repo_status; -mod list_hosts; +#[cfg(feature = "indexer")] mod list_records; -mod list_repos; + +#[cfg(feature = "relay")] +mod request_crawl; #[cfg(feature = "relay")] mod subscribe_repos; pub fn router() -> Router { let r = Router::new() + .route(GetHostStatusRequest::PATH, get(get_host_status::handle)) + .route(ListHostsRequest::PATH, get(list_hosts::handle)) + .route(GetLatestCommitRequest::PATH, get(get_latest_commit::handle)) + .route(GetRepoStatusRequest::PATH, get(get_repo_status::handle)) + .route(ListReposRequest::PATH, get(list_repos::handle)); + + #[cfg(feature = "indexer")] + let r = r .route(GetRecordRequest::PATH, get(get_record::handle)) .route(ListRecordsRequest::PATH, get(list_records::handle)) .route(CountRecords::PATH, get(count_records::handle)) + .route(GetRepoRequest::PATH, get(get_repo::handle)) .route(DescribeRepo::PATH, get(describe_repo::handle)) .route( AtprotoDescribeRepoRequest::PATH, get(com_atproto_describe_repo::handle), - ) - .route(GetHostStatusRequest::PATH, get(get_host_status::handle)) - .route(ListHostsRequest::PATH, get(list_hosts::handle)) - .route(GetLatestCommitRequest::PATH, get(get_latest_commit::handle)) - .route(GetRepoRequest::PATH, get(get_repo::handle)) - .route(GetRepoStatusRequest::PATH, get(get_repo_status::handle)) - .route(ListReposRequest::PATH, get(list_repos::handle)); + ); #[cfg(feature = "relay")] - let r = r.route( - SubscribeReposEndpoint::PATH, - axum::routing::get(subscribe_repos::handle), - ); + let r = r + .route( + SubscribeReposEndpoint::PATH, + axum::routing::get(subscribe_repos::handle), + ) + .route( + RequestCrawlRequest::PATH, + axum::routing::get(subscribe_repos::handle), + ); r } diff --git a/src/api/xrpc/request_crawl.rs b/src/api/xrpc/request_crawl.rs new file mode 100644 index 0000000..c52d0be --- /dev/null +++ b/src/api/xrpc/request_crawl.rs @@ -0,0 +1,24 @@ +use jacquard_api::com_atproto::sync::request_crawl::{ + RequestCrawlError, RequestCrawlRequest, RequestCrawlResponse, +}; +use url::Url; + +use super::*; + +pub async fn handle( + State(hydrant): State, + ExtractXrpc(req): ExtractXrpc, +) -> XrpcResult> { + let nsid = RequestCrawlResponse::NSID; + + let url_str = format!("wss://{}/", req.hostname); + let url = Url::parse(&url_str).map_err(|e| bad_request(nsid, e))?; + + hydrant + .firehose + .add_source(url, true) + .await + .map_err(|e| internal_error(nsid, e))?; + + Ok(StatusCode::OK) +} diff --git a/src/config.rs b/src/config.rs index 8a4a3d0..06fe879 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,12 +1,65 @@ use miette::Result; use serde::{Deserialize, Serialize}; use smol_str::ToSmolStr; +use std::collections::HashMap; use std::fmt; use std::path::PathBuf; use std::str::FromStr; use std::time::Duration; use url::Url; +/// rate limit parameters for a named tier of PDS connections. +/// +/// the per-second limit is `max(per_second_base, accounts * per_second_account_mul)`, +/// giving a floor at `per_second_base` that scales up with the PDS's active account count. +#[derive(Debug, Clone, Copy)] +pub struct RateTier { + /// floor for the per-second limit, regardless of account count. + pub per_second_base: u64, + /// per-second events allowed per active account on this PDS. + pub per_second_account_mul: f64, + /// per-hour limit. + pub per_hour: u64, + /// per-day limit. + pub per_day: u64, +} + +impl RateTier { + /// built-in "trusted" tier: high limits for well-behaved PDS operators. + pub fn trusted() -> Self { + Self { + per_second_base: 5000, + per_second_account_mul: 10.0, + per_hour: 5000 * 3600, + per_day: 5000 * 86400, + } + } + + /// built-in "default" tier: conservative limits for unknown PDS operators. + pub fn default_tier() -> Self { + Self { + per_second_base: 50, + per_second_account_mul: 0.5, + per_hour: 1000 * 3600, + per_day: 1000 * 86400, + } + } + + /// parse `base/mul/hourly/daily` format used by `HYDRANT_RATE_TIERS`. + fn parse(s: &str) -> Option { + let parts: Vec<&str> = s.split('/').collect(); + if parts.len() != 4 { + return None; + } + Some(Self { + per_second_base: parts[0].parse().ok()?, + per_second_account_mul: parts[1].parse().ok()?, + per_hour: parts[2].parse().ok()?, + per_day: parts[3].parse().ok()?, + }) + } +} + /// this is for internal use only, please don't use this macro. #[doc(hidden)] #[macro_export] @@ -309,6 +362,17 @@ pub struct Config { /// set via `HYDRANT_ENABLE_BACKLINKS=true`. pub enable_backlinks: bool, + /// 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, + /// db internals, tune only if you know what you're doing. /// /// size of the fjall block cache in MB. set via `HYDRANT_CACHE_SIZE`. @@ -388,6 +452,13 @@ 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 + }, cache_size: 256, data_compression: Compression::Lz4, journal_compression: Compression::Lz4, @@ -549,6 +620,35 @@ impl Config { let enable_backlinks: bool = cfg!("ENABLE_BACKLINKS", defaults.enable_backlinks); + // start with built-in tiers, then layer in any env-defined overrides. + // format: HYDRANT_RATE_TIERS=name:base/mul/hourly/daily,... + let mut rate_tiers = defaults.rate_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); + } + None => tracing::warn!( + "ignoring invalid rate tier '{name}': expected base/mul/hourly/daily format" + ), + } + } + } + } + + 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()); + let default_mode = CrawlerMode::default_for(full_network); let crawler_sources = match std::env::var("HYDRANT_CRAWLER_URLS") { Ok(s) => s @@ -593,6 +693,8 @@ impl Config { filter_collections, filter_excludes, enable_backlinks, + trusted_hosts, + rate_tiers, cache_size, data_compression, journal_compression, diff --git a/src/control/crawler.rs b/src/control/crawler.rs index 4202d46..3b59b85 100644 --- a/src/control/crawler.rs +++ b/src/control/crawler.rs @@ -149,19 +149,20 @@ impl CrawlerHandle { /// delete all cursor entries associated with the given URL. pub async fn reset_cursor(&self, url: &str) -> Result<()> { - let db = self.state.db.clone(); + let state = self.state.clone(); let point_keys = [keys::crawler_cursor_key(url)]; let by_collection_prefix = keys::by_collection_cursor_prefix(url); tokio::task::spawn_blocking(move || { - let mut batch = db.inner.batch(); + let mut batch = state.db.inner.batch(); for k in point_keys { - batch.remove(&db.cursors, k); + batch.remove(&state.db.cursors, k); } - for entry in db.cursors.prefix(&by_collection_prefix) { + for entry in state.db.cursors.prefix(&by_collection_prefix) { let k = entry.key().into_diagnostic()?; - batch.remove(&db.cursors, k); + batch.remove(&state.db.cursors, k); } - batch.commit().into_diagnostic() + batch.commit().into_diagnostic()?; + state.db.persist() }) .await .into_diagnostic()??; @@ -198,12 +199,15 @@ impl CrawlerHandle { miette::bail!("crawler not yet started: call Hydrant::run() first"); }; - let db = self.state.db.clone(); + let state = self.state.clone(); let key = keys::crawler_source_key(source.url.as_str()); let val = rmp_serde::to_vec(&source.mode).into_diagnostic()?; - tokio::task::spawn_blocking(move || db.crawler.insert(key, val).into_diagnostic()) - .await - .into_diagnostic()??; + tokio::task::spawn_blocking(move || { + state.db.crawler.insert(key, val).into_diagnostic()?; + state.db.persist() + }) + .await + .into_diagnostic()??; let enabled_rx = self.state.crawler_enabled.subscribe(); let handle = spawn_crawler_producer( @@ -249,11 +253,14 @@ impl CrawlerHandle { // remove from DB if it was a persisted source if self.persisted.remove_async(url).await.is_some() { - let db = self.state.db.clone(); + let state = self.state.clone(); let key = keys::crawler_source_key(url.as_str()); - tokio::task::spawn_blocking(move || db.crawler.remove(key).into_diagnostic()) - .await - .into_diagnostic()??; + tokio::task::spawn_blocking(move || { + state.db.crawler.remove(key).into_diagnostic()?; + state.db.persist() + }) + .await + .into_diagnostic()??; } Ok(true) diff --git a/src/control/filter.rs b/src/control/filter.rs index 93f6419..5a6ba8f 100644 --- a/src/control/filter.rs +++ b/src/control/filter.rs @@ -4,7 +4,8 @@ use tracing::error; use miette::{IntoDiagnostic, Result}; use crate::db::filter as db_filter; -use crate::filter::{FilterMode, SetUpdate}; +use crate::filter::FilterMode; +use crate::patch::SetUpdate; use crate::state::AppState; /// a point-in-time snapshot of the filter configuration. returned by all [`FilterControl`] methods. @@ -273,6 +274,7 @@ impl FilterPatch { let filter_ks = self.state.db.filter.clone(); let inner = self.state.db.inner.clone(); let filter_handle = self.state.filter.clone(); + let state = self.state.clone(); let mode = self.mode; let signals = self.signals; let collections = self.collections; @@ -282,6 +284,7 @@ impl FilterPatch { let mut batch = inner.batch(); db_filter::apply_patch(&mut batch, &filter_ks, mode, signals, collections, excludes)?; batch.commit().into_diagnostic()?; + state.db.persist()?; db_filter::load(&filter_ks) }) .await diff --git a/src/control/firehose.rs b/src/control/firehose.rs index 80e7d55..24ab21e 100644 --- a/src/control/firehose.rs +++ b/src/control/firehose.rs @@ -1,7 +1,6 @@ use std::sync::Arc; -use std::sync::atomic::Ordering; -use miette::{Context, IntoDiagnostic, Result}; +use miette::{IntoDiagnostic, Result}; use tokio::sync::watch; use tracing::{error, info}; use url::Url; @@ -62,7 +61,8 @@ pub(super) async fn spawn_firehose_ingestor( state.filter.clone(), enabled, shared.verify_signatures, - ); + ) + .await; let relay_for_log = relay_url.clone(); let abort = tokio::spawn(async move { @@ -88,118 +88,124 @@ pub struct FirehoseHandle { } impl FirehoseHandle { - /// enable the firehose. no-op if already enabled. + pub(super) fn new(state: Arc) -> Self { + Self { + state, + shared: Arc::new(std::sync::OnceLock::new()), + tasks: Arc::new(scc::HashMap::new()), + persisted: Arc::new(scc::HashSet::new()), + } + } + + /// enable firehose ingestion, no-op if already enabled. pub fn enable(&self) { self.state.firehose_enabled.send_replace(true); } - /// disable the firehose. the current message finishes processing before the connection closes. + /// disable firehose ingestion, in-flight messages complete before pausing. pub fn disable(&self) { self.state.firehose_enabled.send_replace(false); } - /// returns the current enabled state of the firehose. + /// returns the current enabled state of firehose ingestion. pub fn is_enabled(&self) -> bool { *self.state.firehose_enabled.borrow() } - /// reset the stored cursor for the given relay URL. - /// - /// clears the `firehose_cursor|{host}|{scheme}` entry from the cursors keyspace and zeroes - /// the in-memory cursor. the next connection will tail live events from the current head. - pub async fn reset_cursor(&self, url: &str) -> Result<()> { - let relay_url = Url::parse(url) - .into_diagnostic() - .wrap_err_with(|| format!("invalid relay url: {url:?}"))?; - let key = keys::firehose_cursor_key_from_url(&relay_url); - let db = self.state.db.clone(); - tokio::task::spawn_blocking(move || db.cursors.remove(key).into_diagnostic()) - .await - .into_diagnostic()??; - - self.state.firehose_cursors.peek_with(&relay_url, |_, c| { - c.store(0, Ordering::SeqCst); - }); - Ok(()) - } - - /// return info on all currently active firehose sources. + /// list all currently active firehose sources. pub async fn list_sources(&self) -> Vec { - let mut sources = Vec::new(); + let mut out = Vec::new(); self.tasks - .iter_async(|url, handle| { - sources.push(FirehoseSourceInfo { + .any_async(|url, handle| { + out.push(FirehoseSourceInfo { url: url.clone(), persisted: self.persisted.contains_sync(url), is_pds: handle.is_pds, }); - true + false }) .await; - sources + out } - /// add a new firehose relay at runtime. - /// - /// the URL is persisted to the database and will be re-spawned on restart. if a relay with - /// the same URL already exists it is replaced: the running task is stopped and a new one - /// is started. any cursor state for that URL is preserved. + /// add a new firehose source at runtime, persisting it to the database. /// - /// returns an error if called before [`Hydrant::run`]. + /// if a source with the same URL already exists, it is replaced: the + /// running task is stopped and a new one is started with the new `is_pds` + /// setting. existing cursor state for the URL is preserved. pub async fn add_source(&self, url: Url, is_pds: bool) -> Result<()> { - let Some(shared) = self.shared.get() else { - miette::bail!("firehose not yet started: call Hydrant::run() first"); - }; + let shared = self + .shared + .get() + .ok_or_else(|| miette::miette!("firehose worker not started"))?; - let db = self.state.db.clone(); + // persist to db first let key = keys::firehose_source_key(url.as_str()); - let value = rmp_serde::to_vec(&crate::db::FirehoseSourceMeta { is_pds }) - .map_err(|e| miette::miette!("failed to serialize firehose source meta: {e}"))?; - tokio::task::spawn_blocking(move || db.crawler.insert(key, value).into_diagnostic()) - .await - .into_diagnostic()??; + tokio::task::spawn_blocking({ + let state = self.state.clone(); + move || { + let mut batch = state.db.inner.batch(); + let value = rmp_serde::to_vec(&db::FirehoseSourceMeta { is_pds }).map_err(|e| { + miette::miette!("failed to serialize firehose source meta: {e}") + })?; + batch.insert(&state.db.crawler, key, &value); + batch.commit().into_diagnostic()?; + state.db.persist() + } + }) + .await + .into_diagnostic()??; + + let _ = self.persisted.insert_async(url.clone()).await; let enabled_rx = self.state.firehose_enabled.subscribe(); let handle = spawn_firehose_ingestor(&url, is_pds, &self.state, shared, enabled_rx).await?; + self.tasks.upsert_async(url, handle).await; - let _ = self.persisted.insert_async(url.clone()).await; - match self.tasks.entry_async(url).await { - scc::hash_map::Entry::Vacant(e) => { - e.insert_entry(handle); - } - scc::hash_map::Entry::Occupied(mut e) => { - *e.get_mut() = handle; - } - } Ok(()) } - /// remove a firehose relay at runtime by URL. - /// - /// aborts the running ingestor task. if the source was added via the API it is removed from - /// the database and will not reappear on restart. `RELAY_HOSTS` sources are only stopped for - /// the current session; they reappear on the next restart. + /// remove a firehose source at runtime. /// - /// returns `true` if the relay was found and removed, `false` if it was not running. - /// returns an error if called before [`Hydrant::run`]. + /// returns `true` if the source was found and removed, `false` otherwise. + /// if the source was added via the API, it is removed from the database; + /// if it came from the static config, only the running task is stopped. pub async fn remove_source(&self, url: &Url) -> Result { - if self.shared.get().is_none() { - miette::bail!("firehose not yet started: call Hydrant::run() first"); + if self.persisted.contains_async(url).await { + let url_str = url.to_string(); + tokio::task::spawn_blocking({ + let state = self.state.clone(); + move || { + state + .db + .crawler + .remove(keys::firehose_source_key(&url_str)) + .into_diagnostic()?; + state.db.persist() + } + }) + .await + .into_diagnostic()??; + self.persisted.remove_async(url).await; } - if self.tasks.remove_async(url).await.is_none() { - return Ok(false); - } + Ok(self.tasks.remove_async(url).await.is_some()) + } - // remove from relay_cursors (persist thread will stop tracking it) - self.state.firehose_cursors.remove_async(url).await; + /// reset the stored firehose cursor for a given URL. + pub async fn reset_cursor(&self, url: &str) -> Result<()> { + let url = Url::parse(url).into_diagnostic()?; + let key = keys::firehose_cursor_key_from_url(&url); + tokio::task::spawn_blocking({ + let state = self.state.clone(); + move || { + state.db.cursors.remove(key).into_diagnostic()?; + state.db.persist() + } + }) + .await + .into_diagnostic()??; - if self.persisted.remove_async(url).await.is_some() { - let db = self.state.db.clone(); - let key = keys::firehose_source_key(url.as_str()); - tokio::task::spawn_blocking(move || db.crawler.remove(key).into_diagnostic()) - .await - .into_diagnostic()??; - } + self.state.firehose_cursors.remove_async(&url).await; - Ok(true) + Ok(()) } } diff --git a/src/control/mod.rs b/src/control/mod.rs index 3ae9010..a1d507d 100644 --- a/src/control/mod.rs +++ b/src/control/mod.rs @@ -3,12 +3,14 @@ pub(crate) mod crawler; pub(crate) mod filter; pub(crate) mod firehose; +pub(crate) mod pds; pub(crate) mod repos; pub(crate) mod stream; pub use crawler::{CrawlerHandle, CrawlerSourceInfo}; pub use filter::{FilterControl, FilterPatch, FilterSnapshot}; pub use firehose::{FirehoseHandle, FirehoseSourceInfo}; +pub use pds::{PdsControl, PdsTierAssignment, PdsTierDefinition}; pub use repos::{ListedRecord, Record, RecordList, RepoHandle, RepoInfo, ReposControl}; use smol_str::{SmolStr, ToSmolStr}; @@ -87,6 +89,7 @@ pub struct Hydrant { pub firehose: FirehoseHandle, pub backfill: BackfillHandle, pub filter: FilterControl, + pub pds: PdsControl, pub repos: ReposControl, pub db: DbControl, #[cfg(feature = "backlinks")] @@ -121,15 +124,15 @@ impl Hydrant { let signals = config .filter_signals .clone() - .map(crate::filter::SetUpdate::Set); + .map(crate::patch::SetUpdate::Set); let collections = config .filter_collections .clone() - .map(crate::filter::SetUpdate::Set); + .map(crate::patch::SetUpdate::Set); let excludes = config .filter_excludes .clone() - .map(crate::filter::SetUpdate::Set); + .map(crate::patch::SetUpdate::Set); tokio::task::spawn_blocking(move || { let mut batch = inner.batch(); @@ -174,14 +177,10 @@ impl Hydrant { tasks: Arc::new(scc::HashMap::new()), persisted: Arc::new(scc::HashSet::new()), }, - firehose: FirehoseHandle { - state: state.clone(), - shared: Arc::new(std::sync::OnceLock::new()), - tasks: Arc::new(scc::HashMap::new()), - persisted: Arc::new(scc::HashSet::new()), - }, + firehose: FirehoseHandle::new(state.clone()), backfill: BackfillHandle(state.clone()), filter: FilterControl(state.clone()), + pds: pds::PdsControl(state.clone()), repos: ReposControl(state.clone()), db: DbControl(state.clone()), #[cfg(feature = "backlinks")] @@ -436,10 +435,10 @@ impl Hydrant { // 11. spawn crawler infrastructure #[cfg(feature = "indexer")] { - use crate::crawler::throttle::Throttler; use crate::crawler::{ CrawlerStats, CrawlerWorker, InFlight, RetryProducer, SignalChecker, }; + use crate::util::throttle::Throttler; let http = reqwest::Client::builder() .user_agent(concat!( @@ -450,7 +449,7 @@ impl Hydrant { .gzip(true) .build() .expect("that reqwest will build"); - let pds_throttler = Throttler::new(); + let pds_throttler = state.throttler.clone(); let in_flight = InFlight::new(); let stats = CrawlerStats::new( state.clone(), @@ -714,7 +713,7 @@ impl Hydrant { /// /// sizes are in bytes, reported per keyspace. pub async fn stats(&self) -> Result { - let db = self.state.db.clone(); + let state = self.state.clone(); let mut counts: BTreeMap<&'static str, u64> = futures::future::join_all( [ @@ -729,29 +728,29 @@ impl Hydrant { ] .into_iter() .map(|name| { - let db = db.clone(); - async move { (name, db.get_count(name).await) } + let state = state.clone(); + async move { (name, state.db.get_count(name).await) } }), ) .await .into_iter() .collect(); - counts.insert("events", db.events.approximate_len() as u64); + counts.insert("events", state.db.events.approximate_len() as u64); let sizes = tokio::task::spawn_blocking(move || { let mut s = BTreeMap::new(); - s.insert("repos", db.repos.disk_space()); - s.insert("records", db.records.disk_space()); - s.insert("blocks", db.blocks.disk_space()); - s.insert("cursors", db.cursors.disk_space()); - s.insert("pending", db.pending.disk_space()); - s.insert("resync", db.resync.disk_space()); - s.insert("resync_buffer", db.resync_buffer.disk_space()); - s.insert("events", db.events.disk_space()); - s.insert("counts", db.counts.disk_space()); - s.insert("filter", db.filter.disk_space()); - s.insert("crawler", db.crawler.disk_space()); + s.insert("repos", state.db.repos.disk_space()); + s.insert("records", state.db.records.disk_space()); + s.insert("blocks", state.db.blocks.disk_space()); + s.insert("cursors", state.db.cursors.disk_space()); + s.insert("pending", state.db.pending.disk_space()); + s.insert("resync", state.db.resync.disk_space()); + s.insert("resync_buffer", state.db.resync_buffer.disk_space()); + s.insert("events", state.db.events.disk_space()); + s.insert("counts", state.db.counts.disk_space()); + s.insert("filter", state.db.filter.disk_space()); + s.insert("crawler", state.db.crawler.disk_space()); s }) .await @@ -788,12 +787,12 @@ impl Hydrant { /// /// returns the seq we are on for this host. pub async fn get_host_status(&self, hostname: &str) -> Result> { - let db = self.state.db.clone(); + let state = self.state.clone(); let hostname = hostname.to_smolstr(); tokio::task::spawn_blocking(move || { let key = keys::firehose_cursor_key(&hostname); - let Some(seq) = db.cursors.get(&key).into_diagnostic()? else { + let Some(seq) = state.db.cursors.get(&key).into_diagnostic()? else { return Ok(None); }; let seq = i64::from_be_bytes( @@ -820,7 +819,7 @@ impl Hydrant { cursor: Option<&str>, limit: usize, ) -> Result<(Vec, Option)> { - let db = self.state.db.clone(); + let state = self.state.clone(); let cursor = cursor.map(str::to_string); tokio::task::spawn_blocking(move || { @@ -836,7 +835,8 @@ impl Hydrant { // fetch one extra item to detect whether there is a next page let mut hosts: Vec = Vec::with_capacity(limit + 1); - for item in db + for item in state + .db .cursors .range((start_bound, std::ops::Bound::Excluded(prefix_end))) .take(limit + 1) @@ -972,9 +972,9 @@ impl DbControl { state .with_ingestion_paused(async || { let train = |name: &'static str| { - let db = state.db.clone(); - tokio::task::spawn_blocking(move || db.train_dict(name)) - .map(|res| res.into_diagnostic().flatten()) + let state = state.clone(); + tokio::task::spawn_blocking(move || state.db.train_dict(name)) + .map(|res: Result<_, _>| res.into_diagnostic().flatten()) }; tokio::try_join!(train("repos"), train("blocks"), train("events")).map(|_| ()) }) diff --git a/src/control/pds.rs b/src/control/pds.rs new file mode 100644 index 0000000..68ff615 --- /dev/null +++ b/src/control/pds.rs @@ -0,0 +1,113 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use miette::{IntoDiagnostic, Result}; +use serde::Serialize; +use smol_str::SmolStr; + +use crate::config::RateTier; +use crate::db::pds_tiers as db_pds; +use crate::state::AppState; + +/// a single PDS-to-tier assignment. +#[derive(Debug, Clone, Serialize)] +pub struct PdsTierAssignment { + pub host: String, + pub tier: String, +} + +/// a rate tier definition, as returned by the API. +#[derive(Debug, Clone, Serialize)] +pub struct PdsTierDefinition { + pub per_second_base: u64, + pub per_second_account_mul: f64, + pub per_hour: u64, + pub per_day: u64, +} + +impl From for PdsTierDefinition { + fn from(t: RateTier) -> Self { + Self { + per_second_base: t.per_second_base, + per_second_account_mul: t.per_second_account_mul, + per_hour: t.per_hour, + per_day: t.per_day, + } + } +} + +/// runtime control over pds related behaviour (eg. ratelimits). +#[derive(Clone)] +pub struct PdsControl(pub(super) Arc); + +impl PdsControl { + /// list all current per-PDS tier assignments. + pub async fn list_assignments(&self) -> Vec { + let snapshot = self.0.pds_tiers.load(); + snapshot + .iter() + .map(|(host, tier)| PdsTierAssignment { + host: host.clone(), + tier: tier.to_string(), + }) + .collect() + } + + /// list all configured rate tier definitions. + pub fn list_rate_tiers(&self) -> HashMap { + self.0 + .rate_tiers + .iter() + .map(|(name, tier)| (name.clone(), PdsTierDefinition::from(*tier))) + .collect() + } + + /// assign `host` to `tier`, persisting the change to the database. + /// returns an error if `tier` is not a known tier name. + pub async fn set_tier(&self, host: String, tier: String) -> Result<()> { + if !self.0.rate_tiers.contains_key(&tier) { + miette::bail!( + "unknown tier '{tier}'; known tiers: {:?}", + self.0.rate_tiers.keys().collect::>() + ); + } + + let state = self.0.clone(); + let host_clone = host.clone(); + let tier_clone = tier.clone(); + tokio::task::spawn_blocking(move || { + let mut batch = state.db.inner.batch(); + db_pds::set(&mut batch, &state.db.filter, &host_clone, &tier_clone); + batch.commit().into_diagnostic()?; + state.db.persist() + }) + .await + .into_diagnostic()??; + + let mut snapshot = (**self.0.pds_tiers.load()).clone(); + snapshot.insert(host, SmolStr::new(&tier)); + self.0.pds_tiers.store(Arc::new(snapshot)); + + Ok(()) + } + + /// remove any explicit tier assignment for `host`, reverting it to the default tier. + pub async fn remove_tier(&self, host: String) -> Result<()> { + let state = self.0.clone(); + let host_clone = host.clone(); + tokio::task::spawn_blocking(move || { + let mut batch = state.db.inner.batch(); + db_pds::remove(&mut batch, &state.db.filter, &host_clone); + batch.commit().into_diagnostic()?; + state.db.persist() + }) + .await + .into_diagnostic()??; + + let mut snapshot = (**self.0.pds_tiers.load()).clone(); + snapshot.remove(&host); + self.0.pds_tiers.store(Arc::new(snapshot)); + + Ok(()) + } +} diff --git a/src/control/repos.rs b/src/control/repos.rs index 893db0f..237b6be 100644 --- a/src/control/repos.rs +++ b/src/control/repos.rs @@ -86,7 +86,7 @@ impl ReposControl { std::ops::Bound::Unbounded }; - let db = self.0.db.clone(); + let state = self.0.clone(); self.0 .db .repos @@ -96,7 +96,8 @@ impl ReposControl { let repo_state = crate::db::deser_repo_state(&v)?.into_static(); let did = TrimmedDid::try_from(k.as_ref())?.to_did(); let metadata_key = keys::repo_metadata_key(&did); - let metadata = db + let metadata = state + .db .repo_metadata .get(&metadata_key) .into_diagnostic()? @@ -122,7 +123,7 @@ impl ReposControl { }; let repos = self.0.db.repos.clone(); - let db = self.0.db.clone(); + let state = self.0.clone(); self.0 .db .pending @@ -144,7 +145,8 @@ impl ReposControl { let repo_state = crate::db::deser_repo_state(bytes.as_ref())?; let did = TrimmedDid::try_from(did_key.as_ref())?.to_did(); let metadata_key = keys::repo_metadata_key(&did); - let metadata = db + let metadata = state + .db .repo_metadata .get(&metadata_key) .into_diagnostic()? @@ -169,7 +171,7 @@ impl ReposControl { }; let repos = self.0.db.repos.clone(); - let db = self.0.db.clone(); + let state = self.0.clone(); self.0 .db .resync @@ -184,7 +186,8 @@ impl ReposControl { let repo_state = crate::db::deser_repo_state(bytes.as_ref())?; let did = TrimmedDid::try_from(did_key.as_ref())?.to_did(); let metadata_key = keys::repo_metadata_key(&did); - let metadata = db + let metadata = state + .db .repo_metadata .get(&metadata_key) .into_diagnostic()? @@ -304,6 +307,7 @@ impl ReposControl { } batch.commit().into_diagnostic()?; + state.db.persist()?; Ok::<_, miette::Report>((queued, transitions)) }) .await @@ -369,6 +373,7 @@ impl ReposControl { } batch.commit().into_diagnostic()?; + state.db.persist()?; Ok::<_, miette::Report>((added, queued, transitions)) }) .await @@ -438,6 +443,7 @@ impl ReposControl { } batch.commit().into_diagnostic()?; + state.db.persist()?; Ok::<_, miette::Report>((untracked, gauge_decrements)) }) .await diff --git a/src/crawler/list_repos.rs b/src/crawler/list_repos.rs index ee23f50..3583d58 100644 --- a/src/crawler/list_repos.rs +++ b/src/crawler/list_repos.rs @@ -1,7 +1,7 @@ -use crate::crawler::throttle::{OrFailure, ThrottleHandle, Throttler}; use crate::db::keys::crawler_cursor_key; use crate::db::{Db, keys}; use crate::state::AppState; +use crate::util::throttle::{OrFailure, ThrottleHandle, Throttler}; use crate::util::{ ErrorForStatus, RetryOutcome, RetryWithBackoff, WatchEnabledExt, parse_retry_after, }; @@ -653,7 +653,7 @@ impl RetryProducer { } async fn process_queue(&self) -> Result> { - let db = self.checker.state.db.clone(); + let state = self.checker.state.clone(); struct ScanResult { ready: Vec>, @@ -675,7 +675,7 @@ impl RetryProducer { let mut next_wake: Option = None; let mut had_more = false; - for guard in db.crawler.prefix(keys::CRAWLER_RETRY_PREFIX) { + for guard in state.db.crawler.prefix(keys::CRAWLER_RETRY_PREFIX) { let (key, val) = guard.into_inner().into_diagnostic()?; let state: RetryState = rmp_serde::from_slice(&val).into_diagnostic()?; let did = keys::crawler_retry_parse_key(&key)?.to_did(); diff --git a/src/crawler/mod.rs b/src/crawler/mod.rs index b496722..050b64c 100644 --- a/src/crawler/mod.rs +++ b/src/crawler/mod.rs @@ -14,10 +14,9 @@ use url::Url; mod by_collection; mod list_repos; -pub mod throttle; mod worker; -use throttle::Throttler; +use crate::util::throttle::Throttler; pub(crate) use by_collection::ByCollectionProducer; pub(crate) use list_repos::{ListReposProducer, RetryProducer, SignalChecker}; diff --git a/src/crawler/worker.rs b/src/crawler/worker.rs index 2242db6..bf81ba3 100644 --- a/src/crawler/worker.rs +++ b/src/crawler/worker.rs @@ -137,35 +137,44 @@ impl CrawlerWorker { // filter already-known repos, build and commit the write batch, then return // the surviving guards so they are dropped on the async side after commit. - let db = self.state.db.clone(); + let app_state = self.state.clone(); let surviving = tokio::time::timeout( BLOCKING_TASK_TIMEOUT, tokio::task::spawn_blocking(move || -> Result> { let mut rng: SmallRng = rand::make_rng(); - let mut batch = db.inner.batch(); + let mut batch = app_state.db.inner.batch(); let mut surviving = Vec::new(); for guard in guards { let did_key = keys::repo_key(&*guard); let metadata_key = keys::repo_metadata_key(&*guard); - if db.repos.contains_key(&did_key).into_diagnostic()? { + if app_state + .db + .repos + .contains_key(&did_key) + .into_diagnostic()? + { continue; } let state = RepoState::backfilling(); let metadata = RepoMetadata::backfilling(rng.next_u64()); - batch.insert(&db.repos, &did_key, ser_repo_state(&state)?); + batch.insert(&app_state.db.repos, &did_key, ser_repo_state(&state)?); batch.insert( - &db.repo_metadata, + &app_state.db.repo_metadata, &metadata_key, crate::db::ser_repo_metadata(&metadata)?, ); - batch.insert(&db.pending, keys::pending_key(metadata.index_id), &did_key); + batch.insert( + &app_state.db.pending, + keys::pending_key(metadata.index_id), + &did_key, + ); // clear any stale retry entry, this DID is confirmed and being enqueued - batch.remove(&db.crawler, keys::crawler_retry_key(&*guard)); + batch.remove(&app_state.db.crawler, keys::crawler_retry_key(&*guard)); trace!(did = %*guard, "enqueuing repo"); surviving.push(guard); } if let Some(cursor) = cursor_update { - batch.insert(&db.cursors, cursor.key, cursor.value); + batch.insert(&app_state.db.cursors, cursor.key, cursor.value); } // todo: repo state overwrites here are acceptable? batch.commit().into_diagnostic()?; @@ -202,12 +211,12 @@ impl CrawlerWorker { } async fn commit_cursor(&self, cursor: CursorUpdate) -> Result<()> { - let db = self.state.db.clone(); + let state = self.state.clone(); tokio::time::timeout( BLOCKING_TASK_TIMEOUT, tokio::task::spawn_blocking(move || { - let mut batch = db.inner.batch(); - batch.insert(&db.cursors, cursor.key, cursor.value); + let mut batch = state.db.inner.batch(); + batch.insert(&state.db.cursors, cursor.key, cursor.value); batch.commit().into_diagnostic() }), ) diff --git a/src/db/filter.rs b/src/db/filter.rs index 487e89f..f5829fc 100644 --- a/src/db/filter.rs +++ b/src/db/filter.rs @@ -1,11 +1,10 @@ use fjall::{Keyspace, OwnedWriteBatch}; -use jacquard_common::IntoStatic; -use jacquard_common::types::nsid::Nsid; use jacquard_common::types::string::Did; use miette::{IntoDiagnostic, Result}; use crate::db::types::TrimmedDid; -use crate::filter::{FilterConfig, FilterMode, SetUpdate}; +use crate::filter::{FilterConfig, FilterMode}; +use crate::patch::SetUpdate; pub const MODE_KEY: &[u8] = b"m"; pub const SIGNAL_PREFIX: u8 = b's'; @@ -113,14 +112,14 @@ pub fn load(ks: &Keyspace) -> Result { for guard in ks.prefix(signal_prefix) { let (k, _) = guard.into_inner().into_diagnostic()?; let val = std::str::from_utf8(&k[signal_prefix.len()..]).into_diagnostic()?; - config.signals.push(Nsid::new(val)?.into_static()); + config.signals.push(val.into()); } let col_prefix = [COLLECTION_PREFIX, SEP]; for guard in ks.prefix(col_prefix) { let (k, _) = guard.into_inner().into_diagnostic()?; let val = std::str::from_utf8(&k[col_prefix.len()..]).into_diagnostic()?; - config.collections.push(Nsid::new(val)?.into_static()); + config.collections.push(val.into()); } Ok(config) @@ -144,6 +143,8 @@ pub fn read_set(ks: &Keyspace, prefix: u8) -> Result> { #[cfg(test)] mod tests { + use smol_str::SmolStr; + use super::*; #[test] @@ -198,14 +199,8 @@ mod tests { let config = load(&ks)?; assert_eq!(config.mode, FilterMode::Filter); - assert_eq!( - config.signals, - vec![Nsid::new("a.b.c").unwrap().into_static()] - ); - assert_eq!( - config.collections, - vec![Nsid::new("d.e.f").unwrap().into_static()] - ); + assert_eq!(config.signals, vec![SmolStr::new("a.b.c")]); + assert_eq!(config.collections, vec![SmolStr::new("d.e.f")]); let excludes = read_set(&ks, EXCLUDE_PREFIX)?; assert_eq!(excludes, vec!["did:plc:yk4q3id7id6p5z3bypvshc64"]); diff --git a/src/db/keys/mod.rs b/src/db/keys/mod.rs index 0b4501c..35f76ef 100644 --- a/src/db/keys/mod.rs +++ b/src/db/keys/mod.rs @@ -231,6 +231,10 @@ pub fn firehose_source_key(url: &str) -> Vec { key } +pub fn pds_account_count_key(host: &str) -> String { + format!("p|{host}") +} + #[cfg(feature = "relay")] /// key format: {SEQ} (u64 big-endian), mirroring event_key pub fn relay_event_key(seq: u64) -> [u8; 8] { diff --git a/src/db/mod.rs b/src/db/mod.rs index 73c9558..94e6b7a 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -29,6 +29,7 @@ pub mod ephemeral; pub mod filter; pub mod keys; pub mod migration; +pub mod pds_tiers; pub mod types; use tokio::sync::broadcast; @@ -38,7 +39,6 @@ fn default_opts() -> KeyspaceCreateOptions { KeyspaceCreateOptions::default() } -#[derive(Clone)] pub struct Db { pub inner: Arc, pub path: std::path::PathBuf, diff --git a/src/db/pds_tiers.rs b/src/db/pds_tiers.rs new file mode 100644 index 0000000..5ae44d6 --- /dev/null +++ b/src/db/pds_tiers.rs @@ -0,0 +1,33 @@ +use fjall::{Keyspace, OwnedWriteBatch}; +use miette::{IntoDiagnostic, Result}; +use smol_str::SmolStr; + +pub const PDS_TIER_PREFIX: &[u8] = b"pt|"; + +// `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 +} + +/// load all PDS tier assignments from the filter keyspace +pub fn load(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))); + } + Ok(out) +} + +pub fn set(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str, tier: &str) { + batch.insert(ks, pds_tier_key(host), tier.as_bytes()); +} + +pub fn remove(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str) { + batch.remove(ks, pds_tier_key(host)); +} diff --git a/src/filter.rs b/src/filter.rs index ee3d837..a987c1e 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -1,5 +1,5 @@ -use jacquard_common::types::nsid::Nsid; use serde::{Deserialize, Serialize}; +use smol_str::SmolStr; use std::sync::Arc; pub(crate) type FilterHandle = Arc>; @@ -8,16 +8,6 @@ pub(crate) fn new_handle(config: FilterConfig) -> FilterHandle { Arc::new(arc_swap::ArcSwap::new(Arc::new(config))) } -/// apply a bool patch or set replacement for a single set update. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(untagged)] -pub(crate) enum SetUpdate { - /// replace the entire set with this list - Set(Vec), - /// patch: true = add, false = remove - Patch(std::collections::HashMap), -} - #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum FilterMode { @@ -28,8 +18,8 @@ pub enum FilterMode { #[derive(Debug, Clone, Serialize)] pub(crate) struct FilterConfig { pub mode: FilterMode, - pub signals: Vec>, - pub collections: Vec>, + pub signals: Vec, + pub collections: Vec, } impl FilterConfig { diff --git a/src/ingest/firehose.rs b/src/ingest/firehose.rs index f8035d3..c6a218a 100644 --- a/src/ingest/firehose.rs +++ b/src/ingest/firehose.rs @@ -3,6 +3,7 @@ use crate::ingest::stream::{FirehoseError, FirehoseStream, SubscribeReposMessage use crate::ingest::{BufferTx, IngestMessage}; use crate::state::AppState; use crate::util::WatchEnabledExt; +use crate::util::throttle::ThrottleHandle; use jacquard_common::IntoStatic; use jacquard_common::types::did::Did; use miette::{IntoDiagnostic, Result}; @@ -21,10 +22,11 @@ pub struct FirehoseIngestor { filter: FilterHandle, enabled: watch::Receiver, _verify_signatures: bool, + throttle: ThrottleHandle, } impl FirehoseIngestor { - pub fn new( + pub async fn new( state: Arc, buffer_tx: BufferTx, relay_host: Url, @@ -33,6 +35,7 @@ impl FirehoseIngestor { enabled: watch::Receiver, verify_signatures: bool, ) -> Self { + let throttle = state.throttler.get_handle(&relay_host).await; Self { state, buffer_tx, @@ -41,11 +44,16 @@ impl FirehoseIngestor { filter, enabled, _verify_signatures: verify_signatures, + throttle, } } #[tracing::instrument(skip(self), fields(relay = %self.relay_host))] pub async fn run(mut self) -> Result<()> { + // extract host as owned String to avoid borrow conflicts with &self inside the loop + let host = self.relay_host.host_str().unwrap_or("").to_string(); + let count_key = crate::db::keys::pds_account_count_key(&host); + loop { self.enabled.wait_enabled("firehose").await; @@ -87,7 +95,14 @@ impl FirehoseIngestor { } }; match decode_frame(&bytes) { - Ok(msg) => self.handle_message(msg).await, + Ok(msg) => { + if self.is_pds { + let accounts = self.state.db.get_count(&count_key).await; + let tier = self.state.pds_tier_for(&host); + self.throttle.wait_for_allow(accounts, &tier).await; + } + self.handle_message(msg).await + }, Err(e) => { match e { // dont disconnect on unknown op or type diff --git a/src/ingest/relay.rs b/src/ingest/relay.rs index d38a0b0..7459d1e 100644 --- a/src/ingest/relay.rs +++ b/src/ingest/relay.rs @@ -274,7 +274,7 @@ impl RelayWorker { } SubscribeReposMessage::Account(account) => { debug!("processing account"); - Self::handle_account(ctx, &mut repo_state, &msg.firehose, *account) + Self::handle_account(ctx, &mut repo_state, &msg.firehose, *account, msg.is_pds) } _ => Ok(()), } @@ -472,8 +472,9 @@ impl RelayWorker { fn handle_account( ctx: &mut WorkerContext, repo_state: &mut RepoState, - #[allow(unused_variables)] firehose: &Url, + firehose: &Url, #[allow(unused_mut)] mut account: Account<'static>, + is_pds: bool, ) -> Result<()> { let event_ms = account.time.0.timestamp_millis(); if repo_state.last_message_time.is_some_and(|t| event_ms <= t) { @@ -483,8 +484,10 @@ impl RelayWorker { repo_state.advance_message_time(event_ms); + // always capture was_active for count tracking, not just in indexer mode + let was_active = repo_state.active; #[cfg(feature = "indexer")] - let (was_active, was_status) = (repo_state.active, repo_state.status.clone()); + let was_status = repo_state.status.clone(); repo_state.active = account.active; if !account.active { @@ -512,6 +515,18 @@ 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); + if !was_active && repo_state.active { + ctx.state.db.update_count(&count_key, 1); + } else if was_active && !repo_state.active { + ctx.state.db.update_count(&count_key, -1); + } + } + } + let repo_key = keys::repo_key(&account.did); #[cfg(feature = "indexer")] @@ -558,14 +573,20 @@ impl WorkerContext<'_> { repo_state: &mut RepoState, source_host: &str, ) -> Result { - let expected = pds_host(repo_state.pds.as_deref()); + let pds_host = |pds: &str| { + Url::parse(pds) + .ok() + .and_then(|u| u.host_str().map(SmolStr::new)) + }; + + let expected = repo_state.pds.as_deref().and_then(pds_host); if expected.as_deref() == Some(source_host) { return Ok(AuthorityOutcome::Authorized); } // try again once self.refresh_doc(did, repo_state)?; - let Some(expected) = pds_host(repo_state.pds.as_deref()) else { + let Some(expected) = repo_state.pds.as_deref().and_then(pds_host) else { miette::bail!("can't get pds host???"); }; @@ -837,6 +858,13 @@ impl WorkerContext<'_> { db.update_count("repos", 1); + // 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); + } + } + Ok(Some(repo_state)) } @@ -861,12 +889,3 @@ enum AuthorityOutcome { /// host did not match even after doc resolution. WrongHost { expected: SmolStr }, } - -fn pds_host(pds: Option<&str>) -> Option { - // todo: add faster host parsing since we only need that - pds.and_then(|pds| Url::parse(pds).ok()).map(|u| { - u.host_str() - .map(SmolStr::new) - .expect("that there is host in pds url") - }) -} diff --git a/src/lib.rs b/src/lib.rs index 7f119b5..779eb00 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,6 +18,7 @@ pub(crate) mod db; pub(crate) mod ingest; #[cfg(feature = "indexer")] pub(crate) mod ops; +pub(crate) mod patch; pub(crate) mod resolver; pub(crate) mod state; pub(crate) mod util; diff --git a/src/patch.rs b/src/patch.rs new file mode 100644 index 0000000..a74c239 --- /dev/null +++ b/src/patch.rs @@ -0,0 +1,11 @@ +use serde::{Deserialize, Serialize}; + +/// apply a bool patch or set replacement for a single set update. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum SetUpdate { + /// replace the entire set with this list + Set(Vec), + /// patch: true = add, false = remove + Patch(std::collections::HashMap), +} diff --git a/src/state.rs b/src/state.rs index f2a20f1..1fb35ac 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,27 +1,38 @@ +use std::collections::HashMap; +use std::sync::Arc; use std::sync::atomic::AtomicI64; use std::time::Duration; +use arc_swap::ArcSwap; use miette::Result; +use smol_str::SmolStr; use tokio::sync::{Notify, watch}; use url::Url; use crate::{ - config::Config, + config::{Config, RateTier}, db::Db, filter::{FilterHandle, new_handle}, resolver::Resolver, + util::throttle::Throttler, }; +/// pds hostname -> tier name. updated atomically via ArcSwap. +pub(crate) type PdsTierHandle = Arc>>; + pub struct AppState { pub db: Db, pub resolver: Resolver, pub(crate) filter: FilterHandle, + pub(crate) pds_tiers: PdsTierHandle, + pub(crate) rate_tiers: HashMap, pub firehose_cursors: scc::HashIndex, pub backfill_notify: Notify, pub crawler_enabled: watch::Sender, pub firehose_enabled: watch::Sender, pub backfill_enabled: watch::Sender, pub ephemeral_ttl: Duration, + pub throttler: Throttler, } impl AppState { @@ -41,6 +52,21 @@ impl AppState { let filter = new_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 tier_map: HashMap = crate::db::pds_tiers::load(&db.filter) + .unwrap_or_default() + .into_iter() + .map(|(host, tier)| (host.to_string(), tier)) + .collect(); + for host in &config.trusted_hosts { + tier_map + .entry(host.clone()) + .or_insert_with(|| SmolStr::new("trusted")); + } + let pds_tiers = Arc::new(ArcSwap::new(Arc::new(tier_map))); + let relay_cursors = scc::HashIndex::new(); let (crawler_enabled, _) = watch::channel(crawler_default); @@ -51,12 +77,15 @@ impl AppState { db, resolver, filter, + pds_tiers, + rate_tiers: config.rate_tiers.clone(), firehose_cursors: relay_cursors, backfill_notify: Notify::new(), crawler_enabled, firehose_enabled, backfill_enabled, ephemeral_ttl: config.ephemeral_ttl.clone(), + throttler: Throttler::new(), }) } @@ -64,6 +93,21 @@ impl AppState { self.backfill_notify.notify_one(); } + /// returns the rate tier for the given PDS hostname. + /// falls back to the "default" tier if no assignment exists or the assigned tier is unknown. + pub fn pds_tier_for(&self, host: &str) -> RateTier { + let default = self + .rate_tiers + .get("default") + .copied() + .unwrap_or_else(RateTier::default_tier); + let snapshot = self.pds_tiers.load(); + snapshot + .get(host) + .and_then(|name| self.rate_tiers.get(name.as_str()).copied()) + .unwrap_or(default) + } + /// pauses the crawler, firehose, and backfill worker, runs `f`, then restores their prior state. /// the restore always happens, even if `f` returns an error. pub async fn with_ingestion_paused(&self, f: F) -> T diff --git a/src/util.rs b/src/util/mod.rs similarity index 99% rename from src/util.rs rename to src/util/mod.rs index 42e0e27..e934b98 100644 --- a/src/util.rs +++ b/src/util/mod.rs @@ -10,6 +10,8 @@ use url::Url; use crate::{db::types::DidKey, types::RepoStatus}; +pub mod throttle; + /// outcome of [`RetryWithBackoff::retry`] when the operation does not succeed. pub enum RetryOutcome { /// ratelimited after exhausting all retries diff --git a/src/crawler/throttle.rs b/src/util/throttle.rs similarity index 61% rename from src/crawler/throttle.rs rename to src/util/throttle.rs index 74f4fb5..2525645 100644 --- a/src/crawler/throttle.rs +++ b/src/util/throttle.rs @@ -1,8 +1,10 @@ +use crate::config::RateTier; +use parking_lot::Mutex; use scc::HashMap; use std::future::Future; use std::sync::Arc; use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering}; -use std::time::Duration; +use std::time::{Duration, Instant}; use tokio::sync::{Notify, Semaphore, SemaphorePermit}; use url::Url; @@ -10,6 +12,13 @@ use url::Url; /// ref pds allows 10 requests per second... so 10 should be fine const PER_PDS_CONCURRENCY: usize = 10; +// per second, hour and day +const DURATIONS: [Duration; 3] = [ + Duration::from_secs(1), + Duration::from_secs(3600), + Duration::from_secs(86400), +]; + #[derive(Clone)] pub struct Throttler { states: Arc>>, @@ -54,6 +63,7 @@ struct State { /// let tasks exit naturally, deferring to the background retry loop. failure_notify: Notify, semaphore: Semaphore, + rate_limiter: RateLimiter, } impl State { @@ -64,6 +74,7 @@ impl State { consecutive_timeouts: AtomicUsize::new(0), failure_notify: Notify::new(), semaphore: Semaphore::new(PER_PDS_CONCURRENCY), + rate_limiter: RateLimiter::new(), } } } @@ -92,9 +103,6 @@ impl ThrottleHandle { /// called on a 429 response. `retry_after_secs` comes from the `Retry-After` /// header if present; falls back to 60s. uses `fetch_max` so concurrent callers /// don't race each other back to a shorter window. - /// - /// deliberately does NOT notify waiters — 429s are soft and tasks should exit - /// naturally via the `Retry` result rather than being cancelled. pub fn record_ratelimit(&self, retry_after_secs: Option) { let secs = retry_after_secs.unwrap_or(60) as i64; let until = chrono::Utc::now().timestamp() + secs; @@ -155,7 +163,6 @@ impl ThrottleHandle { } /// resolves when this PDS gets a hard failure notification. - /// used by `or_throttle` and the semaphore acquire select to cancel in-flight work. pub async fn wait_for_failure(&self) { loop { let notified = self.state.failure_notify.notified(); @@ -165,6 +172,108 @@ impl ThrottleHandle { notified.await; } } + + /// waits until the rate tier's limits allow more events for this PDS. + /// sleeps precisely until the most restrictive window opens rather than polling. + pub async fn wait_for_allow(&self, num_accounts: u64, tier: &RateTier) { + let limits = limits_for(num_accounts, tier); + while let Some(wait) = self.state.rate_limiter.try_acquire(limits) { + tokio::time::sleep(wait).await; + } + } +} + +fn limits_for(num_accounts: u64, tier: &RateTier) -> [u64; 3] { + let per_sec = tier + .per_second_base + .max((num_accounts as f64 * tier.per_second_account_mul) as u64); + [per_sec, tier.per_hour, tier.per_day] +} + +struct WindowState { + count: u64, + prev_count: u64, + window_start: Instant, +} + +impl WindowState { + fn new() -> Self { + Self { + count: 0, + prev_count: 0, + window_start: Instant::now(), + } + } + + fn rotate(&mut self, dur: Duration) { + let elapsed = self.window_start.elapsed(); + if elapsed >= dur { + let n = (elapsed.as_nanos() / dur.as_nanos()).max(1) as u32; + self.prev_count = if n == 1 { self.count } else { 0 }; + self.count = 0; + self.window_start += dur * n; + } + } + + /// returns how long to sleep before this window would allow one more event. + /// Duration::ZERO means allow now. + fn wait_needed(&self, dur: Duration, limit: u64) -> Duration { + let elapsed = self.window_start.elapsed(); + let remaining = dur.saturating_sub(elapsed); + let weight = remaining.as_secs_f64() / dur.as_secs_f64(); + let effective = self.count as f64 + self.prev_count as f64 * weight; + + if effective < limit as f64 { + return Duration::ZERO; + } + + if self.prev_count == 0 || self.count as f64 >= limit as f64 { + // must wait for a full window rotation + remaining + Duration::from_millis(1) + } else { + let secs = remaining.as_secs_f64() + - dur.as_secs_f64() * (limit as f64 - self.count as f64) / self.prev_count as f64; + Duration::from_secs_f64(secs.max(0.0)) + Duration::from_micros(500) + } + } +} + +struct RateLimiter { + // parking_lot::Mutex — uncontended path never touches the kernel + windows: Mutex<[WindowState; 3]>, +} + +impl RateLimiter { + fn new() -> Self { + Self { + windows: Mutex::new([WindowState::new(), WindowState::new(), WindowState::new()]), + } + } + + /// returns None if the slot was acquired, or Some(sleep_for) if limited. + fn try_acquire(&self, limits: [u64; 3]) -> Option { + let mut windows = self.windows.lock(); + + windows + .iter_mut() + .zip(DURATIONS) + .for_each(|(w, d)| w.rotate(d)); + + let max_wait = windows + .iter() + .zip(DURATIONS) + .zip(limits) + .map(|((w, dur), limit)| w.wait_needed(dur, limit)) + .max() + .unwrap_or(Duration::ZERO); + + if max_wait.is_zero() { + windows.iter_mut().for_each(|w| w.count += 1); + None + } else { + Some(max_wait) + } + } } /// adds a method for racing the future against a hard-failure notification. diff --git a/tests/api.nu b/tests/api.nu index 3022ab3..5d4a1f6 100644 --- a/tests/api.nu +++ b/tests/api.nu @@ -250,6 +250,281 @@ def test-firehose-sources [url: string, pid: int] { print "firehose source tests passed!" } +def test-pds-tiers [url: string, pid: int] { + print "=== test: pds tier management ===" + + # initial state: no assignments, built-in rate tiers present + print " GET /pds/tiers (expect empty assignments, built-in rate_tiers)..." + let initial = (http get $"($url)/pds/tiers") + if ($initial.assignments | length) != 0 { + fail $"expected empty assignments, got ($initial.assignments | length)" $pid + } + if not ("default" in $initial.rate_tiers) { + fail "expected 'default' tier in rate_tiers" $pid + } + if not ("trusted" in $initial.rate_tiers) { + fail "expected 'trusted' tier in rate_tiers" $pid + } + print " ok: empty assignments and built-in rate tiers present" + + # GET /pds/rate-tiers returns the same definitions with the right fields + print " GET /pds/rate-tiers (check structure)..." + 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"] { + if not ($field in $tier) { + fail $"($tier_name) tier missing field ($field)" $pid + } + } + } + # trusted tier must have higher per-second limit than default + if ($rate_tiers.trusted.per_second_base) <= ($rate_tiers.default.per_second_base) { + fail $"expected trusted.per_second_base > default, got ($rate_tiers.trusted.per_second_base) vs ($rate_tiers.default.per_second_base)" $pid + } + print " ok: rate tier definitions have correct fields and expected ordering" + + # assign a host to the trusted tier + print " PUT /pds/tiers (assign to trusted)..." + http put -f -e -t application/json $"($url)/pds/tiers" { + host: "pds.example.com", + tier: "trusted" + } | assert-status 200 "PUT /pds/tiers" $pid + let after_assign = (http get $"($url)/pds/tiers") + if ($after_assign.assignments | length) != 1 { + fail $"expected 1 assignment, got ($after_assign.assignments | length)" $pid + } + let a = ($after_assign.assignments | first) + if $a.host != "pds.example.com" { + fail $"expected host=pds.example.com, got ($a.host)" $pid + } + if $a.tier != "trusted" { + fail $"expected tier=trusted, got ($a.tier)" $pid + } + print $" ok: assignment created host=($a.host), tier=($a.tier)" + + # re-assigning the same host to a different tier updates without creating a duplicate + print " PUT /pds/tiers (re-assign to default)..." + http put -f -e -t application/json $"($url)/pds/tiers" { + host: "pds.example.com", + tier: "default" + } | assert-status 200 "PUT /pds/tiers re-assign" $pid + let after_reassign = (http get $"($url)/pds/tiers") + if ($after_reassign.assignments | length) != 1 { + fail $"expected 1 assignment after re-assign, got ($after_reassign.assignments | length)" $pid + } + if ($after_reassign.assignments | first).tier != "default" { + fail $"expected tier=default after re-assign, got (($after_reassign.assignments | first).tier)" $pid + } + print " ok: re-assign updates tier without creating a duplicate" + + # assigning an unknown tier name is rejected with 400 + print " PUT /pds/tiers (unknown tier, expect 400)..." + http put -f -e -t application/json $"($url)/pds/tiers" { + host: "pds.example.com", + tier: "nonexistent" + } | assert-status 400 "PUT /pds/tiers unknown tier" $pid + let after_bad = (http get $"($url)/pds/tiers") + if ($after_bad.assignments | length) != 1 { + fail "expected assignment count unchanged after rejected request" $pid + } + if ($after_bad.assignments | first).tier != "default" { + fail "expected tier unchanged after rejected request" $pid + } + print " ok: unknown tier name rejected with 400, existing assignment unchanged" + + # add a second host to verify multi-assignment listing works + print " PUT /pds/tiers (second host)..." + http put -f -e -t application/json $"($url)/pds/tiers" { + host: "other.example.com", + tier: "trusted" + } | assert-status 200 "PUT /pds/tiers second host" $pid + let after_second = (http get $"($url)/pds/tiers") + if ($after_second.assignments | length) != 2 { + fail $"expected 2 assignments, got ($after_second.assignments | length)" $pid + } + print " ok: two distinct hosts listed independently" + + # 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 + let after_del = (http get $"($url)/pds/tiers") + if ($after_del.assignments | length) != 1 { + fail $"expected 1 assignment after delete, got ($after_del.assignments | length)" $pid + } + if ($after_del.assignments | first).host != "other.example.com" { + fail "expected only other.example.com to remain after delete" $pid + } + 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 + + # 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 + let after_idempotent = (http get $"($url)/pds/tiers") + if ($after_idempotent.assignments | length) != 0 { + fail "expected empty assignments after cleanup" $pid + } + print " ok: delete of non-existent host is idempotent" + + print "pds tier management tests passed!" +} + +# verify that tier assignments are written to the database and survive a restart. +def test-pds-tier-persistence [binary: string, db_path: string, port: int] { + print "=== test: pds tier assignments persist across restart ===" + + let url = $"http://localhost:($port)" + + let instance = (with-env { HYDRANT_CRAWLER_URLS: "", HYDRANT_RELAY_HOSTS: "" } { + start-hydrant $binary $db_path $port + }) + if not (wait-for-api $url) { + fail "hydrant did not start" + } + + print " assigning host to trusted tier..." + http put -t application/json $"($url)/pds/tiers" { + host: "persist.example.com", + tier: "trusted" + } + + let before = (http get $"($url)/pds/tiers") + if ($before.assignments | length) != 1 { + fail "assignment was not created" $instance.pid + } + + print " restarting hydrant..." + kill $instance.pid + sleep 2sec + + let instance2 = (with-env { HYDRANT_CRAWLER_URLS: "", HYDRANT_RELAY_HOSTS: "" } { + start-hydrant $binary $db_path $port + }) + if not (wait-for-api $url) { + fail "hydrant did not restart" $instance2.pid + } + + print " checking assignment survived restart..." + let after = (http get $"($url)/pds/tiers") + if ($after.assignments | length) != 1 { + fail $"expected 1 assignment after restart, got ($after.assignments | length)" $instance2.pid + } + let a = ($after.assignments | first) + if $a.host != "persist.example.com" { + fail $"expected host=persist.example.com after restart, got ($a.host)" $instance2.pid + } + if $a.tier != "trusted" { + fail $"expected tier=trusted after restart, got ($a.tier)" $instance2.pid + } + print " ok: tier assignment persisted across restart" + + kill $instance2.pid + print "pds tier persistence test passed!" +} + +# verify that HYDRANT_TRUSTED_HOSTS pre-assigns hosts to the trusted tier at startup. +def test-pds-trusted-hosts [binary: string, db_path: string, port: int] { + print "=== test: HYDRANT_TRUSTED_HOSTS pre-assigns tier at startup ===" + + let url = $"http://localhost:($port)" + let host_a = "alpha.example.com" + let host_b = "beta.example.com" + + let instance = (with-env { + HYDRANT_CRAWLER_URLS: "", + HYDRANT_RELAY_HOSTS: "", + HYDRANT_TRUSTED_HOSTS: $"($host_a),($host_b)" + } { + start-hydrant $binary $db_path $port + }) + if not (wait-for-api $url) { + fail "hydrant did not start" + } + + print " checking pre-assigned trusted hosts..." + let tiers = (http get $"($url)/pds/tiers") + let assignments = $tiers.assignments + + for host in [$host_a, $host_b] { + let match = ($assignments | where host == $host) + if ($match | length) != 1 { + fail $"expected assignment for ($host) from HYDRANT_TRUSTED_HOSTS, got ($assignments)" $instance.pid + } + if ($match | first).tier != "trusted" { + fail $"expected tier=trusted for ($host), got (($match | first).tier)" $instance.pid + } + } + print $" ok: ($host_a) and ($host_b) pre-assigned to trusted tier" + + kill $instance.pid + print "trusted hosts startup test passed!" +} + +# verify that a custom tier defined via HYDRANT_RATE_TIERS is visible and assignable. +def test-pds-custom-rate-tier [binary: string, db_path: string, port: int] { + print "=== test: custom rate tier via HYDRANT_RATE_TIERS ===" + + let url = $"http://localhost:($port)" + + # custom:100/1.0/360000/8640000 — base=100, mul=1.0, hourly=360000, daily=8640000 + let instance = (with-env { + HYDRANT_CRAWLER_URLS: "", + HYDRANT_RELAY_HOSTS: "", + HYDRANT_RATE_TIERS: "custom:100/1.0/360000/8640000" + } { + start-hydrant $binary $db_path $port + }) + if not (wait-for-api $url) { + fail "hydrant did not start" + } + + # custom tier should appear alongside the built-in tiers + print " checking custom tier is listed in /pds/rate-tiers..." + let rate_tiers = (http get $"($url)/pds/rate-tiers") + if not ("custom" in $rate_tiers) { + fail "expected 'custom' tier in rate_tiers" $instance.pid + } + if not ("default" in $rate_tiers) { + fail "built-in 'default' tier should still be present alongside custom tier" $instance.pid + } + let custom = ($rate_tiers | get custom) + if $custom.per_second_base != 100 { + fail $"expected custom.per_second_base=100, got ($custom.per_second_base)" $instance.pid + } + if $custom.per_hour != 360000 { + fail $"expected custom.per_hour=360000, got ($custom.per_hour)" $instance.pid + } + print $" ok: custom tier listed with correct parameters" + + # a host can be assigned to the custom tier + print " assigning host to custom tier..." + http put -f -e -t application/json $"($url)/pds/tiers" { + host: "custom.example.com", + tier: "custom" + } | assert-status 200 "PUT /pds/tiers custom tier" $instance.pid + let after = (http get $"($url)/pds/tiers") + let match = ($after.assignments | where host == "custom.example.com") + if ($match | length) != 1 { + fail "expected assignment for custom.example.com" $instance.pid + } + if ($match | first).tier != "custom" { + fail $"expected tier=custom, got (($match | first).tier)" $instance.pid + } + print " ok: host assigned to custom tier successfully" + + kill $instance.pid + print "custom rate tier test passed!" +} + def main [] { let port = resolve-test-port 3007 let url = $"http://localhost:($port)" @@ -268,6 +543,7 @@ def main [] { test-crawler-sources $url $instance.pid test-firehose-sources $url $instance.pid + test-pds-tiers $url $instance.pid kill $instance.pid sleep 2sec @@ -282,6 +558,24 @@ def main [] { print $"db: ($db_config)" test-config-source-not-persisted $binary $db_config $port + sleep 1sec + + let db_pds_persist = (mktemp -d -t hydrant_api.XXXXXX) + print $"db: ($db_pds_persist)" + test-pds-tier-persistence $binary $db_pds_persist $port + + sleep 1sec + + let db_pds_trusted = (mktemp -d -t hydrant_api.XXXXXX) + print $"db: ($db_pds_trusted)" + test-pds-trusted-hosts $binary $db_pds_trusted $port + + sleep 1sec + + let db_pds_custom = (mktemp -d -t hydrant_api.XXXXXX) + print $"db: ($db_pds_custom)" + test-pds-custom-rate-tier $binary $db_pds_custom $port + print "" print "all api tests passed!" } -- 2.51.2