diff --git a/AGENTS.md b/AGENTS.md index 68e2f8c..05951f1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -118,7 +118,7 @@ Hydrant uses multiple `fjall` keyspaces: - `redactions`: Maps `{record key} 00 {CID bytes}` -> empty for operator-erased versions. Markers are permanent and suppress the same record+CID on replay even after an intervening version or history compaction. - `event_bodies` (indexer stream): Finite compatibility archive, `{record key} 00 {CID bytes}` -> raw DAG-CBOR. v10 uses it only for legacy events that cannot resolve through `heads` or `history`; new writes never add entries. - `events`: Maps `{ID}` (u64 BE) -> `StoredEvent` (MessagePack). This is the source for the JSON stream API. Ephemeral mode stores create/update bodies inline as `StoredData::Block`, so event TTL removes identity and body bytes together without creating record heads or history. -- `cursors`: Maps per-relay cursor keys -> `Value` (u64/i64 BE Bytes). Keys: `firehose_cursor|{relay}`, `crawler_cursor|{relay}`, `by_collection_cursor|{url}|{collection}`, and `stream_next|{counter}` stream position marks, which every sequenced commit writes so positions never repeat after their rows are pruned or when they have none (`db::sequencer`). +- `cursors`: Maps per-relay cursor keys -> `Value` (u64/i64 BE Bytes). Keys: `firehose_cursor|{relay}`, `crawler_cursor|{relay}`, `by_collection_cursor|{url}|{collection}`, `new_pds_crawl|{host}` (the page cursor of a newly followed PDS's one-off `listRepos` walk, deleted with its last page), `seed_listed|{seed url}` (set once a seed relay's `listHosts` was read to the end, after which hosts new on it get that walk), and `stream_next|{counter}` stream position marks, which every sequenced commit writes so positions never repeat after their rows are pruned or when they have none (`db::sequencer`). - `pending`: Queue of `{ID}` (u64 BE) -> `Empty` (Backfill queue). - `resync`: Maps `{DID}` -> `ResyncState` (MessagePack) for retry logic/tombstones. - `resync_buffer`: Maps `{DID}|{Rev}` -> `Commit` (MessagePack). Used to buffer live events during backfill. diff --git a/docs/concepts/relay.md b/docs/concepts/relay.md index afd5634..f85048a 100644 --- a/docs/concepts/relay.md +++ b/docs/concepts/relay.md @@ -38,6 +38,8 @@ banned hosts (`status: "banned"`) are skipped. all other statuses are included s a host we have no cursor for starts from the seq the seed relay reported for it, or from the live head when it reported none, so its retained history is not replayed. hosts that already have a cursor keep it. +in indexer mode that live tail would miss repos that already lived on a new PDS, so a host that is new to the network also gets one `listRepos` walk, its repos going through the same signal check as crawler sources. a seed's hosts only count as new once its `listHosts` has been read to the end once, so a fresh node doesn't walk every PDS it seeds. a PDS admitted through `com.atproto.sync.requestCrawl` always gets the walk. it runs while the crawler is enabled, resumes from its last page after a restart, and is dropped if the host is banned or removed as a source. + seeding runs from latest cursor on restart so new PDS' added to the upstream relay since the last start are picked up automatically (if they haven't through firehose). sources that are already running are detected and skipped, so re-seeding is idempotent. ## crawler sources diff --git a/src/control/firehose.rs b/src/control/firehose.rs index 560cd30..f1378f1 100644 --- a/src/control/firehose.rs +++ b/src/control/firehose.rs @@ -12,6 +12,7 @@ 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")] @@ -431,11 +432,12 @@ impl FirehoseHandle { return Ok(Err(PdsAdmissionRejection::DailyLimitReached)); }; - self.spawn_admitted_source(&url, host.as_str(), false, reservation) + // 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 } - async fn admit_seeded_source(&self, url: Url) -> Result { + async fn admit_seeded_source(&self, url: Url, queue_crawl: bool) -> Result { let (host, url) = match canonical_pds_url(&url) { Ok(source) => source, Err(error) => return Ok(Err(PdsAdmissionRejection::InvalidHost(error))), @@ -449,8 +451,14 @@ impl FirehoseHandle { return Ok(early); } - self.spawn_admitted_source(&url, host.as_str(), true, PdsDailyReservation::Unlimited) - .await + self.spawn_admitted_source( + &url, + host.as_str(), + true, + PdsDailyReservation::Unlimited, + queue_crawl, + ) + .await } async fn check_admission_preconditions( @@ -503,6 +511,7 @@ impl FirehoseHandle { host: &str, delay_startup: bool, reservation: PdsDailyReservation, + queue_crawl: bool, ) -> Result { let shared = match self.shared.get() { Some(shared) => shared, @@ -525,7 +534,7 @@ impl FirehoseHandle { }; match self - .persist_validated_source(&source, host, reservation) + .persist_validated_source(&source, host, reservation, queue_crawl) .await { Ok(true) => {} @@ -551,6 +560,9 @@ impl FirehoseHandle { 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)) } @@ -559,6 +571,7 @@ impl FirehoseHandle { source: &FirehoseSource, host: &str, reservation: PdsDailyReservation, + queue_crawl: bool, ) -> Result { let key = keys::firehose_source_key(source.url.as_str()); let marker_key = keys::trusted_firehose_source_key(source.url.as_str()); @@ -569,6 +582,7 @@ impl FirehoseHandle { 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| { @@ -593,6 +607,9 @@ impl FirehoseHandle { } 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!( @@ -610,7 +627,12 @@ impl FirehoseHandle { /// /// seeded sources bypass the public crawl request daily limit. /// startup stays staggered so a page cannot stampede remote hosts. - pub(super) async fn add_seeded_sources(&self, urls: Vec) -> Result { + /// `queue_crawl` is set once the seed is past bootstrap, see [`new_pds_crawl`]. + pub(super) async fn add_seeded_sources( + &self, + urls: Vec, + queue_crawl: bool, + ) -> Result { let mut adds = SeededSourceAdds { added: 0, known: 0, @@ -621,7 +643,7 @@ impl FirehoseHandle { adds.known += 1; continue; } - match self.admit_seeded_source(url.clone()).await? { + match self.admit_seeded_source(url.clone(), queue_crawl).await? { Ok(PdsAdmission::StartedNew) => adds.added += 1, Ok(PdsAdmission::RestartedKnown | PdsAdmission::AlreadyRunning) => { adds.known += 1; @@ -783,6 +805,58 @@ mod tests { .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(()) + } + /// `upsert_async` returns the previous handle, whose drop cancels its task. #[tokio::test] async fn replacing_a_source_cancels_the_previous_ingestor() -> Result<()> { @@ -925,7 +999,7 @@ mod tests { let adds = hydrant .firehose - .add_seeded_sources(vec![url.clone()]) + .add_seeded_sources(vec![url.clone()], false) .await?; assert_eq!(adds.added, 0); @@ -946,7 +1020,7 @@ mod tests { let adds = hydrant .firehose - .add_seeded_sources(vec![first.clone(), second.clone()]) + .add_seeded_sources(vec![first.clone(), second.clone()], false) .await?; assert_eq!(adds.added, 2); diff --git a/src/control/hydrant/run.rs b/src/control/hydrant/run.rs index 111d557..6dfd0c4 100644 --- a/src/control/hydrant/run.rs +++ b/src/control/hydrant/run.rs @@ -464,6 +464,21 @@ impl Hydrant { .expect("crawler shared already set"); let shared = crawler.shared.get().unwrap(); + tokio::spawn( + crate::crawler::NewPdsCrawler { + lister: crate::crawler::PdsRepoLister { + checker: shared.checker.clone(), + in_flight: shared.in_flight.clone(), + tx: shared.tx.clone(), + enabled: state.crawler_enabled.subscribe(), + stats: shared.stats.clone(), + component: "new PDS crawler", + }, + http: crate::net::public_http(), + } + .run(), + ); + // spawn initial sources from config for source in config.crawler_sources.iter() { let enabled_rx = state.crawler_enabled.subscribe(); diff --git a/src/control/mod.rs b/src/control/mod.rs index a5f9a04..23c5740 100644 --- a/src/control/mod.rs +++ b/src/control/mod.rs @@ -4,6 +4,7 @@ pub(crate) mod crawler; pub(crate) mod filter; pub(crate) mod firehose; +pub(crate) mod new_pds_crawl; pub(crate) mod pds; pub(crate) mod repos; mod seed; diff --git a/src/control/new_pds_crawl.rs b/src/control/new_pds_crawl.rs new file mode 100644 index 0000000..34ac160 --- /dev/null +++ b/src/control/new_pds_crawl.rs @@ -0,0 +1,81 @@ +//! the one-off `listRepos` walk an indexer owes a PDS it starts following. +//! +//! a new PDS is tailed from the relay's seq, so repos that already lived on it +//! never reach us as events and nothing else lists them. only hosts that are new +//! to the network get the walk: a seed relay's hosts count as bootstrap until its +//! `listHosts` has been read to the end once, otherwise a fresh node would walk +//! every PDS out there. a relay has nothing to backfill, so there it's a no-op. + +#[cfg(feature = "indexer")] +mod enabled { + use miette::{IntoDiagnostic, Result}; + use smol_str::SmolStr; + use url::Url; + + use crate::db::{Db, keys}; + use crate::state::AppState; + + /// stage the walk in the batch that admits the host, so a restart can't lose it. + pub(crate) fn queue(batch: &mut fjall::OwnedWriteBatch, db: &Db, host: &str) -> Result<()> { + batch.insert( + &db.cursors, + keys::new_pds_crawl_key(host), + encode_cursor(None)?, + ); + Ok(()) + } + + /// call once the batch from [`queue`] has committed. + pub(crate) fn wake(state: &AppState) { + state.new_pds_crawl_notify.notify_one(); + } + + pub(crate) fn seed_listed(db: &Db, seed: &Url) -> Result { + db.cursors + .contains_key(keys::seed_listed_key(seed.as_str())) + .into_diagnostic() + } + + pub(crate) fn mark_seed_listed(db: &Db, seed: &Url) -> Result<()> { + db.cursors + .insert(keys::seed_listed_key(seed.as_str()), b"") + .into_diagnostic() + } + + pub(crate) fn encode_cursor(cursor: Option<&SmolStr>) -> Result> { + rmp_serde::to_vec(&cursor).into_diagnostic() + } + + pub(crate) fn decode_cursor(bytes: &[u8]) -> Result> { + rmp_serde::from_slice(bytes).into_diagnostic() + } +} + +#[cfg(feature = "indexer")] +pub(crate) use enabled::*; + +#[cfg(not(feature = "indexer"))] +mod noop { + use miette::Result; + use url::Url; + + use crate::db::Db; + use crate::state::AppState; + + pub(crate) fn queue(_batch: &mut fjall::OwnedWriteBatch, _db: &Db, _host: &str) -> Result<()> { + Ok(()) + } + + pub(crate) fn wake(_state: &AppState) {} + + pub(crate) fn seed_listed(_db: &Db, _seed: &Url) -> Result { + Ok(false) + } + + pub(crate) fn mark_seed_listed(_db: &Db, _seed: &Url) -> Result<()> { + Ok(()) + } +} + +#[cfg(not(feature = "indexer"))] +pub(crate) use noop::*; diff --git a/src/control/seed.rs b/src/control/seed.rs index 3a43e10..0a494b8 100644 --- a/src/control/seed.rs +++ b/src/control/seed.rs @@ -9,6 +9,7 @@ use tracing::{debug, info, warn}; use url::Url; use super::firehose::FirehoseHandle; +use super::new_pds_crawl; use crate::db::keys; use crate::pds_discovery::{canonical_pds_host, list_hosts_url}; use crate::state::AppState; @@ -76,6 +77,15 @@ async fn seed_one( let mut cursor: Option = None; info!("seeding firehose sources from listHosts"); + // read once, so the pass that first reaches the end still counts as bootstrap + let queue_crawl = firehose.is_some() + && new_pds_crawl::seed_listed(&state.db, seed_url).unwrap_or_else(|error| { + warn!(%error, "failed to read the seed listed marker, treating its hosts as bootstrap"); + false + }); + let mut reached_end = false; + let mut admitted_every_page = true; + let mut total = 0usize; let mut added = 0usize; @@ -169,7 +179,7 @@ async fn seed_one( seed_sources.push(wss_url); } - match firehose.add_seeded_sources(seed_sources).await { + match firehose.add_seeded_sources(seed_sources, queue_crawl).await { Ok(adds) => { page_added = adds.added; added += adds.added; @@ -177,6 +187,7 @@ async fn seed_one( rejected = adds.rejected; } Err(error) => { + admitted_every_page = false; warn!(%error, "failed to add seeded firehose sources"); } } @@ -198,10 +209,21 @@ async fn seed_one( cursor = next_cursor; if cursor.is_none() { + reached_end = true; break; } } + // hosts a failed page dropped would be admitted later as if they were new + if firehose.is_some() && reached_end && admitted_every_page && !queue_crawl { + match new_pds_crawl::mark_seed_listed(&state.db, seed_url) { + Ok(()) => { + info!("seed listed to the end, hosts new on it from now on get a listRepos walk") + } + Err(error) => warn!(%error, "failed to write the seed listed marker"), + } + } + info!( total, added, "finished seeding firehose sources from listHosts" @@ -329,6 +351,123 @@ mod tests { .transpose() } + /// serves `listHosts` with the first host on page one and the rest on page two. + #[cfg(feature = "indexer")] + struct MockRelay { + hosts: std::sync::Mutex>, + fail_second_page: std::sync::atomic::AtomicBool, + } + + #[cfg(feature = "indexer")] + async fn mock_relay() -> miette::Result<(Url, Arc)> { + use axum::extract::{Query, State}; + use axum::response::{IntoResponse, Response}; + use axum::{Json, Router, http::StatusCode, routing::get}; + use std::collections::HashMap; + + async fn list_hosts( + State(relay): State>, + Query(query): Query>, + ) -> Response { + let hosts = relay.hosts.lock().unwrap().clone(); + let host = |hostname: &str| serde_json::json!({ "hostname": hostname, "seq": 1, "status": "active" }); + if !query.contains_key("cursor") { + return Json(serde_json::json!({ "cursor": "1", "hosts": [host(hosts[0])] })) + .into_response(); + } + if relay.fail_second_page.load(Ordering::SeqCst) { + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + let rest: Vec<_> = hosts[1..].iter().map(|hostname| host(hostname)).collect(); + Json(serde_json::json!({ "hosts": rest })).into_response() + } + + let relay = Arc::new(MockRelay { + hosts: std::sync::Mutex::new(Vec::new()), + fail_second_page: std::sync::atomic::AtomicBool::new(false), + }); + let app = Router::new() + .route("/xrpc/com.atproto.sync.listHosts", get(list_hosts)) + .with_state(relay.clone()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .into_diagnostic()?; + let addr = listener.local_addr().into_diagnostic()?; + tokio::spawn(async move { axum::serve(listener, app).await }); + Ok(( + Url::parse(&format!("http://{addr}/")).into_diagnostic()?, + relay, + )) + } + + #[cfg(feature = "indexer")] + #[tokio::test] + async fn seed_hosts_are_bootstrap_until_the_seed_is_listed_to_the_end() -> miette::Result<()> { + let tmp = tempdir().into_diagnostic()?; + let hydrant = Hydrant::new(Config { + database_path: tmp.path().to_path_buf(), + ..Default::default() + }) + .await?; + hydrant.firehose.startable_for_test(); + hydrant.firehose.disable_spawns_for_test(); + let (seed, relay) = mock_relay().await?; + let http = seed_http_client(); + let pass = || seed_one(&seed, Some(hydrant.firehose.clone()), &hydrant.state, &http); + let known = |host: &str| -> miette::Result { + Ok(hydrant + .firehose + .is_source_known(&Url::parse(&format!("wss://{host}/")).into_diagnostic()?)) + }; + let queued = |host: &str| { + hydrant + .state + .db + .cursors + .contains_key(keys::new_pds_crawl_key(host)) + .into_diagnostic() + }; + + *relay.hosts.lock().unwrap() = vec!["a.example", "b.example"]; + relay.fail_second_page.store(true, Ordering::SeqCst); + pass().await; + assert!(known("a.example")?); + assert!(!new_pds_crawl::seed_listed(&hydrant.state.db, &seed)?); + + // b is new to us but not to the network, the first full listing only now reaches it + relay.fail_second_page.store(false, Ordering::SeqCst); + pass().await; + assert!(known("b.example")?); + assert!(!queued("a.example")?); + assert!(!queued("b.example")?); + assert!(new_pds_crawl::seed_listed(&hydrant.state.db, &seed)?); + + *relay.hosts.lock().unwrap() = vec!["a.example", "b.example", "c.example"]; + pass().await; + assert!(queued("c.example")?); + assert!(!queued("a.example")?); + assert!(!queued("b.example")?); + Ok(()) + } + + #[cfg(feature = "indexer")] + #[tokio::test] + async fn snapshot_refresh_never_marks_a_seed_listed() -> miette::Result<()> { + let tmp = tempdir().into_diagnostic()?; + let state = Arc::new(AppState::new(&Config { + database_path: tmp.path().to_path_buf(), + ..Default::default() + })?); + let (seed, relay) = mock_relay().await?; + *relay.hosts.lock().unwrap() = vec!["a.example", "b.example"]; + + // a refresh admits nothing, so marking the seed here would make the next pass treat every host as new + refresh_seed_snapshots(std::slice::from_ref(&seed), &state).await; + assert_eq!(persisted_cursor(&state, "b.example")?, Some(1)); + assert!(!new_pds_crawl::seed_listed(&state.db, &seed)?); + Ok(()) + } + #[tokio::test] async fn apply_seed_snapshot_persists_statuses_and_seeds_cursors() -> miette::Result<()> { let tmp = tempdir().into_diagnostic()?; diff --git a/src/crawler/mod.rs b/src/crawler/mod.rs index 4c71c58..49c8a34 100644 --- a/src/crawler/mod.rs +++ b/src/crawler/mod.rs @@ -14,6 +14,7 @@ use url::Url; mod by_collection; mod list_repos; +mod new_pds; mod pds_repos; mod relay_first; mod worker; @@ -22,6 +23,8 @@ use crate::util::throttle::Throttler; pub(crate) use by_collection::ByCollectionProducer; pub(crate) use list_repos::{ListReposProducer, RetryProducer, SignalChecker}; +pub(crate) use new_pds::NewPdsCrawler; +pub(crate) use pds_repos::PdsRepoLister; pub(crate) use relay_first::RelayFirstProducer; pub(crate) use worker::{CrawlerBatch, CrawlerWorker}; diff --git a/src/crawler/new_pds.rs b/src/crawler/new_pds.rs new file mode 100644 index 0000000..8aa59ff --- /dev/null +++ b/src/crawler/new_pds.rs @@ -0,0 +1,373 @@ +use std::time::Duration; + +use miette::{IntoDiagnostic, Result}; +use smol_str::SmolStr; +use tokio::task::JoinSet; +use tracing::{error, info, warn}; +use url::Url; + +use super::pds_repos::{PdsListing, PdsRepoLister}; +use super::worker::CursorUpdate; +use crate::control::new_pds_crawl; +use crate::db::keys; +use crate::net::{PublicHost, PublicHttpClient}; +use crate::util::WatchEnabledExt; + +const MAX_CONCURRENT_CRAWLS: usize = 4; +const RETRY_DELAY: Duration = Duration::from_secs(5 * 60); + +/// runs the one-off `listRepos` walk queued for each PDS we start following, +/// see [`new_pds_crawl`]. +#[derive(Clone)] +pub(crate) struct NewPdsCrawler { + pub(crate) lister: PdsRepoLister, + pub(crate) http: PublicHttpClient, +} + +impl NewPdsCrawler { + pub(crate) async fn run(mut self) { + loop { + self.lister + .enabled + .wait_enabled(self.lister.component) + .await; + let all_done = self.crawl_queued().await.unwrap_or_else(|err| { + error!(err = ?err, "failed to read queued new PDS crawls"); + false + }); + + // a wake that lands mid round is kept by the notify, so the next round sees its host + let state = self.lister.checker.state.clone(); + let woken = state.new_pds_crawl_notify.notified(); + if all_done { + woken.await; + } else { + tokio::select! { + _ = woken => {} + _ = tokio::time::sleep(RETRY_DELAY) => {} + } + } + } + } + + /// returns whether every queued host got through. + async fn crawl_queued(&self) -> Result { + let queued = self.lister.checker.state.db.run(load_queued).await?; + let mut all_done = true; + let mut crawls = JoinSet::new(); + for (host, cursor) in queued { + if crawls.len() >= MAX_CONCURRENT_CRAWLS + && let Some(result) = crawls.join_next().await + { + all_done &= crawl_succeeded(result); + } + let crawler = self.clone(); + crawls.spawn(async move { + let result = crawler.crawl(&host, cursor).await; + (host, result) + }); + } + while let Some(result) = crawls.join_next().await { + all_done &= crawl_succeeded(result); + } + Ok(all_done) + } + + async fn crawl(self, host: &str, cursor: Option) -> Result<()> { + let key = keys::new_pds_crawl_key(host); + let public = match PublicHost::parse(host) { + Ok(public) => public, + Err(error) => { + warn!(%host, %error, "queued PDS is not a public host, dropping its crawl"); + return self.forget(key).await; + } + }; + // an operator removing the source means they don't want this host either + let source_key = keys::firehose_source_key(public.firehose_url().as_str()); + let followed = self + .lister + .checker + .state + .db + .run(move |db| db.crawler.contains_key(source_key).into_diagnostic()) + .await?; + if !followed { + info!(%host, "PDS is no longer followed, dropping its crawl"); + return self.forget(key).await; + } + + info!(%host, resuming = cursor.is_some(), "listing repos of a newly followed PDS"); + let pds = public.https_endpoint().url().clone(); + self.list_host(host, key, &pds, cursor).await + } + + async fn list_host( + mut self, + host: &str, + key: Vec, + pds: &Url, + cursor: Option, + ) -> Result<()> { + let http = self.http.clone(); + // we don't know its size, and one account paces it at the tier's base rate + let listing = self + .lister + .list(pds, 1, &http, cursor, |next| { + Ok(vec![CursorUpdate { + key: key.clone(), + // the last page deletes the key in the same commit as its repos + value: next + .map(|next| new_pds_crawl::encode_cursor(Some(next))) + .transpose()?, + }]) + }) + .await?; + match listing { + PdsListing::Complete => info!(%host, "listed repos of a newly followed PDS"), + PdsListing::Banned => { + info!(%host, "PDS was banned, dropping its crawl"); + self.forget(key).await?; + } + } + Ok(()) + } + + async fn forget(&self, key: Vec) -> Result<()> { + self.lister + .send_cursor_updates(vec![CursorUpdate { key, value: None }]) + .await + } +} + +fn load_queued(db: &crate::db::Db) -> Result)>> { + db.cursors + .prefix(keys::NEW_PDS_CRAWL_PREFIX) + .map(|entry| { + let (key, value) = entry.into_inner().into_diagnostic()?; + let host = std::str::from_utf8(&key[keys::NEW_PDS_CRAWL_PREFIX.len()..]) + .into_diagnostic()? + .to_owned(); + Ok((host, new_pds_crawl::decode_cursor(&value)?)) + }) + .collect() +} + +fn crawl_succeeded( + result: std::result::Result<(String, Result<()>), tokio::task::JoinError>, +) -> bool { + match result { + Ok((_, Ok(()))) => true, + Ok((host, Err(err))) => { + warn!(%host, err = ?err, "new PDS crawl failed, retrying later"); + false + } + Err(err) => { + error!(err = ?err, "new PDS crawl task failed"); + false + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::Config; + use crate::crawler::{CrawlerStats, CrawlerWorker, InFlight, SignalChecker}; + use crate::state::AppState; + use axum::extract::{Query, State}; + use axum::response::{IntoResponse, Response}; + use axum::{Json, Router, http::StatusCode, routing::get}; + use jacquard_common::types::string::Did; + use serde_json::json; + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use tempfile::tempdir; + + const HEAD: &str = "bafyreidf747c4x3lps3k4n357l3a3r57k3k465743k573k465743k5"; + + struct MockPds { + requests: AtomicUsize, + fail_second_page: AtomicBool, + } + + /// serves two `listRepos` pages, one unknown repo on each. + async fn mock_pds() -> Result<(std::net::SocketAddr, Arc)> { + async fn list_repos( + State(pds): State>, + Query(query): Query>, + ) -> Response { + pds.requests.fetch_add(1, Ordering::SeqCst); + if !query.contains_key("cursor") { + return Json(json!({ + "cursor": "next", + "repos": [{"did": "did:web:first.example", "head": HEAD, "rev": "3j6s746574657"}] + })) + .into_response(); + } + if pds.fail_second_page.load(Ordering::SeqCst) { + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + Json(json!({ + "repos": [{"did": "did:web:second.example", "head": HEAD, "rev": "3j6s746574657"}] + })) + .into_response() + } + + let pds = Arc::new(MockPds { + requests: AtomicUsize::new(0), + fail_second_page: AtomicBool::new(false), + }); + let app = Router::new() + .route("/xrpc/com.atproto.sync.listRepos", get(list_repos)) + .with_state(pds.clone()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .into_diagnostic()?; + let addr = listener.local_addr().into_diagnostic()?; + tokio::spawn(async move { axum::serve(listener, app).await }); + Ok((addr, pds)) + } + + /// a crawler whose `pds.example` points at `addr`, with a running worker. + fn crawler(state: &Arc, addr: std::net::SocketAddr) -> Result { + let stats = CrawlerStats::new(state.clone(), Vec::new(), state.throttler.clone()); + let (worker, tx) = CrawlerWorker::new(state.clone(), 100, 50, stats.clone()); + tokio::spawn(worker.run()); + Ok(NewPdsCrawler { + lister: PdsRepoLister { + checker: SignalChecker { + http: PublicHttpClient::test(reqwest::Client::new()), + state: state.clone(), + throttler: state.throttler.clone(), + }, + in_flight: InFlight::new(), + tx, + enabled: state.crawler_enabled.subscribe(), + stats, + component: "new PDS crawler", + }, + http: PublicHttpClient::test( + reqwest::Client::builder() + .resolve("pds.example", addr) + .build() + .into_diagnostic()?, + ), + }) + } + + fn test_state(dir: &tempfile::TempDir) -> Result> { + Ok(Arc::new(AppState::new(&Config { + database_path: dir.path().to_path_buf(), + full_network: true, + enable_crawler: Some(true), + ..Default::default() + })?)) + } + + fn queue(state: &AppState, host: &str, cursor: Option<&str>) -> Result<()> { + let cursor = cursor.map(SmolStr::new); + state + .db + .cursors + .insert( + keys::new_pds_crawl_key(host), + new_pds_crawl::encode_cursor(cursor.as_ref())?, + ) + .into_diagnostic() + } + + fn tracked(state: &AppState, did: &'static str) -> Result { + state + .db + .repos + .contains_key(keys::repo_key(&Did::new_static(did)?)) + .into_diagnostic() + } + + /// the worker commits batches on its own task, so wait for it to catch up. + async fn eventually(check: impl Fn() -> Result) -> Result<()> { + tokio::time::timeout(Duration::from_secs(5), async { + while !check()? { + tokio::time::sleep(Duration::from_millis(10)).await; + } + Ok(()) + }) + .await + .into_diagnostic()? + } + + #[tokio::test] + async fn a_broken_walk_keeps_its_cursor_and_the_last_page_clears_it() -> Result<()> { + let dir = tempdir().into_diagnostic()?; + let state = test_state(&dir)?; + let (addr, pds) = mock_pds().await?; + let crawler = crawler(&state, addr)?; + let url = Url::parse("http://pds.example/").into_diagnostic()?; + let key = keys::new_pds_crawl_key("pds.example"); + queue(&state, "pds.example", None)?; + + pds.fail_second_page.store(true, Ordering::SeqCst); + assert!( + crawler + .clone() + .list_host("pds.example", key.clone(), &url, None) + .await + .is_err() + ); + eventually(|| { + Ok(tracked(&state, "did:web:first.example")? + && load_queued(&state.db)? == vec![("pds.example".into(), Some("next".into()))]) + }) + .await?; + + // the 500 put the PDS in backoff, which the producer waits out between rounds + pds.fail_second_page.store(false, Ordering::SeqCst); + state.throttler.get_handle(&url).await.record_success(); + let [(host, cursor)] = <[_; 1]>::try_from(load_queued(&state.db)?).unwrap(); + crawler.list_host(&host, key, &url, cursor).await?; + eventually(|| { + Ok(tracked(&state, "did:web:second.example")? && load_queued(&state.db)?.is_empty()) + }) + .await?; + // first page, the failed second page, then the second page again + assert_eq!(pds.requests.load(Ordering::SeqCst), 3); + Ok(()) + } + + #[tokio::test] + async fn banned_unfollowed_and_private_hosts_drop_their_crawl() -> Result<()> { + let dir = tempdir().into_diagnostic()?; + let state = test_state(&dir)?; + let (addr, pds) = mock_pds().await?; + let crawler = crawler(&state, addr)?; + + queue(&state, "pds.example", None)?; + state.pds_meta.rcu(|meta| { + let mut next = (**meta).clone(); + next.update_host_entry("pds.example", |entry| { + entry.status = crate::pds_meta::HostStatus::Banned; + }); + next + }); + let url = Url::parse("http://pds.example/").into_diagnostic()?; + crawler + .clone() + .list_host( + "pds.example", + keys::new_pds_crawl_key("pds.example"), + &url, + None, + ) + .await?; + + queue(&state, "gone.example", None)?; + crawler.clone().crawl("gone.example", None).await?; + queue(&state, "127.0.0.1", None)?; + crawler.crawl("127.0.0.1", None).await?; + + eventually(|| Ok(load_queued(&state.db)?.is_empty())).await?; + assert_eq!(pds.requests.load(Ordering::SeqCst), 0); + Ok(()) + } +} diff --git a/src/db/keys/indexer.rs b/src/db/keys/indexer.rs index 4a8b3c7..af9af45 100644 --- a/src/db/keys/indexer.rs +++ b/src/db/keys/indexer.rs @@ -291,6 +291,28 @@ pub fn relay_first_pds_cursor_key(relay: &str, host: &str) -> Vec { relay_first_cursor_key(relay, &format!("pds|{host}")) } +/// key format: `new_pds_crawl|`, a PDS we started following that still +/// owes us one `listRepos` walk. the value is its page cursor. +pub const NEW_PDS_CRAWL_PREFIX: &[u8] = b"new_pds_crawl|"; + +pub fn new_pds_crawl_key(host: &str) -> Vec { + let mut key = Vec::with_capacity(NEW_PDS_CRAWL_PREFIX.len() + host.len()); + key.extend_from_slice(NEW_PDS_CRAWL_PREFIX); + key.extend_from_slice(host.as_bytes()); + key +} + +/// key format: `seed_listed|`, set once a seed relay's `listHosts` +/// has been read to the end with every host on it admitted. +pub const SEED_LISTED_PREFIX: &[u8] = b"seed_listed|"; + +pub fn seed_listed_key(seed: &str) -> Vec { + let mut key = Vec::with_capacity(SEED_LISTED_PREFIX.len() + seed.len()); + key.extend_from_slice(SEED_LISTED_PREFIX); + key.extend_from_slice(seed.as_bytes()); + key +} + pub const BY_COLLECTION_CURSOR_PREFIX: &[u8] = b"by_collection_cursor|"; /// prefix for all by-collection cursors belonging to a given index URL. diff --git a/src/state.rs b/src/state.rs index 93f3712..76c897f 100644 --- a/src/state.rs +++ b/src/state.rs @@ -52,6 +52,8 @@ pub struct AppState { #[cfg(feature = "indexer")] pub backfill_notify: Notify, #[cfg(feature = "indexer")] + pub(crate) new_pds_crawl_notify: Notify, + #[cfg(feature = "indexer")] pub crawler_enabled: watch::Sender, #[cfg(feature = "indexer")] pub backfill_enabled: watch::Sender, @@ -152,6 +154,8 @@ impl AppState { #[cfg(feature = "indexer")] backfill_notify: Notify::new(), #[cfg(feature = "indexer")] + new_pds_crawl_notify: Notify::new(), + #[cfg(feature = "indexer")] crawler_enabled, firehose_enabled, #[cfg(feature = "indexer")]