use 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 super::new_pds_crawl;
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>,
#[cfg(test)]
last_spawn_delayed: 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)),
#[cfg(test)]
last_spawn_delayed: Arc::new(std::sync::atomic::AtomicBool::new(false)),
}
}
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)]
{
self.last_spawn_delayed
.store(delay_startup, Ordering::Relaxed);
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;
if !source.is_pds {
self.state.track_dispatched(&source.url).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 state = self.state.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.
if tasks.remove_if_async(&relay_url, |h| h.id == id).await.is_some() {
state.forget_dispatched(&relay_url).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.
pub(crate) async fn add_source_validated(&self, url: Url) -> Result<PdsAdmissionResult> {
self.admit_validated_source(url).await
}
async fn admit_validated_source(&self, url: Url) -> 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;
// a request that restarts a known source waits out the startup jitter too, so a host
// that keeps failing can't get an immediate connect each time it asks
if let Some(early) = self
.check_admission_preconditions(&url, host.as_str(), true)
.await?
{
#[cfg(feature = "indexer")]
if early.is_ok() {
self.request_rewalk(host.as_str()).await?;
}
return Ok(early);
}
let Some(reservation) = self.state.pds_daily_limit.reserve() else {
return Ok(Err(PdsAdmissionRejection::DailyLimitReached));
};
// a PDS asking to be followed can already hold repos, so it always gets the walk
self.spawn_admitted_source(&url, host.as_str(), false, reservation, true)
.await
}
#[cfg(feature = "indexer")]
async fn request_rewalk(&self, host: &str) -> Result<()> {
let now = chrono::Utc::now().timestamp();
let window = self.state.rewalk_window.as_secs() as i64;
let limit = self.state.rewalk_bytes;
let walked = host.to_owned();
let rewalk = self
.state
.db
.run(move |db| new_pds_crawl::request_rewalk(db, &walked, now, window, limit))
.await?;
debug!(%host, ?rewalk, "requested a rewalk");
if rewalk == new_pds_crawl::Rewalk::Queued {
new_pds_crawl::wake(&self.state);
}
Ok(())
}
async fn admit_seeded_source(&self, url: Url, queue_crawl: 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 let Some(early) = self
.check_admission_preconditions(&url, host.as_str(), true)
.await?
{
return Ok(early);
}
self.spawn_admitted_source(
&url,
host.as_str(),
true,
PdsDailyReservation::Unlimited,
queue_crawl,
)
.await
}
async fn check_admission_preconditions(
&self,
url: &Url,
host: &str,
delay_startup: bool,
) -> Result<Option<PdsAdmissionResult>> {
if self.state.pds_meta.load().is_banned(host) {
return Ok(Some(Err(PdsAdmissionRejection::HostBanned)));
}
#[cfg(test)]
self.pause_validated_admission_after_ban_check_if_requested()
.await;
// a request wakes an offline host's task instead of waiting out its backoff
if self.is_source_running(url)
&& self.state.pds_meta.load().status(host) != crate::pds_meta::HostStatus::Offline
{
return Ok(Some(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(Some(Ok(PdsAdmission::RestartedKnown)));
}
Ok(None)
}
async fn spawn_admitted_source(
&self,
url: &Url,
host: &str,
delay_startup: bool,
reservation: PdsDailyReservation,
queue_crawl: bool,
) -> Result<PdsAdmissionResult> {
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: url.clone(),
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, reservation, queue_crawl)
.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;
if queue_crawl {
new_pds_crawl::wake(&self.state);
}
Ok(Ok(PdsAdmission::StartedNew))
}
async fn persist_validated_source(
&self,
source: &FirehoseSource,
host: &str,
reservation: PdsDailyReservation,
queue_crawl: bool,
) -> 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();
let crawl_host = queue_crawl.then(|| host.to_owned());
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);
if let Some(host) = &crawl_host {
new_pds_crawl::queue(&mut batch, db, host)?;
}
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.
///
/// seeded sources bypass the public crawl request daily limit.
/// startup stays staggered so a page cannot stampede remote hosts.
/// `queue_crawl` is set once the seed is past bootstrap, see [`new_pds_crawl`].
pub(super) async fn add_seeded_sources(
&self,
urls: Vec<Url>,
queue_crawl: bool,
) -> 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_seeded_source(url.clone(), queue_crawl).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;
}
self.state.forget_dispatched(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
}
#[cfg(feature = "indexer")]
#[tokio::test]
async fn only_newly_admitted_hosts_queue_a_repo_listing() -> Result<()> {
let dir = tempfile::tempdir().into_diagnostic()?;
let hydrant = Hydrant::new(test_config(&dir)).await?;
startable(&hydrant);
hydrant.firehose.disable_spawns_for_test();
let queued = |host: &str| {
hydrant
.state
.db
.cursors
.contains_key(keys::new_pds_crawl_key(host))
.into_diagnostic()
};
let requested = Url::parse("wss://requested.example/").into_diagnostic()?;
assert_eq!(
hydrant.firehose.add_source_validated(requested).await?,
Ok(PdsAdmission::StartedNew)
);
assert!(queued("requested.example")?);
let bootstrap = Url::parse("wss://bootstrap.example/").into_diagnostic()?;
let fresh = Url::parse("wss://fresh.example/").into_diagnostic()?;
hydrant
.firehose
.add_seeded_sources(vec![bootstrap], false)
.await?;
hydrant
.firehose
.add_seeded_sources(vec![fresh.clone()], true)
.await?;
assert!(!queued("bootstrap.example")?);
assert!(queued("fresh.example")?);
// seeing a followed host again doesn't make it new
hydrant
.state
.db
.cursors
.remove(keys::new_pds_crawl_key("fresh.example"))
.into_diagnostic()?;
let adds = hydrant
.firehose
.add_seeded_sources(vec![fresh], true)
.await?;
assert_eq!(adds.known, 1);
assert!(!queued("fresh.example")?);
Ok(())
}
#[cfg(feature = "indexer")]
fn walking(hydrant: &Hydrant, host: &str) -> Result<bool> {
let cursors = &hydrant.state.db.cursors;
Ok(cursors
.contains_key(keys::new_pds_crawl_key(host))
.into_diagnostic()?
|| cursors
.contains_key(keys::rewalk_key(host))
.into_diagnostic()?)
}
#[cfg(feature = "indexer")]
fn finish_walk(hydrant: &Hydrant, host: &str) -> Result<()> {
let cursors = &hydrant.state.db.cursors;
cursors
.remove(keys::new_pds_crawl_key(host))
.into_diagnostic()?;
cursors.remove(keys::rewalk_key(host)).into_diagnostic()
}
#[cfg(feature = "indexer")]
#[tokio::test]
async fn a_known_host_gets_rewalks_until_its_bytes_are_spent_outside_the_daily_limit()
-> Result<()> {
let dir = tempfile::tempdir().into_diagnostic()?;
let hydrant = Hydrant::new(test_config_with_limit(&dir, 2)).await?;
startable(&hydrant);
hydrant.firehose.disable_spawns_for_test();
let url = Url::parse("wss://known.example/").into_diagnostic()?;
assert_eq!(
hydrant.firehose.add_source_validated(url.clone()).await?,
Ok(PdsAdmission::StartedNew)
);
finish_walk(&hydrant, "known.example")?;
for _ in 0..4 {
assert_eq!(
hydrant.firehose.add_source_validated(url.clone()).await?,
Ok(PdsAdmission::RestartedKnown)
);
assert!(
hydrant
.state
.db
.cursors
.contains_key(keys::rewalk_key("known.example"))
.into_diagnostic()?
);
finish_walk(&hydrant, "known.example")?;
}
hydrant
.state
.db
.cursors
.insert(
keys::rewalk_allowance_key("known.example"),
rmp_serde::to_vec(&Some(crate::types::Allowance {
since: chrono::Utc::now().timestamp(),
used: hydrant.state.rewalk_bytes,
}))
.into_diagnostic()?,
)
.into_diagnostic()?;
assert_eq!(
hydrant.firehose.add_source_validated(url).await?,
Ok(PdsAdmission::RestartedKnown)
);
assert!(!walking(&hydrant, "known.example")?);
assert!(matches!(
crate::db::load_pds_daily_adds(&hydrant.state.db)?,
Some((_, 1))
));
let other = Url::parse("wss://other.example/").into_diagnostic()?;
assert_eq!(
hydrant.firehose.add_source_validated(other).await?,
Ok(PdsAdmission::StartedNew)
);
Ok(())
}
#[cfg(feature = "indexer")]
#[tokio::test]
async fn requests_during_a_walk_neither_queue_nor_count() -> 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://busy.example/").into_diagnostic()?;
assert_eq!(
hydrant.firehose.add_source_validated(url.clone()).await?,
Ok(PdsAdmission::StartedNew)
);
// the new host's own walk, then a rewalk, both still running
for walk in [keys::new_pds_crawl_key, keys::rewalk_key] {
hydrant
.state
.db
.cursors
.insert(
walk("busy.example"),
crate::control::new_pds_crawl::encode_cursor(Some(&"page".into()))?,
)
.into_diagnostic()?;
for _ in 0..5 {
assert_eq!(
hydrant.firehose.add_source_validated(url.clone()).await?,
Ok(PdsAdmission::RestartedKnown)
);
}
assert_eq!(
crate::control::new_pds_crawl::decode_cursor(
&hydrant
.state
.db
.cursors
.get(walk("busy.example"))
.into_diagnostic()?
.expect("the running walk stays queued")
)?
.as_deref(),
Some("page"),
"a request must not restart a running walk"
);
finish_walk(&hydrant, "busy.example")?;
}
for _ in 0..3 {
assert_eq!(
hydrant.firehose.add_source_validated(url.clone()).await?,
Ok(PdsAdmission::RestartedKnown)
);
assert!(
walking(&hydrant, "busy.example")?,
"the allowance is still whole"
);
finish_walk(&hydrant, "busy.example")?;
}
Ok(())
}
#[cfg(feature = "indexer")]
#[tokio::test]
async fn a_banned_known_host_gets_no_rewalk() -> 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://turned.example/").into_diagnostic()?;
assert_eq!(
hydrant.firehose.add_source_validated(url.clone()).await?,
Ok(PdsAdmission::StartedNew)
);
finish_walk(&hydrant, "turned.example")?;
hydrant.pds.ban("turned.example").await?;
assert_eq!(
hydrant.firehose.add_source_validated(url).await?,
Err(PdsAdmissionRejection::HostBanned)
);
assert!(!walking(&hydrant, "turned.example")?);
Ok(())
}
/// `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(())
}
/// the crawler waits on a relay's dispatched time, which an offline relay doesn't move
#[cfg(feature = "indexer")]
#[tokio::test]
async fn an_offline_relay_stops_counting_for_the_crawler() -> Result<()> {
let dir = tempfile::tempdir().into_diagnostic()?;
let hydrant = Hydrant::new(test_config(&dir)).await?;
startable(&hydrant);
// nothing listens there, and the test firehose goes offline after one failure
let url = Url::parse("ws://127.0.0.1:1/").into_diagnostic()?;
let counted = || {
hydrant
.state
.firehose_dispatched
.peek_with(&url, |_, _| ())
.is_some()
};
hydrant.firehose.add_source(url.clone(), false).await?;
assert!(counted());
for _ in 0..500 {
if !counted() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
assert!(!counted(), "the relay never went offline");
assert!(
hydrant.firehose.is_source_running(&url),
"and it keeps retrying"
);
hydrant.firehose.tasks.remove_async(&url).await;
Ok(())
}
#[tokio::test]
async fn pds_info_reads_a_stored_big_endian_cursor() -> Result<()> {
let dir = tempfile::tempdir().into_diagnostic()?;
let hydrant = Hydrant::new(test_config(&dir)).await?;
let url = Url::parse("wss://stored.example/").into_diagnostic()?;
hydrant
.state
.db
.cursors
.insert(
keys::firehose_cursor_key_from_url(&url),
[0, 0, 0, 0, 0, 0, 1, 2],
)
.into_diagnostic()?;
let meta = hydrant.state.pds_meta.load();
let info = hydrant.firehose.pds_info(&url, &meta);
assert_eq!(info.map(|info| info.seq), Some(258));
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 a_validated_restart_waits_out_the_startup_jitter() -> 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://jittered.example/").into_diagnostic()?;
let delayed = || hydrant.firehose.last_spawn_delayed.load(Ordering::Relaxed);
hydrant.firehose.add_source(url.clone(), true).await?;
assert!(!delayed(), "an operator adding a source connects at once");
assert_eq!(
hydrant.firehose.add_source_validated(url.clone()).await?,
Ok(PdsAdmission::RestartedKnown)
);
assert!(
delayed(),
"a requested restart must not skip the startup jitter"
);
Ok(())
}
/// a host that keeps failing could otherwise call requestCrawl to wipe its own backoff
#[tokio::test]
async fn a_validated_request_restarts_a_source_only_once_its_offline() -> 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://retrying.example/").into_diagnostic()?;
let failing = || hydrant.state.throttler.snapshot(&url).is_failing();
hydrant.firehose.add_source(url.clone(), true).await?;
hydrant
.state
.throttler
.get_handle(&url)
.await
.record_failure();
let cancel = CancellationToken::new();
let _ = hydrant
.firehose
.tasks
.upsert_async(
url.clone(),
FirehoseIngestorHandle {
id: 0,
cancel: cancel.clone(),
},
)
.await;
assert_eq!(
hydrant.firehose.add_source_validated(url.clone()).await?,
Ok(PdsAdmission::AlreadyRunning)
);
assert!(failing(), "a source still retrying keeps its backoff");
assert!(!cancel.is_cancelled(), "and its task is left running");
// past max_failures the task marks the host offline and keeps retrying slowly
crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "retrying.example", |h| {
h.status = crate::pds_meta::HostStatus::Offline
});
assert_eq!(
hydrant.firehose.add_source_validated(url.clone()).await?,
Ok(PdsAdmission::RestartedKnown)
);
assert!(!failing(), "a source that gave up starts over");
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()], false)
.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_does_not_spend_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()], false)
.await?;
assert_eq!(adds.added, 2);
assert_eq!(adds.rejected, 0);
assert_eq!(persisted_sources(&hydrant).await?.len(), 2);
assert_eq!(crate::db::load_pds_daily_adds(&hydrant.state.db)?, None);
// public validated requests still enforce the daily limit independently
let third = Url::parse("wss://third-public.example/").into_diagnostic()?;
assert_eq!(
hydrant.firehose.add_source_validated(third.clone()).await?,
Ok(PdsAdmission::StartedNew)
);
let fourth = Url::parse("wss://fourth-public.example/").into_diagnostic()?;
assert_eq!(
hydrant
.firehose
.add_source_validated(fourth.clone())
.await?,
Err(PdsAdmissionRejection::DailyLimitReached)
);
hydrant.firehose.tasks.remove_async(&first).await;
hydrant.firehose.tasks.remove_async(&second).await;
hydrant.firehose.tasks.remove_async(&third).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(())
}
}