very fast at protocol indexer with flexible filtering, xrpc queries, cursor-backed event stream, and more, built on fjall
rust fjall at-protocol atproto indexer
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374use 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<bool> { 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<SmolStr>) -> 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<u8>, pds: &Url, cursor: Option<SmolStr>, ) -> 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<u8>) -> Result<()> { self.lister .send_cursor_updates(vec![CursorUpdate { key, value: None }]) .await }}
fn load_queued(db: &crate::db::Db) -> Result<Vec<(String, Option<SmolStr>)>> { 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<MockPds>)> { async fn list_repos( State(pds): State<Arc<MockPds>>, Query(query): Query<HashMap<String, String>>, ) -> 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<AppState>, addr: std::net::SocketAddr) -> Result<NewPdsCrawler> { 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<Arc<AppState>> { 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<bool> { 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<bool>) -> 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(()) }}