use std::collections::HashMap;
use std::future::Future;
use std::sync::atomic::AtomicI64;
#[cfg(feature = "indexer")]
use std::sync::atomic::Ordering;
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::cursor::SeqTracker;
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::{drain::Drain, stop::Stop, throttle::Throttler},
};
const IDENTITY_RECONCILE_COOL_OFF: Duration = Duration::from_secs(30);
#[cfg(all(test, feature = "indexer"))]
pub(crate) type AfterChunk = std::sync::Arc<dyn Fn(usize) + Send + Sync>;
/// one on one core, else one fewer than the cores, so live events always have a core that no
/// backfill step is using
#[cfg(feature = "indexer")]
fn bulk_permits() -> usize {
std::thread::available_parallelism().map_or(1, |cores| cores.get().saturating_sub(1).max(1))
}
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>,
/// one per source url for the life of the process, so an ingestor that replaces another
/// orders its frames behind the ones the old one left in flight instead of publishing a
/// second watermark into the same cursor
firehose_trackers: scc::HashMap<Url, std::sync::Arc<SeqTracker>>,
pub(crate) lookups: crate::ingest::relay::lookup::Lookups,
/// the repo row changes the relay handed the indexer and it hasn't applied yet
#[cfg(feature = "indexer")]
pub(crate) relay_rows: crate::ingest::relay::overlay::Overlay,
/// the event time, in ms, of the newest frame from each relay firehose source that's done
/// along with every frame read before it, 0 until one is, see [`Self::firehose_dispatched`].
/// a pds's own firehose isn't here, since it goes quiet whenever its accounts do. only the
/// crawler reads it, so relay builds keep none
#[cfg(feature = "indexer")]
pub(crate) firehose_dispatched: 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(crate) new_pds_crawl_notify: Notify,
/// how long a host's rewalk allowance lasts, see [`crate::control::new_pds_crawl`]
#[cfg(feature = "indexer")]
pub(crate) rewalk_window: Duration,
/// the `listRepos` bytes a host's rewalks may read per `rewalk_window`
#[cfg(feature = "indexer")]
pub(crate) rewalk_bytes: u32,
#[cfg(feature = "indexer")]
pub(crate) rev_clock_skew_secs: i64,
#[cfg(feature = "indexer")]
pub(crate) refetch_budget: crate::ingest::indexer::refetch::RefetchBudget,
#[cfg(feature = "indexer")]
pub(crate) refetch_denied: std::sync::atomic::AtomicU64,
#[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,
/// the crawler stopped queueing because the pending queue is full
#[cfg(feature = "indexer")]
pub(crate) crawler_throttled: std::sync::atomic::AtomicBool,
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,
#[cfg(feature = "indexer")]
pub(crate) backfill_budget: crate::backfill::budget::BackfillBudget,
/// records a full backfill writes per step on its indexer shard
#[cfg(feature = "indexer")]
pub(crate) backfill_chunk_records: usize,
/// where whole cars too big to hold in memory go
#[cfg(feature = "indexer")]
pub(crate) spills: crate::backfill::spill::Spills,
/// how many indexer shards may write a backfill step at once. one on one core, so a big
/// backfill per shard can't crowd out live events across the box
#[cfg(feature = "indexer")]
pub(crate) bulk_permits: Semaphore,
/// the indexer shards, once hydrant runs them, for the api and the crawler to hand them
/// work that only the repo's owner may write
#[cfg(feature = "indexer")]
pub(crate) indexer: parking_lot::Mutex<Option<crate::ingest::indexer::IndexerTx>>,
/// indexer shard threads still running, since once hydrant stops a write may only run in
/// place after they're gone
#[cfg(feature = "indexer")]
pub(crate) shards_running: std::sync::atomic::AtomicUsize,
/// runs after each backfill step a task sends, with the steps sent so far, so a test can
/// act between them
#[cfg(all(test, feature = "indexer"))]
pub(crate) after_chunk: parking_lot::Mutex<Option<AfterChunk>>,
/// runs on the shard after each drain step, with the steps done so far
#[cfg(all(test, feature = "indexer"))]
pub(crate) after_drain_step: parking_lot::Mutex<Option<AfterChunk>>,
/// runs on the shard after each step of an erase or a wipe, with its steps done so far
#[cfg(all(test, feature = "indexer"))]
pub(crate) after_job_step: parking_lot::Mutex<Option<AfterChunk>>,
pub get_repo_semaphore: Semaphore,
pub maintenance_timeout: Duration,
pub(crate) stop: Stop,
pub(crate) drain: Drain,
}
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_trackers: scc::HashMap::new(),
lookups: Default::default(),
#[cfg(feature = "indexer")]
relay_rows: Default::default(),
#[cfg(feature = "indexer")]
firehose_dispatched: scc::HashIndex::new(),
firehose_stats: FirehoseStats::default(),
#[cfg(feature = "indexer")]
backfill_notify: Notify::new(),
#[cfg(feature = "indexer")]
new_pds_crawl_notify: Notify::new(),
#[cfg(feature = "indexer")]
rewalk_window: config.rewalk_window,
#[cfg(feature = "indexer")]
rewalk_bytes: config.rewalk_mb.saturating_mul(1024 * 1024),
#[cfg(feature = "indexer")]
rev_clock_skew_secs: config.rev_clock_skew_secs,
#[cfg(feature = "indexer")]
refetch_budget: crate::ingest::indexer::refetch::RefetchBudget::new(config),
#[cfg(feature = "indexer")]
refetch_denied: std::sync::atomic::AtomicU64::new(0),
#[cfg(feature = "indexer")]
crawler_enabled,
firehose_enabled,
#[cfg(feature = "indexer")]
backfill_enabled,
#[cfg(feature = "indexer")]
backfill_drained: std::sync::atomic::AtomicUsize::new(0),
#[cfg(feature = "indexer")]
crawler_throttled: std::sync::atomic::AtomicBool::new(false),
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,
#[cfg(feature = "indexer")]
backfill_budget: crate::backfill::budget::BackfillBudget::new(
config.backfill_memory_bytes,
),
#[cfg(feature = "indexer")]
backfill_chunk_records: match config.backfill_chunk_records {
0 => usize::MAX,
records => records,
},
#[cfg(feature = "indexer")]
bulk_permits: Semaphore::new(bulk_permits()),
#[cfg(feature = "indexer")]
spills: crate::backfill::spill::Spills::open(
config
.backfill_spill_dir
.clone()
.unwrap_or_else(|| config.database_path.join("spill")),
config.backfill_spill_bytes,
config.backfill_spills,
)?,
#[cfg(feature = "indexer")]
indexer: parking_lot::Mutex::new(None),
#[cfg(feature = "indexer")]
shards_running: std::sync::atomic::AtomicUsize::new(0),
#[cfg(all(test, feature = "indexer"))]
after_chunk: parking_lot::Mutex::new(None),
#[cfg(all(test, feature = "indexer"))]
after_drain_step: parking_lot::Mutex::new(None),
#[cfg(all(test, feature = "indexer"))]
after_job_step: parking_lot::Mutex::new(None),
get_repo_semaphore: Semaphore::new(config.get_repo_concurrency_limit),
maintenance_timeout: config.maintenance_timeout,
stop: Stop::default(),
drain: Drain::default(),
throttler: Throttler::new(config.per_pds_concurrency),
pds_daily_limit,
})
}
pub(crate) async fn firehose_tracker(&self, url: &Url) -> std::sync::Arc<SeqTracker> {
self.firehose_trackers
.entry_async(url.clone())
.await
.or_insert_with(|| SeqTracker::new(crate::ingest::relay::sink::PUBLISHES_WATERMARK))
.get()
.clone()
}
#[cfg(feature = "indexer_stream")]
pub(crate) fn firehose_backlog(&self) -> usize {
let mut backlog = 0;
self.firehose_trackers.iter_sync(|_, tracker| {
backlog += tracker.pending();
true
});
backlog
}
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);
}
/// starts counting how far the relay firehose `firehose` got, for the crawler
#[cfg(feature = "indexer")]
pub(crate) async fn track_dispatched(&self, firehose: &Url) {
let _ = self
.firehose_dispatched
.insert_async(firehose.clone(), AtomicI64::new(0))
.await;
}
#[cfg(not(feature = "indexer"))]
pub(crate) async fn track_dispatched(&self, _firehose: &Url) {}
/// stops counting `firehose`, which never moves again once it's gone, so the crawler would
/// wait on it for nothing
#[cfg(feature = "indexer")]
pub(crate) async fn forget_dispatched(&self, firehose: &Url) {
self.firehose_dispatched.remove_async(firehose).await;
}
#[cfg(not(feature = "indexer"))]
pub(crate) async fn forget_dispatched(&self, _firehose: &Url) {}
/// notes that `firehose` is done with everything up to an event from `time_ms`. that time is
/// whatever the event's pds put there, so it counts as now at the most, or one event from the
/// future would end the crawler's waits on this source for good
#[cfg(feature = "indexer")]
pub(crate) fn note_dispatched(&self, firehose: &Url, time_ms: i64) {
let time_ms = time_ms.min(chrono::Utc::now().timestamp_millis());
self.firehose_dispatched
.peek_with(firehose, |_, dispatched| {
dispatched.fetch_max(time_ms, Ordering::Relaxed)
});
}
#[cfg(not(feature = "indexer"))]
pub(crate) fn note_dispatched(&self, _firehose: &Url, _time_ms: i64) {}
#[cfg(feature = "indexer")]
/// how far, in ms of event time, the relay firehose for `host` has handed events on. a host
/// hydrant doesn't subscribe to falls back to the relay furthest behind, since its accounts
/// may come through any of them. `None` when no relay firehose runs, paused included, since
/// nothing moves then
pub(crate) fn firehose_dispatched(&self, host: &str) -> Option<i64> {
if !*self.firehose_enabled.borrow() {
return None;
}
let (mut own, mut slowest) = (None, None::<i64>);
self.firehose_dispatched.iter_sync(|url, dispatched| {
let dispatched = dispatched.load(Ordering::Relaxed);
if url.host_str() == Some(host) {
own = Some(dispatched);
}
slowest = Some(slowest.map_or(dispatched, |slowest| slowest.min(dispatched)));
true
});
own.or(slowest)
}
/// 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)
}
/// checks the tier name, because in an indexer build the default tier has trusted limits
#[cfg(feature = "indexer")]
pub(crate) fn is_trusted_host(&self, host: &str) -> 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_name(host, override_name) == "trusted"
}
#[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(())
}
/// a paused firehose moves nothing, so there's nothing for the crawler to wait for, and a
/// future event time counts as now
#[cfg(feature = "indexer")]
#[test]
fn a_paused_firehose_has_no_dispatch_time() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let state = AppState::new(&Config {
database_path: tmp.path().to_path_buf(),
enable_firehose: true,
..Config::default()
})?;
let relay = Url::parse("wss://relay.example/").into_diagnostic()?;
let _ = state
.firehose_dispatched
.insert_sync(relay.clone(), AtomicI64::new(0));
let future = chrono::Utc::now().timestamp_millis() + 3_600_000;
state.note_dispatched(&relay, future);
let dispatched = state.firehose_dispatched("relay.example").unwrap();
assert!(dispatched < future, "a future event time was taken as is");
state.firehose_enabled.send_replace(false);
assert_eq!(state.firehose_dispatched("relay.example"), None);
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(())
}
}