Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434use std::collections::HashMap;use std::future::Future;use std::sync::atomic::AtomicI64;use std::time::{Duration, Instant};
use jacquard_common::types::string::{Did, Handle};use miette::Result;use scc::HashCache;use smol_str::SmolStr;#[cfg(feature = "indexer")]use tokio::sync::Notify;use tokio::sync::{Mutex, Semaphore, watch};use url::Url;
use crate::ingest::firehose_stats::FirehoseStats;use crate::pds_daily_limit::PdsDailyLimit;use crate::{ config::Config, db::Db, filter::{FilterHandle, new_handle as new_filter_handle}, pds_meta::{PdsMeta, PdsMetaHandle, TierPolicy, new_handle as new_pds_handle}, resolver::Resolver, types::DocIdentity, util::throttle::Throttler,};
const IDENTITY_RECONCILE_COOL_OFF: Duration = Duration::from_secs(30);
struct IdentityReconcileCooldown { identity: DocIdentity<'static>, requested_handle: Option<Handle>, gone: bool, expires_at: Instant,}
pub struct AppState { pub db: Db, pub resolver: Resolver, identity_reconcile_cooldowns: HashCache<Did, IdentityReconcileCooldown>, pub throttler: Throttler, pub(crate) filter: FilterHandle, pub(crate) pds_meta: PdsMetaHandle, /// Serializes validated PDS admission. pub(crate) pds_admission: Mutex<()>, /// serializes durable host-status writes with their in-memory mirrors. pub(crate) pds_status_write: std::sync::Arc<parking_lot::Mutex<()>>, pub(crate) pds_daily_limit: PdsDailyLimit, pub(crate) tier_policy: TierPolicy, pub firehose_cursors: scc::HashIndex<Url, AtomicI64>, pub(crate) firehose_stats: FirehoseStats, pub firehose_enabled: watch::Sender<bool>, #[cfg(feature = "indexer")] pub backfill_notify: Notify, #[cfg(feature = "indexer")] pub crawler_enabled: watch::Sender<bool>, #[cfg(feature = "indexer")] pub backfill_enabled: watch::Sender<bool>, #[cfg(feature = "indexer")] pub backfill_drained: std::sync::atomic::AtomicUsize, pub ephemeral: bool, pub ephemeral_ttl: Duration, /// bounded record-history retention; `None` keeps history permanently pub history_ttl: Option<Duration>, /// replay misses where a record body was likely trimmed by history /// retention (as opposed to a storage bug) pub history_trim_misses: std::sync::atomic::AtomicU64, pub only_index_links: bool, pub verify_cids: bool, pub max_car_body_bytes: usize, pub get_repo_semaphore: Semaphore, pub maintenance_timeout: Duration,}
impl AppState { pub fn new(config: &Config) -> Result<Self> { let db = Db::open(config)?; let resolver = Resolver::new(config.plc_urls.clone(), config.identity_cache_size); let filter_config = crate::db::filter::load(&db.filter)?;
#[cfg(feature = "indexer")] let crawler_default = match config.enable_crawler { Some(b) => b, // default: enabled if full-network mode, or if crawler sources are configured None => { filter_config.mode == crate::filter::FilterMode::Full || !config.crawler_sources.is_empty() } };
let filter = new_filter_handle(filter_config);
let tiers: HashMap<String, SmolStr> = crate::db::pds_meta::load_tiers(&db.filter) .unwrap_or_default() .into_iter() .map(|(host, tier)| (host.to_string(), tier)) .collect();
let statuses: HashMap<String, crate::pds_meta::HostStatus> = crate::db::pds_meta::load_statuses(&db.filter) .unwrap_or_default() .into_iter() .map(|(host, stat)| (host.to_string(), stat)) .collect(); let operator_bans = crate::db::pds_meta::load_bans(&db.filter).unwrap_or_default();
let mut hosts = HashMap::new(); for (host, tier) in tiers { hosts .entry(host) .or_insert_with(crate::pds_meta::HostDesc::default) .tier = Some(tier); } for (host, stat) in statuses { hosts .entry(host) .or_insert_with(crate::pds_meta::HostDesc::default) .status = stat; } for host in operator_bans { hosts .entry(host.to_string()) .or_insert_with(crate::pds_meta::HostDesc::default) .operator_ban = true; }
let pds_meta = new_pds_handle(PdsMeta { hosts });
let relay_cursors = scc::HashIndex::new(); let pds_daily_limit = PdsDailyLimit::new(config.new_host_limit, crate::db::load_pds_daily_adds(&db)?);
#[cfg(feature = "indexer")] let (crawler_enabled, _) = watch::channel(crawler_default); let (firehose_enabled, _) = watch::channel(config.enable_firehose); #[cfg(feature = "indexer")] let (backfill_enabled, _) = watch::channel(config.enable_backfill);
Ok(Self { db, resolver, identity_reconcile_cooldowns: HashCache::with_capacity( std::cmp::min(1000, (config.identity_cache_size / 100) as usize), config.identity_cache_size as usize, ), filter, pds_meta, pds_admission: Mutex::new(()), pds_status_write: std::sync::Arc::new(parking_lot::Mutex::new(())), tier_policy: config.tier_policy.clone(), firehose_cursors: relay_cursors, firehose_stats: FirehoseStats::default(), #[cfg(feature = "indexer")] backfill_notify: Notify::new(), #[cfg(feature = "indexer")] crawler_enabled, firehose_enabled, #[cfg(feature = "indexer")] backfill_enabled, #[cfg(feature = "indexer")] backfill_drained: std::sync::atomic::AtomicUsize::new(0), ephemeral: config.ephemeral, ephemeral_ttl: config.ephemeral_ttl, history_ttl: config.history_ttl, history_trim_misses: std::sync::atomic::AtomicU64::new(0), only_index_links: config.only_index_links, verify_cids: config.verify_cids, max_car_body_bytes: config.max_car_body_bytes, get_repo_semaphore: Semaphore::new(config.get_repo_concurrency_limit), maintenance_timeout: config.maintenance_timeout, // public backfill is direct-only; one socket pool must not inflate // per-PDS pacing as if proxies provided independent egress IPs. throttler: Throttler::new(1, config.per_pds_concurrency), pds_daily_limit, }) }
pub(crate) async fn identity_reconcile_is_cooled_down( &self, did: &Did, identity: &DocIdentity<'static>, requested_handle: Option<&Handle>, gone: bool, ) -> bool { let now = Instant::now(); self.identity_reconcile_cooldowns .read_async(did, |_, cooldown| { cooldown.expires_at > now && &cooldown.identity == identity && cooldown.requested_handle.as_ref() == requested_handle && cooldown.gone == gone }) .await .unwrap_or(false) }
pub(crate) async fn cool_down_identity_reconcile( &self, did: Did, identity: DocIdentity<'static>, requested_handle: Option<Handle>, gone: bool, ) { let cooldown = IdentityReconcileCooldown { identity, requested_handle, gone, expires_at: Instant::now() + IDENTITY_RECONCILE_COOL_OFF, }; let _ = self .identity_reconcile_cooldowns .entry_async(did) .await .put_entry(cooldown); }
pub fn stores_record_bodies(&self) -> bool { !self.ephemeral && !self.only_index_links }
#[cfg(feature = "indexer")] pub(crate) fn ensure_operator_body_deletion_supported(&self) -> Result<()> { if self.ephemeral { miette::bail!("operator body deletion is not supported in ephemeral mode"); } if !self.stores_record_bodies() { miette::bail!("record bodies are not stored in links-only mode"); } Ok(()) }
#[cfg(feature = "indexer")] pub fn notify_backfill(&self) { self.backfill_notify.notify_one(); }
/// 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<F, Fut, T>(&self, f: F) -> T where F: FnOnce() -> Fut, Fut: Future<Output = T>, { #[cfg(feature = "indexer")] let crawler_was = *self.crawler_enabled.borrow(); let firehose_was = *self.firehose_enabled.borrow(); #[cfg(feature = "indexer")] let backfill_was = *self.backfill_enabled.borrow();
#[cfg(feature = "indexer")] self.crawler_enabled.send_replace(false); self.firehose_enabled.send_replace(false); #[cfg(feature = "indexer")] self.backfill_enabled.send_replace(false);
struct IngestionPauseGuard<'a> { state: &'a AppState, #[cfg(feature = "indexer")] crawler_was: bool, firehose_was: bool, #[cfg(feature = "indexer")] backfill_was: bool, }
impl<'a> Drop for IngestionPauseGuard<'a> { fn drop(&mut self) { #[cfg(feature = "indexer")] self.state.crawler_enabled.send_replace(self.crawler_was); self.state.firehose_enabled.send_replace(self.firehose_was); #[cfg(feature = "indexer")] self.state.backfill_enabled.send_replace(self.backfill_was); tracing::info!("ingestion resumed"); } }
let _guard = IngestionPauseGuard { state: self, #[cfg(feature = "indexer")] crawler_was, firehose_was, #[cfg(feature = "indexer")] backfill_was, };
f().await }
/// applies an account limit status transition for `host`, writing to `batch` and updating /// in-memory state. call this after any event that changes the active account count for a PDS. pub(crate) fn apply_host_limit_status( &self, batch: &mut fjall::OwnedWriteBatch, host: &str, count: u64, ) { use crate::db::pds_meta as db_pds; use crate::pds_meta::PdsMeta; use tracing::{debug, error};
let _status_write = self.pds_status_write.lock(); let (current_status, limit) = { let meta = self.pds_meta.load(); let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref()); let tier = self.tier_policy.resolve(host, override_name); (meta.status(host), tier.account_limit) };
debug!(%host, ?current_status, ?limit, count, "apply_host_limit_status");
let Some(new_status) = current_status.check_limit_transition(count, limit) else { return; };
debug!(%host, count, ?limit, ?new_status, "account count crossed limit, shifting status");
if let Err(e) = db_pds::set_status(batch, &self.db.filter, host, new_status) { error!(%host, err = %e, "failed to write host status"); return; }
PdsMeta::update_host(&self.pds_meta, host, |h| h.status = new_status); }
/// checks whether `host` is at or over its account limit at the given count. /// does not modify any state. pub(crate) fn is_over_account_limit(&self, host: &str, count: u64) -> bool { let meta = self.pds_meta.load(); let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref()); self.tier_policy .resolve(host, override_name) .account_limit .is_some_and(|l| count >= l) }
#[cfg(feature = "indexer")] pub(crate) fn resolve_pds_tier(&self, host: &str) -> crate::config::RateTier { let meta = self.pds_meta.load(); let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref()); self.tier_policy.resolve(host, override_name) }}
#[cfg(test)]mod tests { use super::*; use crate::config::Config; use crate::pds_meta::HostStatus; use miette::IntoDiagnostic; use std::time::Duration; use tokio::time::timeout;
#[test] fn banned_host_survives_an_account_limit_transition() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let cfg = Config { database_path: tmp.path().to_path_buf(), ..Config::default() }; let state = AppState::new(&cfg)?; let host = "banned-limit.example";
let mut ban = state.db.inner.batch(); crate::db::pds_meta::set_ban(&mut ban, &state.db.filter, host); ban.commit().into_diagnostic()?; state.db.persist()?; PdsMeta::update_host(&state.pds_meta, host, |entry| { entry.operator_ban = true; });
let mut batch = state.db.inner.batch(); state.apply_host_limit_status(&mut batch, host, u64::MAX); batch.commit().into_diagnostic()?; state.db.persist()?;
assert!(state.pds_meta.load().is_banned(host)); assert_eq!(state.pds_meta.load().status(host), HostStatus::Throttled); assert!( state .db .filter .contains_key(crate::db::pds_meta::pds_ban_key(host)) .into_diagnostic()? ); Ok(()) }
#[tokio::test] async fn test_with_ingestion_paused_restore_on_cancel() -> Result<()> { let temp = tempfile::tempdir().unwrap(); let mut config = Config::default(); config.database_path = temp.path().to_path_buf(); #[cfg(feature = "indexer")] { config.enable_crawler = Some(true); config.enable_backfill = true; } config.enable_firehose = true;
let state = AppState::new(&config)?;
#[cfg(feature = "indexer")] { assert!(*state.crawler_enabled.borrow()); assert!(*state.backfill_enabled.borrow()); } assert!(*state.firehose_enabled.borrow());
let result = timeout( Duration::from_millis(50), state.with_ingestion_paused(async || { #[cfg(feature = "indexer")] { assert!(!*state.crawler_enabled.borrow()); assert!(!*state.backfill_enabled.borrow()); } assert!(!*state.firehose_enabled.borrow());
tokio::time::sleep(Duration::from_secs(10)).await; }), ) .await;
assert!(result.is_err()); // should be timed out
#[cfg(feature = "indexer")] { assert!(*state.crawler_enabled.borrow()); assert!(*state.backfill_enabled.borrow()); } assert!(*state.firehose_enabled.borrow());
Ok(()) }}