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!"
}