Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113use std::sync::Arc;#[cfg(test)]use std::sync::Mutex;use std::sync::atomic::{AtomicUsize, Ordering};use std::time::Duration;
use miette::{IntoDiagnostic, Result};use rand::RngExt;#[cfg(test)]use tokio::sync::Notify;use tokio_util::sync::CancellationToken;use tracing::{debug, error, info};use url::Url;
use crate::config::FirehoseSource;use crate::db::{self, keys};#[cfg(feature = "firehose-diagnostics")]use crate::ingest::firehose_stats::{FirehoseStatsSnapshot, RelayWorkerStatsSnapshot};use crate::ingest::{BufferTx, firehose::FirehoseIngestor};use crate::pds_daily_limit::PdsDailyReservation;use crate::pds_discovery::{PdsHostError, canonical_pds_url};use crate::state::AppState;
pub(super) struct FirehoseIngestorHandle { id: usize, cancel: CancellationToken,}
impl Drop for FirehoseIngestorHandle { fn drop(&mut self) { self.cancel.cancel(); }}
pub(super) struct FirehoseShared { pub(super) buffer_tx: BufferTx, pub(super) verify_signatures: bool, pub(super) max_failures: usize,}
/// a successful validated PDS admission.#[derive(Debug, Clone, Copy, PartialEq, Eq)]pub(crate) enum PdsAdmission { StartedNew, RestartedKnown, AlreadyRunning,}
/// a public policy rejection from validated PDS admission.#[derive(Debug, Clone, PartialEq, Eq)]pub(crate) enum PdsAdmissionRejection { InvalidHost(PdsHostError), HostBanned, DailyLimitReached,}
pub(crate) type PdsAdmissionResult = std::result::Result<PdsAdmission, PdsAdmissionRejection>;
pub(super) struct SeededSourceAdds { pub(super) added: usize, pub(super) known: usize, pub(super) rejected: usize,}
/// a snapshot of a single firehose relay's runtime state.#[derive(Debug, Clone, serde::Serialize)]pub struct FirehosePdsInfo { pub host: String, pub seq: i64, pub account_count: u64, pub status: &'static str,}
/// details for the most recent recorded firehose source failure.#[derive(Debug, Clone, serde::Serialize)]pub struct FirehoseFailureInfo { pub at: i64, pub kind: String, pub detail: String,}
/// a snapshot of a single firehose relay's runtime state.#[derive(Debug, Clone, serde::Serialize)]pub struct FirehoseSourceInfo { pub url: Url, /// true when this is a direct PDS connection; enables host authority enforcement. pub is_pds: bool, /// true when outbound sockets enforce the public-address boundary. pub public: bool, pub running: bool, pub failing: bool, pub throttled: bool, pub consecutive_failures: usize, #[serde(skip_serializing_if = "Option::is_none")] pub throttled_until: Option<i64>, #[serde(skip_serializing_if = "Option::is_none")] pub retry_in_secs: Option<u64>, #[serde(skip_serializing_if = "Option::is_none")] pub last_failure: Option<FirehoseFailureInfo>, #[cfg(feature = "firehose-diagnostics")] #[serde(skip_serializing_if = "Option::is_none")] pub stats: Option<FirehoseStatsSnapshot>, #[serde(skip_serializing_if = "Option::is_none")] pub host_status: Option<&'static str>, #[serde(skip_serializing_if = "Option::is_none")] pub pds: Option<FirehosePdsInfo>,}
/// feature-gated runtime diagnostics for firehose ingestion internals.#[cfg(feature = "firehose-diagnostics")]#[derive(Debug, Clone, serde::Serialize)]pub struct FirehoseDiagnosticsInfo { pub relay_worker: RelayWorkerStatsSnapshot,}
/// runtime control over the firehose ingestor component.#[derive(Clone)]pub struct FirehoseHandle { pub(super) state: Arc<AppState>, /// set once by [`Hydrant::run`]; `None` means run() has not been called yet. pub(super) shared: Arc<std::sync::OnceLock<FirehoseShared>>, /// per-relay running tasks, keyed by url. pub(super) tasks: Arc<scc::HashMap<Url, FirehoseIngestorHandle>>, /// known source urls → is_pds flag; includes API-added (db-persisted) and static config sources. pub(super) known_sources: Arc<scc::HashMap<Url, bool>>, /// Sources admitted through the public boundary. Persisted legacy rows /// without a trusted marker are inserted here; absence is reserved for /// current operator/admin sources. pub(super) public_sources: Arc<scc::HashMap<Url, bool>>, /// ids assigned to spawned tasks next_task_id: Arc<AtomicUsize>, #[cfg(test)] admission_pause: Arc<Mutex<Option<(Arc<Notify>, Arc<Notify>)>>>, #[cfg(test)] spawn_enabled: Arc<std::sync::atomic::AtomicBool>,}
impl FirehoseHandle { pub(super) fn new(state: Arc<AppState>) -> Self { Self { state, shared: Arc::new(std::sync::OnceLock::new()), tasks: Arc::new(scc::HashMap::new()), known_sources: Arc::new(scc::HashMap::new()), public_sources: Arc::new(scc::HashMap::new()), next_task_id: Arc::new(AtomicUsize::new(0)), #[cfg(test)] admission_pause: Arc::new(Mutex::new(None)), #[cfg(test)] spawn_enabled: Arc::new(std::sync::atomic::AtomicBool::new(true)), } }
pub(super) async fn spawn_firehose_ingestor( &self, source: &FirehoseSource, shared: &FirehoseShared, delay_startup: bool, ) -> Result<()> { let public = self.public_sources.contains_sync(&source.url); let prepared = self .prepare_firehose_ingestor(source, shared, public) .await?; self.spawn_prepared_firehose_ingestor(source, prepared, delay_startup) .await; Ok(()) }
async fn prepare_firehose_ingestor( &self, source: &FirehoseSource, shared: &FirehoseShared, public: bool, ) -> Result<(Option<i64>, FirehoseIngestor)> { let start = db::get_firehose_cursor(&self.state.db, &source.url).await?; let ingestor = FirehoseIngestor::new( self.state.clone(), shared.buffer_tx.clone(), source.url.clone(), source.is_pds, public, self.state.filter.clone(), self.state.firehose_enabled.subscribe(), shared.verify_signatures, shared.max_failures, ) .await; Ok((start, ingestor)) }
async fn spawn_prepared_firehose_ingestor( &self, source: &FirehoseSource, (start, ingestor): (Option<i64>, FirehoseIngestor), delay_startup: bool, ) { #[cfg(test)] if !self.spawn_enabled.load(Ordering::Relaxed) { drop(ingestor); return; } use std::sync::atomic::AtomicI64;
// existing in-memory cursor state takes precedence. let _ = self .state .firehose_cursors .insert_async(source.url.clone(), AtomicI64::new(start.unwrap_or(0))) .await; info!(relay = %source.url, source.is_pds, cursor = ?start, "starting firehose ingestor");
let id = self.next_task_id.fetch_add(1, Ordering::Relaxed); let cancel = CancellationToken::new(); tokio::spawn({ let relay_url = source.url.clone(); let is_pds = source.is_pds; let tasks = self.tasks.clone(); let token = cancel.clone(); async move { // jitter connection start so a seeded page does not create a thundering herd. if delay_startup { let max_jitter_ms = if is_pds && !cfg!(debug_assertions) { 60_000 } else { 2_000 }; let jitter_ms = rand::rng().random_range(0u64..max_jitter_ms); debug!( relay = %relay_url, is_pds, jitter_ms, "delaying firehose ingestor startup" ); tokio::select! { _ = tokio::time::sleep(Duration::from_millis(jitter_ms)) => {} _ = token.cancelled() => { info!(relay = %relay_url, "firehose ingestor cancelled"); return; } } } tokio::select! { res = ingestor.run() => { // only remove our own entry because an upsert can replace us. tasks.remove_if_async(&relay_url, |h| h.id == id).await; match res { Ok(()) => info!(relay = %relay_url, "firehose shut down!"), Err(e) => error!(relay = %relay_url, err = %e, "firehose ingestor exited with error"), } }, _ = token.cancelled() => { info!(relay = %relay_url, "firehose ingestor cancelled"); } } } });
let handle = FirehoseIngestorHandle { id, cancel }; // `upsert_async` returns the old handle. ignoring it drops and cancels the old task. self.tasks.upsert_async(source.url.clone(), handle).await; }
#[cfg(test)] pub(crate) fn disable_spawns_for_test(&self) { self.spawn_enabled.store(false, Ordering::Relaxed); }
#[cfg(test)] pub(crate) fn startable_for_test(&self) { let (buffer_tx, _) = BufferTx::channel(1); assert!( self.shared .set(FirehoseShared { buffer_tx, verify_signatures: false, max_failures: 1, }) .is_ok(), "test firehose shared must be set once" ); }
#[cfg(test)] pub(crate) fn pause_validated_admission_after_ban_check(&self) -> (Arc<Notify>, Arc<Notify>) { let pause = (Arc::new(Notify::new()), Arc::new(Notify::new())); *self .admission_pause .lock() .expect("admission test pause lock must not be poisoned") = Some(pause.clone()); pause }
#[cfg(test)] async fn pause_validated_admission_after_ban_check_if_requested(&self) { let pause = self .admission_pause .lock() .expect("admission test pause lock must not be poisoned") .clone(); if let Some((checked, resume)) = pause { checked.notify_one(); resume.notified().await; } }
/// enable firehose ingestion, no-op if already enabled. pub fn enable(&self) { self.state.firehose_enabled.send_replace(true); } /// 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 firehose ingestion. pub fn is_enabled(&self) -> bool { *self.state.firehose_enabled.borrow() }
/// returns `true` if this URL is already a known firehose source. /// either currently running or persisted (e.g. the host is offline but was previously added). pub fn is_source_known(&self, url: &Url) -> bool { self.known_sources.contains_sync(url) }
/// return `true` if this source has a running firehose task (eg. its not offline). pub fn is_source_running(&self, url: &Url) -> bool { self.tasks.contains_sync(url) }
/// list all known firehose sources, including offline ones pending retry. pub async fn list_sources(&self) -> Vec<FirehoseSourceInfo> { let now = chrono::Utc::now().timestamp(); let meta = self.state.pds_meta.load(); let mut out = Vec::with_capacity(self.known_sources.capacity()); self.known_sources .iter_async(|url, &is_pds| { let running = self.tasks.contains_sync(url); let throttle = self.state.throttler.snapshot(url); let pds = is_pds.then(|| self.pds_info(url, &meta)).flatten(); let host_status = pds.as_ref().map(|pds| pds.status); #[cfg(feature = "firehose-diagnostics")] let stats = self.state.firehose_stats.snapshot(url); let last_failure = throttle .last_failure .clone() .map(|failure| FirehoseFailureInfo { at: failure.at, kind: failure.kind, detail: failure.detail, }); out.push(FirehoseSourceInfo { url: url.clone(), is_pds, public: self.public_sources.contains_sync(url), running, failing: throttle.is_failing(), throttled: throttle.is_throttled(now), consecutive_failures: throttle.consecutive_failures, throttled_until: (throttle.throttled_until != 0) .then_some(throttle.throttled_until), retry_in_secs: throttle.retry_in_secs(now), last_failure, #[cfg(feature = "firehose-diagnostics")] stats, host_status, pds, }); true }) .await; out.sort_unstable_by(|a, b| a.url.as_str().cmp(b.url.as_str())); out }
#[cfg(feature = "firehose-diagnostics")] pub fn diagnostics(&self) -> FirehoseDiagnosticsInfo { FirehoseDiagnosticsInfo { relay_worker: self.state.firehose_stats.relay_worker_snapshot(), } }
fn pds_info(&self, url: &Url, meta: &crate::pds_meta::PdsMeta) -> Option<FirehosePdsInfo> { let host = url.host_str()?; let seq = self .state .firehose_cursors .peek_with(url, |_, cursor| cursor.load(Ordering::SeqCst)) .or_else(|| { self.state .db .cursors .get(keys::firehose_cursor_key_from_url(url)) .ok() .flatten() .and_then(|bytes| bytes.as_ref().try_into().ok().map(i64::from_be_bytes)) }) .unwrap_or(0); let account_count = self .state .db .get_count_sync(&keys::pds_account_count_key(host)); Some(FirehosePdsInfo { host: host.to_string(), seq, account_count, status: meta.status(host).as_str(), }) }
/// add a direct PDS source through the public discovery policy. #[cfg_attr(not(feature = "relay"), allow(dead_code))] pub(crate) async fn add_source_validated(&self, url: Url) -> Result<PdsAdmissionResult> { self.admit_validated_source(url, false).await }
async fn admit_validated_source( &self, url: Url, delay_startup: bool, ) -> Result<PdsAdmissionResult> { let (host, url) = match canonical_pds_url(&url) { Ok(source) => source, Err(error) => return Ok(Err(PdsAdmissionRejection::InvalidHost(error))), }; let _admission = self.state.pds_admission.lock().await;
if self.state.pds_meta.load().is_banned(host.as_str()) { return Ok(Err(PdsAdmissionRejection::HostBanned)); } #[cfg(test)] self.pause_validated_admission_after_ban_check_if_requested() .await; if self.is_source_running(&url) { return Ok(Ok(PdsAdmission::AlreadyRunning)); } if self.is_source_known(&url) { // Do not route this through admin add_source: that trusted override // intentionally clears public provenance. A validated request stays // public across restart and persistence. let shared = self .shared .get() .ok_or_else(|| miette::miette!("firehose worker not started"))?; let is_pds = self .known_sources .read_async(&url, |_, &is_pds| is_pds) .await .unwrap_or(true); let source = FirehoseSource { url: url.clone(), is_pds, }; let throttle = self.state.throttler.get_handle(&url).await; throttle.record_success(); let prepared = self .prepare_firehose_ingestor(&source, shared, self.public_sources.contains_sync(&url)) .await?; self.spawn_prepared_firehose_ingestor(&source, prepared, delay_startup) .await; return Ok(Ok(PdsAdmission::RestartedKnown)); }
let Some(reservation) = self.state.pds_daily_limit.reserve() else { return Ok(Err(PdsAdmissionRejection::DailyLimitReached)); }; let shared = match self.shared.get() { Some(shared) => shared, None => { self.state.pds_daily_limit.refund(reservation); return Err(miette::miette!("firehose worker not started")); } }; let source = FirehoseSource { url, is_pds: true }; // all fallible startup preparation happens before taking a quota slot's write. let prepared = match self.prepare_firehose_ingestor(&source, shared, true).await { Ok(prepared) => prepared, Err(error) => { self.state.pds_daily_limit.refund(reservation); return Err(error); } };
match self .persist_validated_source(&source, host.as_str(), reservation) .await { Ok(true) => {} Ok(false) => { self.state.pds_daily_limit.refund(reservation); return Ok(Err(PdsAdmissionRejection::HostBanned)); } Err(error) => { self.state.pds_daily_limit.refund(reservation); return Err(error); } }
let _ = self .known_sources .insert_async(source.url.clone(), true) .await; let _ = self .public_sources .insert_async(source.url.clone(), true) .await; let throttle = self.state.throttler.get_handle(&source.url).await; throttle.record_success(); self.spawn_prepared_firehose_ingestor(&source, prepared, delay_startup) .await; Ok(Ok(PdsAdmission::StartedNew)) }
async fn persist_validated_source( &self, source: &FirehoseSource, host: &str, reservation: PdsDailyReservation, ) -> Result<bool> { let key = keys::firehose_source_key(source.url.as_str()); let marker_key = keys::trusted_firehose_source_key(source.url.as_str()); let status_key = db::pds_meta::pds_status_key(host); let ban_key = db::pds_meta::pds_ban_key(host); let value = rmp_serde::to_vec(&db::FirehoseSourceMeta { is_pds: true }).map_err(|error| { miette::miette!("failed to serialize firehose source meta: {error}") })?; let source_url = source.url.clone(); self.state .db .run(move |db| { let operator_banned = db.filter.contains_key(ban_key).into_diagnostic()?; let observed_banned = db .filter .get(status_key) .into_diagnostic()? .map(|status| { rmp_serde::from_slice::<crate::pds_meta::HostStatus>(status.as_ref()) .into_diagnostic() }) .transpose()? .is_some_and(|status| status == crate::pds_meta::HostStatus::Banned); if operator_banned || observed_banned { return Ok(false); }
let mut batch = db.inner.batch(); if let Some((day, count)) = reservation.persisted() { db::stage_pds_daily_adds(&mut batch, db, day, count); } batch.insert(&db.crawler, key, value); batch.remove(&db.crawler, marker_key); batch.commit().into_diagnostic()?; if let Err(error) = db.persist() { tracing::error!( source = %source_url, err = %error, "PDS source committed but durability persist failed" ); } Ok(true) }) .await }
/// add PDS sources discovered from a seed relay. /// /// each accepted source persists its own quota reservation with its source /// record. startup stays staggered so a page cannot stampede remote hosts. pub(super) async fn add_seeded_sources(&self, urls: Vec<Url>) -> Result<SeededSourceAdds> { let mut adds = SeededSourceAdds { added: 0, known: 0, rejected: 0, }; for url in urls { if self.is_source_known(&url) { adds.known += 1; continue; } match self.admit_validated_source(url.clone(), true).await? { Ok(PdsAdmission::StartedNew) => adds.added += 1, Ok(PdsAdmission::RestartedKnown | PdsAdmission::AlreadyRunning) => { adds.known += 1; } Err(rejection) => { adds.rejected += 1; tracing::warn!(source = %url, ?rejection, "rejected seeded PDS source"); } } } Ok(adds) }
/// add a new firehose source at runtime, persisting it to the database. /// /// 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 _admission = self.state.pds_admission.lock().await; let shared = self .shared .get() .ok_or_else(|| miette::miette!("firehose worker not started"))?;
// persist to db first let key = keys::firehose_source_key(url.as_str()); let marker_key = keys::trusted_firehose_source_key(url.as_str()); self.state .db .run(move |db| { let mut batch = 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(&db.crawler, key, &value); batch.insert(&db.crawler, marker_key, []); batch.commit().into_diagnostic()?; db.persist() }) .await?; self.public_sources.remove_async(&url).await;
let _ = self.known_sources.upsert_async(url.clone(), is_pds).await;
// reset failure state so the fresh task gets a clean slate. // if the previous task exited after max failures, the failure counter // would otherwise cause the new task to exit immediately. let throttle = self.state.throttler.get_handle(&url).await; throttle.record_success();
self.spawn_firehose_ingestor(&FirehoseSource { url, is_pds }, shared, false) .await?;
Ok(()) }
/// remove a firehose source at runtime. /// /// 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<bool> { let _admission = self.state.pds_admission.lock().await; if self.known_sources.contains_async(url).await { let url_str = url.to_string(); let marker_key = keys::trusted_firehose_source_key(&url_str); self.state .db .run(move |db| { db.crawler .remove(keys::firehose_source_key(&url_str)) .into_diagnostic()?; db.crawler.remove(marker_key).into_diagnostic()?; db.persist() }) .await?; self.known_sources.remove_async(url).await; self.public_sources.remove_async(url).await; }
Ok(self.tasks.remove_async(url).await.is_some()) }
/// restart an offline firehose source without touching the database or daily limits. pub(super) async fn restart_source(&self, url: Url) -> Result<()> { let _admission = self.state.pds_admission.lock().await; let Some(is_pds) = self .known_sources .read_async(&url, |_, &is_pds| is_pds) .await else { return Ok(()); }; if self.tasks.contains_sync(&url) { return Ok(()); } let shared = self .shared .get() .ok_or_else(|| miette::miette!("firehose worker not started"))?;
// clear the failure counter so the new task isn't immediately terminated let throttle = self.state.throttler.get_handle(&url).await; throttle.record_success();
self.spawn_firehose_ingestor(&FirehoseSource { url, is_pds }, shared, true) .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); self.state .db .run(move |db| { db.cursors.remove(key).into_diagnostic()?; db.persist() }) .await?;
self.state.firehose_cursors.remove_async(&url).await;
Ok(()) }}
#[cfg(test)]mod tests { use super::*; use crate::config::Config; use crate::control::Hydrant;
fn test_config(dir: &tempfile::TempDir) -> Config { Config { database_path: dir.path().to_path_buf(), ..Default::default() } }
fn test_config_with_limit(dir: &tempfile::TempDir, new_host_limit: u64) -> Config { Config { database_path: dir.path().to_path_buf(), new_host_limit: Some(new_host_limit), ..Default::default() } }
fn startable(hydrant: &Hydrant) { hydrant.firehose.startable_for_test(); }
async fn persisted_sources(hydrant: &Hydrant) -> Result<Vec<(FirehoseSource, bool)>> { hydrant .state .db .run(crate::db::load_persisted_firehose_sources) .await }
/// `upsert_async` returns the previous handle, whose drop cancels its task. #[tokio::test] async fn replacing_a_source_cancels_the_previous_ingestor() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config(&dir)).await?; startable(&hydrant); let url = Url::parse("wss://replace.example/").into_diagnostic()?;
hydrant.firehose.add_source(url.clone(), true).await?; let (first_id, first_token) = hydrant .firehose .tasks .read_async(&url, |_, handle| (handle.id, handle.cancel.clone())) .await .expect("first source must be running"); assert!(!first_token.is_cancelled(), "a fresh handle is live");
hydrant.firehose.add_source(url.clone(), true).await?;
let (second_id, second_token) = hydrant .firehose .tasks .read_async(&url, |_, handle| (handle.id, handle.cancel.clone())) .await .expect("a source must still be running"); assert_ne!(first_id, second_id, "the replacement mints a new task id"); assert!( first_token.is_cancelled(), "replacing the handle must cancel the previous ingestor" ); assert!(!second_token.is_cancelled(), "the new ingestor stays live"); assert_eq!(hydrant.firehose.tasks.len(), 1, "one handle per url"); assert!(hydrant.firehose.is_source_running(&url));
hydrant.firehose.tasks.remove_async(&url).await; assert!(second_token.is_cancelled(), "dropping the handle stops it"); assert!(!hydrant.firehose.is_source_running(&url)); Ok(()) }
#[tokio::test] async fn admin_override_clears_public_provenance_in_memory_and_db() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config(&dir)).await?; startable(&hydrant); hydrant.firehose.disable_spawns_for_test(); let url = Url::parse("wss://public.example/").into_diagnostic()?;
assert_eq!( hydrant.firehose.add_source_validated(url.clone()).await?, Ok(PdsAdmission::StartedNew) ); assert!(hydrant.firehose.public_sources.contains_sync(&url)); assert!( persisted_sources(&hydrant) .await? .iter() .any(|(source, trusted)| { source.url == url && !trusted }) );
hydrant.firehose.add_source(url.clone(), false).await?; assert!( !hydrant.firehose.public_sources.contains_sync(&url), "trusted admin replacement must not inherit public provenance" ); assert!( persisted_sources(&hydrant) .await? .iter() .any(|(source, trusted)| { source.url == url && *trusted }) ); Ok(()) }
#[tokio::test] async fn validated_request_preserves_a_known_trusted_source() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config(&dir)).await?; startable(&hydrant); hydrant.firehose.disable_spawns_for_test(); let url = Url::parse("wss://operator.example/").into_diagnostic()?;
hydrant.firehose.add_source(url.clone(), false).await?; let throttle = hydrant.state.throttler.get_handle(&url).await; throttle.record_failure(); assert!(hydrant.state.throttler.snapshot(&url).is_failing());
assert_eq!( hydrant.firehose.add_source_validated(url.clone()).await?, Ok(PdsAdmission::RestartedKnown) ); assert!(!hydrant.state.throttler.snapshot(&url).is_failing()); assert_eq!( hydrant .firehose .known_sources .read_async(&url, |_, is_pds| *is_pds) .await, Some(false) ); assert!(!hydrant.firehose.public_sources.contains_sync(&url)); assert!( persisted_sources(&hydrant) .await? .iter() .any(|(source, trusted)| source.url == url && *trusted && !source.is_pds) ); Ok(()) }
#[tokio::test] async fn validated_admission_rejects_invalid_and_private_hosts() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config(&dir)).await?;
for url in [ "wss://127.0.0.1/", "ws://public.example/", "wss://public.example/path", ] { assert!(matches!( hydrant .firehose .add_source_validated(Url::parse(url).into_diagnostic()?) .await?, Err(PdsAdmissionRejection::InvalidHost(_)) )); } assert!(persisted_sources(&hydrant).await?.is_empty()); Ok(()) }
#[tokio::test] async fn seeded_banned_host_is_not_persisted_or_started() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config(&dir)).await?; startable(&hydrant); hydrant.pds.ban("blocked.example").await?; let url = Url::parse("wss://blocked.example/").into_diagnostic()?;
let adds = hydrant .firehose .add_seeded_sources(vec![url.clone()]) .await?;
assert_eq!(adds.added, 0); assert_eq!(adds.rejected, 1); assert!(!hydrant.firehose.is_source_known(&url)); assert!(!hydrant.firehose.is_source_running(&url)); assert!(persisted_sources(&hydrant).await?.is_empty()); Ok(()) }
#[tokio::test] async fn seeded_admission_stops_at_the_daily_limit() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?; startable(&hydrant); let first = Url::parse("wss://first-seed.example/").into_diagnostic()?; let second = Url::parse("wss://second-seed.example/").into_diagnostic()?;
let adds = hydrant .firehose .add_seeded_sources(vec![first.clone(), second.clone()]) .await?;
assert_eq!(adds.added, 1); assert_eq!(adds.rejected, 1); assert_eq!(persisted_sources(&hydrant).await?.len(), 1); assert!(matches!( crate::db::load_pds_daily_adds(&hydrant.state.db)?, Some((_, 1)) )); hydrant.firehose.tasks.remove_async(&first).await; Ok(()) }
#[tokio::test] async fn committed_admission_is_not_refunded_when_persist_fails() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?; startable(&hydrant); let url = Url::parse("wss://post-commit-persist.example/").into_diagnostic()?; hydrant.state.db.fail_next_persist_for_test();
assert_eq!( hydrant.firehose.add_source_validated(url.clone()).await?, Ok(PdsAdmission::StartedNew) ); assert!(hydrant.firehose.is_source_running(&url)); assert!(matches!( crate::db::load_pds_daily_adds(&hydrant.state.db)?, Some((_, 1)) )); assert_eq!( hydrant .firehose .add_source_validated( Url::parse("wss://second-post-commit.example/").into_diagnostic()? ) .await?, Err(PdsAdmissionRejection::DailyLimitReached) );
hydrant.firehose.tasks.remove_async(&url).await; Ok(()) }
#[tokio::test] async fn running_validated_source_does_not_spend_another_slot() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?; startable(&hydrant); let url = Url::parse("wss://running.example/").into_diagnostic()?;
assert_eq!( hydrant.firehose.add_source_validated(url.clone()).await?, Ok(PdsAdmission::StartedNew) ); assert_eq!( hydrant.firehose.add_source_validated(url.clone()).await?, Ok(PdsAdmission::AlreadyRunning) ); assert!(matches!( crate::db::load_pds_daily_adds(&hydrant.state.db)?, Some((_, 1)) )); hydrant.firehose.tasks.remove_async(&url).await; Ok(()) }
#[tokio::test] async fn legacy_raw_ban_is_enforced_and_removed_by_unban() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let cfg = test_config(&dir); let raw_host = "BANNED.EXAMPLE."; let url = Url::parse("wss://banned.example/").into_diagnostic()?;
{ let hydrant = Hydrant::new(cfg.clone()).await?; let mut batch = hydrant.state.db.inner.batch(); crate::db::pds_meta::set_status( &mut batch, &hydrant.state.db.filter, raw_host, crate::pds_meta::HostStatus::Banned, )?; batch.commit().into_diagnostic()?; crate::db::migration::rewind_version_for_test(&hydrant.state.db, 9)?; hydrant.state.db.persist()?; }
let hydrant = Hydrant::new(cfg.clone()).await?; hydrant.firehose.startable_for_test(); assert_eq!( hydrant.state.pds_meta.load().status("banned.example"), crate::pds_meta::HostStatus::Active ); assert!( hydrant .state .db .filter .contains_key(crate::db::pds_meta::pds_ban_key("banned.example")) .into_diagnostic()? ); assert_eq!( hydrant.firehose.add_source_validated(url.clone()).await?, Err(PdsAdmissionRejection::HostBanned) );
hydrant.pds.unban(raw_host).await?; assert!(!hydrant.pds.is_banned("banned.example")); assert!( !hydrant .state .db .filter .contains_key(crate::db::pds_meta::pds_ban_key("banned.example")) .into_diagnostic()? ); drop(hydrant);
let reopened = Hydrant::new(cfg).await?; assert!(!reopened.pds.is_banned("banned.example")); Ok(()) }
#[tokio::test] async fn durable_ban_prevents_validated_admission() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?; startable(&hydrant); let host = "durably-banned.example"; let url = Url::parse("wss://durably-banned.example/").into_diagnostic()?;
hydrant .state .db .run(move |db| { let mut batch = db.inner.batch(); crate::db::pds_meta::set_status( &mut batch, &db.filter, host, crate::pds_meta::HostStatus::Banned, )?; batch.commit().into_diagnostic()?; db.persist() }) .await?; assert!( !hydrant.pds.is_banned(host), "the durable check is independent of the RCU snapshot" );
assert_eq!( hydrant.firehose.add_source_validated(url.clone()).await?, Err(PdsAdmissionRejection::HostBanned) ); assert!(!hydrant.firehose.is_source_known(&url)); assert!(!hydrant.firehose.is_source_running(&url)); assert!(persisted_sources(&hydrant).await?.is_empty()); assert_eq!(crate::db::load_pds_daily_adds(&hydrant.state.db)?, None); Ok(()) }
#[tokio::test] async fn inflight_admission_starts_before_its_ban_is_published() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config(&dir)).await?; startable(&hydrant); let host = "inflight-ban.example"; let url = Url::parse("wss://inflight-ban.example/").into_diagnostic()?; let (checked, resume) = hydrant.firehose.pause_validated_admission_after_ban_check();
let admission = tokio::spawn({ let firehose = hydrant.firehose.clone(); let url = url.clone(); async move { firehose.add_source_validated(url).await } }); checked.notified().await;
let mut ban = tokio::spawn({ let pds = hydrant.pds.clone(); async move { pds.ban(host).await } }); assert!( tokio::time::timeout(Duration::from_millis(100), &mut ban) .await .is_err(), "the ban cannot publish while admission owns its lock" );
resume.notify_one(); assert_eq!( admission .await .map_err(|error| miette::miette!("admission task failed: {error}"))??, Ok(PdsAdmission::StartedNew) ); assert!(hydrant.firehose.is_source_running(&url));
ban.await .map_err(|error| miette::miette!("ban task failed: {error}"))??; assert!(hydrant.pds.is_banned(host)); hydrant.firehose.tasks.remove_async(&url).await; Ok(()) }}