Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339use std::sync::Arc;
use miette::{IntoDiagnostic, Result};use tokio::sync::{mpsc, watch};use tracing::{error, info};use url::Url;
use crate::db::keys;use crate::state::AppState;
pub(super) struct ProducerHandle { mode: crate::config::CrawlerMode, abort: tokio::task::AbortHandle,}
impl Drop for ProducerHandle { fn drop(&mut self) { self.abort.abort(); }}
pub(super) struct CrawlerShared { pub(super) http: reqwest::Client, pub(super) checker: crate::crawler::SignalChecker, pub(super) in_flight: crate::crawler::InFlight, pub(super) tx: mpsc::Sender<crate::crawler::CrawlerBatch>, pub(super) stats: crate::crawler::CrawlerStats,}
/// a snapshot of a single crawler source's runtime state.#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]pub struct CrawlerSourceInfo { pub url: Url, pub mode: crate::config::CrawlerMode,}
#[allow(clippy::too_many_arguments)]pub(super) fn spawn_crawler_producer( source: &crate::config::CrawlerSource, http: &reqwest::Client, state: &Arc<AppState>, checker: &crate::crawler::SignalChecker, in_flight: &crate::crawler::InFlight, tx: &mpsc::Sender<crate::crawler::CrawlerBatch>, stats: &crate::crawler::CrawlerStats, enabled: watch::Receiver<bool>,) -> ProducerHandle { use crate::config::CrawlerMode; use crate::crawler::{ByCollectionProducer, ListReposProducer, RelayFirstProducer}; use std::time::Duration; use tracing::Instrument;
let abort = match source.mode { CrawlerMode::ListRepos => { info!(relay = %source.url, enabled = *state.crawler_enabled.borrow(), "starting relay crawler"); let span = tracing::info_span!("crawl", url = %source.url); tokio::spawn( ListReposProducer { url: source.url.clone(), checker: checker.clone(), in_flight: in_flight.clone(), tx: tx.clone(), enabled, stats: stats.clone(), completion_key: None, last_completion: None, } .run() .instrument(span), ) .abort_handle() } CrawlerMode::RelayFirst => { info!(relay = %source.url, enabled = *state.crawler_enabled.borrow(), "starting relay-first crawler"); let span = tracing::info_span!("relay_first", url = %source.url); tokio::spawn( RelayFirstProducer { url: source.url.clone(), checker: checker.clone(), in_flight: in_flight.clone(), tx: tx.clone(), enabled, stats: stats.clone(), } .run() .instrument(span), ) .abort_handle() } CrawlerMode::ByCollection => { info!( host = source.url.host_str(), enabled = *state.crawler_enabled.borrow(), "starting by-collection crawler" ); let span = tracing::info_span!("by_collection", host = source.url.host_str()); let http = http.clone(); let state = state.clone(); let in_flight = in_flight.clone(); let tx = tx.clone(); let stats = stats.clone(); let url = source.url.clone(); tokio::spawn( async move { loop { let producer = ByCollectionProducer { index_url: url.clone(), http: http.clone(), state: state.clone(), in_flight: in_flight.clone(), tx: tx.clone(), enabled: enabled.clone(), stats: stats.clone(), }; if let Err(e) = producer.run().await { error!(err = ?e, "by-collection crawler fatal error, restarting in 30s"); tokio::time::sleep(Duration::from_secs(30)).await; } } } .instrument(span), ) .abort_handle() } }; ProducerHandle { mode: source.mode, abort, }}
/// runtime control over the crawler component.////// the crawler walks `com.atproto.sync.listRepos` on each configured relay to discover/// repositories that have never emitted a firehose event. in `filter` mode it also/// checks each discovered repo against the configured signal collections before/// enqueuing it for backfill.////// disabling the crawler does not affect in-progress repo checks. each one completes/// its current PDS request before pausing.#[derive(Clone)]pub struct CrawlerHandle { 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<CrawlerShared>>, /// per-source running tasks, keyed by url. pub(super) tasks: Arc<scc::HashMap<Url, ProducerHandle>>, /// set of urls persisted in the database (dynamically added sources). pub(super) persisted: Arc<scc::HashSet<Url>>,}
impl CrawlerHandle { /// enable the crawler (enables all configured producers). no-op if already enabled. pub fn enable(&self) { self.state.crawler_enabled.send_replace(true); } /// disable the crawler (disables all configured producers). /// in-progress repo checks finish before the crawler pauses. pub fn disable(&self) { self.state.crawler_enabled.send_replace(false); } /// returns the current enabled state of the crawler. pub fn is_enabled(&self) -> bool { *self.state.crawler_enabled.borrow() }
/// delete all cursor entries associated with the given URL. pub async fn reset_cursor(&self, url: &str) -> Result<()> { let point_keys = [keys::crawler_cursor_key(url)]; let by_collection_prefix = keys::by_collection_cursor_prefix(url); let relay_first_prefix = keys::relay_first_cursor_prefix(url); self.state .db .run(move |db| { let mut batch = db.inner.batch(); for k in point_keys { batch.remove(&db.cursors, k); } for entry in db.cursors.prefix(&by_collection_prefix) { let k = entry.key().into_diagnostic()?; batch.remove(&db.cursors, k); } for entry in db.cursors.prefix(&relay_first_prefix) { let k = entry.key().into_diagnostic()?; batch.remove(&db.cursors, k); } batch.commit().into_diagnostic()?; db.persist() }) .await?; Ok(()) }
/// return info on all currently active crawler sources. /// /// returns an empty list if called before [`Hydrant::run`]. pub async fn list_sources(&self) -> Vec<CrawlerSourceInfo> { let mut sources = Vec::new(); self.tasks .iter_async(|url, h| { sources.push(CrawlerSourceInfo { url: url.clone(), mode: h.mode, }); true }) .await; sources }
/// add a new crawler source at runtime. /// /// the source is persisted to the database and will be re-spawned on restart. /// if a source with the same URL already exists, it is replaced (the old task is /// aborted and a new one is started with the new mode). /// /// returns an error if called before [`Hydrant::run`]. pub async fn add_source(&self, source: crate::config::CrawlerSource) -> Result<()> { let Some(shared) = self.shared.get() else { miette::bail!("crawler not yet started: call Hydrant::run() first"); };
let key = keys::crawler_source_key(source.url.as_str()); let val = rmp_serde::to_vec(&source.mode).into_diagnostic()?; self.state .db .run(move |db| { db.crawler.insert(key, val).into_diagnostic()?; db.persist() }) .await?;
let enabled_rx = self.state.crawler_enabled.subscribe(); let handle = spawn_crawler_producer( &source, &shared.http, &self.state, &shared.checker, &shared.in_flight, &shared.tx, &shared.stats, enabled_rx, );
let _ = self.persisted.insert_async(source.url.clone()).await; match self.tasks.entry_async(source.url).await { scc::hash_map::Entry::Vacant(e) => { e.insert_entry(handle); } scc::hash_map::Entry::Occupied(mut e) => { *e.get_mut() = handle; } } Ok(()) }
/// remove a crawler source at runtime by URL. /// /// aborts the running producer task and removes the source from the database if it /// was dynamically added. config-sourced entries are aborted but not persisted, so /// they will reappear on restart. /// /// returns `true` if a source with the given URL was found and removed. /// returns an error if called before [`Hydrant::run`]. pub async fn remove_source(&self, url: &Url) -> Result<bool> { if self.shared.get().is_none() { miette::bail!("crawler not yet started: call Hydrant::run() first"); }
// dropping the ProducerHandle aborts the task via Drop if self.tasks.remove_async(url).await.is_none() { return Ok(false); }
// remove from DB if it was a persisted source if self.persisted.remove_async(url).await.is_some() { let key = keys::crawler_source_key(url.as_str()); self.state .db .run(move |db| { db.crawler.remove(key).into_diagnostic()?; db.persist() }) .await?; }
Ok(true) }}
#[cfg(test)]mod tests { use super::*; use crate::config::Config; use tempfile::tempdir;
#[tokio::test] async fn reset_cursor_clears_relay_first_progress_for_only_that_source() -> Result<()> { let tmp = tempdir().into_diagnostic()?; let config = Config { database_path: tmp.path().to_path_buf(), ..Default::default() }; let state = Arc::new(AppState::new(&config)?); let handle = CrawlerHandle { state: state.clone(), shared: Arc::new(std::sync::OnceLock::new()), tasks: Arc::new(scc::HashMap::new()), persisted: Arc::new(scc::HashSet::new()), }; let relay = "https://relay.example/"; let other = "https://other.example/"; let relay_keys = [ keys::crawler_cursor_key(relay), keys::relay_first_pass_key(relay), keys::relay_first_handled_key(relay), keys::relay_first_hosts_cursor_key(relay), keys::relay_first_pds_cursor_key(relay, "pds.example"), ]; for key in &relay_keys { state.db.cursors.insert(key, b"value").into_diagnostic()?; } let other_key = keys::relay_first_pass_key(other); state .db .cursors .insert(&other_key, b"other") .into_diagnostic()?;
handle.reset_cursor(relay).await?;
for key in relay_keys { assert!(state.db.cursors.get(key).into_diagnostic()?.is_none()); } assert!(state.db.cursors.get(other_key).into_diagnostic()?.is_some()); Ok(()) }}