use jacquard_common::{IntoStatic, types::{string::Did, tid::Tid}};
use miette::{Context, IntoDiagnostic, Result};
use reqwest::StatusCode;
use serde::Deserialize;
use smol_str::{SmolStr, ToSmolStr};
use super::super::worker::ListedRepoStatus;
#[derive(Deserialize)]
pub(crate) struct ListedRepo<'a> {
pub active: Option<bool>,
pub did: Did,
pub rev: Tid,
#[serde(borrow)]
pub status: Option<crate::ingest::stream::AccountStatus<'a>>,
}
#[derive(Deserialize)]
pub(crate) struct ListReposOutput<'a> {
pub cursor: Option<SmolStr>,
#[serde(borrow)]
pub repos: Vec<ListedRepo<'a>>,
}
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, watch};
use tracing::{debug, error, info, warn};
use url::Url;
use super::super::worker::{
CrawlerBatch, CursorUpdate, DispatchWait, RepoListing, WAITING_PAGES, WaitingPages, send_page,
};
use super::super::{CrawlerStats, InFlight, base_url};
use super::checker::SignalChecker;
use crate::db::keys::crawler_cursor_key;
use crate::db::{Db, keys};
use crate::state::AppState;
use crate::util::WatchEnabledExt;
use crate::util::{ErrorForStatus, RetryOutcome, RetryWithBackoff};
const BLOCKING_TASK_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_LIST_REPOS_BODY_BYTES: usize = 10 * 1024 * 1024;
pub(crate) struct ListReposProducer {
pub(crate) url: Url,
pub(crate) checker: SignalChecker,
pub(crate) in_flight: InFlight,
pub(crate) tx: mpsc::Sender<CrawlerBatch>,
pub(crate) enabled: watch::Receiver<bool>,
pub(crate) stats: CrawlerStats,
pub(crate) completion_key: Option<Vec<u8>>,
pub(crate) last_completion: Option<u64>,
}
impl ListReposProducer {
pub(crate) async fn run(mut self) -> Result<()> {
let waiting = WaitingPages::spawn(
self.tx.clone(),
self.checker.state.clone(),
DispatchWait::DEFAULT,
WAITING_PAGES,
);
loop {
if let Err(e) = self.crawl(&waiting).await {
error!(err = ?e, relay = %self.url, "fatal relay crawl error, restarting in 30s");
tokio::time::sleep(Duration::from_secs(30)).await;
}
}
}
async fn get_cursor(&self) -> Result<Option<SmolStr>> {
let key = crawler_cursor_key(self.url.as_str());
let cursor_bytes = self
.checker
.state
.db
.run(move |db| db.cursors.get(key).into_diagnostic())
.await?;
Ok(cursor_bytes
.as_deref()
.and_then(|b| rmp_serde::from_slice::<SmolStr>(b).ok()))
}
async fn next_completion(&mut self, pass_complete: bool) -> Result<Option<(Vec<u8>, u64)>> {
let Some(key) = self.completion_key.clone().filter(|_| pass_complete) else {
return Ok(None);
};
let lookup_key = key.clone();
let current = self
.checker
.state
.db
.run(move |db| {
db.cursors
.get(lookup_key)
.into_diagnostic()?
.map(|value| {
value
.as_ref()
.try_into()
.into_diagnostic()
.map(u64::from_be_bytes)
})
.transpose()
})
.await?;
let next = next_completion_number(current, self.last_completion)?;
self.last_completion = Some(next);
Ok(Some((key, next)))
}
async fn crawl(&mut self, waiting: &WaitingPages) -> Result<()> {
let db = self.checker.state.db.clone();
let base = base_url(&self.url)?;
waiting.begin_crawl();
let mut cursor = self.get_cursor().await?;
match &cursor {
Some(c) => info!(cursor = %c, "resuming"),
None => info!("starting from scratch"),
}
loop {
self.enabled.wait_enabled("crawler").await;
let mut url = base
.join("/xrpc/com.atproto.sync.listRepos")
.into_diagnostic()?;
url.query_pairs_mut().append_pair("limit", "1000");
if let Some(c) = &cursor {
url.query_pairs_mut().append_pair("cursor", c.as_str());
}
// the relay source is operator configuration; only the PDS
// returned by it enters the public checker below.
let fetch_result = (|| {
crate::net::trusted_http()
.get(url.clone())
.send()
.error_for_status()
})
.retry(5, |e: &reqwest::Error, attempt| {
matches!(e.status(), Some(StatusCode::TOO_MANY_REQUESTS))
.then(|| Duration::from_secs(1 << attempt.min(5)))
})
.await;
let res = match fetch_result {
Ok(r) => r,
Err(RetryOutcome::Ratelimited) => {
warn!("rate limited by relay after retries");
miette::bail!("rate limited by relay after retries");
}
Err(RetryOutcome::Failed(e)) => {
error!(err = %e, "crawler failed to fetch listRepos");
return Err(e)
.into_diagnostic()
.wrap_err("crawler failed to fetch listRepos");
}
};
// the listing is a snapshot from about now, which the firehose may not have caught
// up to yet
let read_at = chrono::Utc::now().timestamp_millis();
let bytes = match crate::util::read_limited_bytes(res, MAX_LIST_REPOS_BODY_BYTES).await
{
Ok(bytes) => bytes,
Err(error) => {
error!(err = %error, "can't read bounded listRepos response");
return Err(error).wrap_err("can't read listRepos response");
}
};
let filter = self.checker.state.filter.load_full();
struct ParseResult {
unknown_dids: Vec<Did>,
listings: Vec<RepoListing>,
/// listings that claim a newer rev than stored, see [`send_page`]
newer: Vec<RepoListing>,
cursor: Option<SmolStr>,
count: usize,
}
let parse_result = {
let repos = db.repos.clone();
let filter_ks = db.filter.clone();
let crawler_ks = db.crawler.clone();
let persist_inactive_unknown = self.completion_key.is_some();
let skew_secs = self.checker.state.rev_clock_skew_secs;
match tokio::time::timeout(
BLOCKING_TASK_TIMEOUT,
// CPU-bound JSON parsing + read-only DB scan, stays on spawn_blocking
tokio::task::spawn_blocking(move || -> Result<ParseResult> {
let output = serde_json::from_slice::<ListReposOutput>(&bytes)
.into_diagnostic()
.wrap_err("failed to parse listRepos response")?;
let count = output.repos.len();
let next_cursor = output.cursor;
let mut unknown = Vec::new();
let mut listings = Vec::new();
let mut newer = Vec::new();
for repo in output.repos {
let excl_key = crate::db::filter::exclude_key(repo.did.as_str())?;
if filter_ks.contains_key(&excl_key).into_diagnostic()? {
continue;
}
let did_key = keys::repo_key(&repo.did);
if let Some(stored) = repos.get(&did_key).into_diagnostic()? {
if let Some(active) = repo.active {
let listing = RepoListing {
did: repo.did.into_static(),
rev: crate::db::types::DbTid::from(&repo.rev),
active,
status: repo.status.into_static(),
create_if_missing: false,
};
let stored = crate::db::deser_repo_state(&stored)?;
if listing.claims_newer(&stored, skew_secs) {
newer.push(listing);
} else {
listings.push(listing);
}
}
continue;
}
let retry_key = keys::crawler_retry_key(&repo.did);
if crawler_ks.contains_key(&retry_key).into_diagnostic()? {
continue;
}
if persist_inactive_unknown
&& should_persist_inactive_unknown(
repo.active,
repo.status.is_some(),
)
{
listings.push(RepoListing {
did: repo.did.into_static(),
rev: crate::db::types::DbTid::from(&repo.rev),
active: false,
status: repo.status.into_static(),
create_if_missing: true,
});
continue;
}
if should_enqueue_unknown(repo.active) {
unknown.push(repo.did.into_static());
}
}
Ok(ParseResult {
unknown_dids: unknown,
listings,
newer,
cursor: next_cursor,
count,
})
}),
)
.await
{
Ok(res) => res
.into_diagnostic()
.wrap_err("blocking task failed to parse listRepos")?,
Err(_) => {
error!(
"spawn_blocking task for parsing listRepos timed out after {}s",
BLOCKING_TASK_TIMEOUT.as_secs()
);
miette::bail!("spawn_blocking task for parsing listRepos timed out");
}
}
};
let ParseResult {
unknown_dids,
listings,
newer,
cursor: next_cursor,
count,
} = parse_result.wrap_err("error while crawling")?;
debug!(count, "fetched repos");
self.stats.record_crawled(count);
let pass_complete;
if let Some(new_cursor) = next_cursor {
cursor = Some(new_cursor.as_str().into());
pass_complete = false;
} else {
info!("reached end of list.");
cursor = None;
pass_complete = true;
}
let completion = self.next_completion(pass_complete).await?;
let cursor_updates =
pass_cursor_updates(self.url.as_str(), cursor.as_ref(), completion)?;
let in_flight = self.in_flight.acquire(unknown_dids).await;
let confirmed = if filter.check_signals() && !in_flight.is_empty() {
let mut retry_batch = db.inner.batch();
let confirmed = self
.checker
.check_signals_batch(in_flight, &filter, &mut retry_batch, &HashMap::new())
.await?;
tokio::time::timeout(
BLOCKING_TASK_TIMEOUT,
db.run(move |_db| retry_batch.commit().into_diagnostic()),
)
.await
.map_err(|_| miette::miette!("retry state commit timed out"))??;
confirmed
} else {
in_flight
};
let page = CrawlerBatch {
guards: confirmed,
listings,
cursor_updates,
applied: None,
};
send_page(
&self.tx,
waiting,
base.host_str().unwrap_or_default(),
read_at,
page,
newer,
)
.await?;
if pass_complete {
info!("enumeration complete, sleeping 1h before next pass");
tokio::select! {
_ = tokio::time::sleep(Duration::from_secs(3600)) => {
info!("resuming after 1h sleep");
}
_ = self.enabled.changed() => {
if !*self.enabled.borrow() { return Ok(()); }
}
}
}
}
}
}
fn should_enqueue_unknown(active: Option<bool>) -> bool {
active != Some(false)
}
fn should_persist_inactive_unknown(active: Option<bool>, has_status: bool) -> bool {
active == Some(false) || (active.is_none() && has_status)
}
fn next_completion_number(persisted: Option<u64>, last_issued: Option<u64>) -> Result<u64> {
persisted
.into_iter()
.chain(last_issued)
.max()
.unwrap_or(0)
.checked_add(1)
.ok_or_else(|| miette::miette!("relay-first pass counter overflow"))
}
fn pass_cursor_updates(
relay: &str,
cursor: Option<&SmolStr>,
completion: Option<(Vec<u8>, u64)>,
) -> Result<Vec<CursorUpdate>> {
let cursor_value = cursor
.map(|cursor| {
rmp_serde::to_vec(cursor.as_str())
.into_diagnostic()
.wrap_err("cant serialize cursor")
})
.transpose()?;
let mut updates = vec![CursorUpdate {
key: crawler_cursor_key(relay),
value: cursor_value,
}];
if let Some((key, pass)) = completion {
updates.push(CursorUpdate {
key,
value: Some(pass.to_be_bytes().to_vec()),
});
}
Ok(updates)
}
#[cfg(test)]
mod tests {
use super::{
ListReposProducer, cursor_display, next_completion_number, pass_cursor_updates,
should_enqueue_unknown, should_persist_inactive_unknown,
};
use crate::config::Config;
use crate::crawler::{CrawlerStats, CrawlerWorker, InFlight, SignalChecker};
use crate::db::keys;
use crate::state::AppState;
use axum::{Json, Router, routing::get};
use jacquard_common::types::string::Did;
use miette::{IntoDiagnostic, Result};
use serde_json::json;
use std::sync::Arc;
use std::time::Duration;
use tempfile::tempdir;
use url::Url;
#[test]
fn inactive_unknown_repo_is_not_backfilled() {
assert!(!should_enqueue_unknown(Some(false)));
assert!(should_enqueue_unknown(Some(true)));
assert!(should_enqueue_unknown(None));
}
#[test]
fn relay_first_persists_authoritative_inactive_unknown_repo() {
assert!(should_persist_inactive_unknown(Some(false), false));
assert!(should_persist_inactive_unknown(None, true));
assert!(!should_persist_inactive_unknown(Some(true), true));
assert!(!should_persist_inactive_unknown(None, false));
}
#[test]
fn relay_first_marker_is_only_written_with_the_terminal_cursor() {
let completion_key = b"relay_first_cursor|relay|pass";
let cursor = "next".into();
let mid_pass = pass_cursor_updates("https://relay.example", Some(&cursor), None).unwrap();
assert_eq!(mid_pass.len(), 1);
assert!(mid_pass[0].value.is_some());
let terminal = pass_cursor_updates(
"https://relay.example",
None,
Some((completion_key.to_vec(), 7)),
)
.unwrap();
assert_eq!(terminal.len(), 2);
assert!(terminal[0].value.is_none());
assert_eq!(terminal[1].key, completion_key);
assert_eq!(
terminal[1].value.as_deref(),
Some(7u64.to_be_bytes().as_slice())
);
}
#[test]
fn relay_first_marker_advances_past_an_uncommitted_marker() {
assert_eq!(next_completion_number(Some(5), None).unwrap(), 6);
assert_eq!(next_completion_number(Some(5), Some(6)).unwrap(), 7);
assert!(next_completion_number(Some(u64::MAX), None).is_err());
}
#[tokio::test]
async fn relay_first_marker_continues_from_the_stored_big_endian_pass() -> Result<()> {
let tmp = tempdir().into_diagnostic()?;
let state = Arc::new(AppState::new(&Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})?);
let pass_key = keys::relay_first_pass_key("http://relay.example/");
state
.db
.cursors
.insert(&pass_key, [0, 0, 0, 0, 0, 0, 1, 2])
.into_diagnostic()?;
let (tx, _rx) = tokio::sync::mpsc::channel(1);
let mut producer = ListReposProducer {
url: Url::parse("http://relay.example/").into_diagnostic()?,
checker: SignalChecker {
http: crate::net::PublicHttpClient::test(reqwest::Client::new()),
state: state.clone(),
throttler: state.throttler.clone(),
},
in_flight: InFlight::new(),
tx,
enabled: state.crawler_enabled.subscribe(),
stats: CrawlerStats::new(state.clone(), Vec::new(), state.throttler.clone()),
completion_key: Some(pass_key.clone()),
last_completion: None,
};
assert_eq!(
producer.next_completion(true).await?,
Some((pass_key.clone(), 259))
);
state
.db
.cursors
.insert(&pass_key, [0; 7])
.into_diagnostic()?;
assert!(producer.next_completion(true).await.is_err());
Ok(())
}
#[tokio::test]
async fn crawler_cursor_display_decodes_the_msgpack_cursor() -> Result<()> {
let tmp = tempdir().into_diagnostic()?;
let state = AppState::new(&Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})?;
let relay = Url::parse("https://relay.example/").into_diagnostic()?;
assert_eq!(cursor_display(&state, &relay).await, "none");
let key = keys::crawler_cursor_key(relay.as_str());
state
.db
.cursors
.insert(&key, b"\xa3abc")
.into_diagnostic()?;
assert_eq!(cursor_display(&state, &relay).await, "abc");
state.db.cursors.insert(&key, b"abc").into_diagnostic()?;
assert_eq!(cursor_display(&state, &relay).await, "none");
Ok(())
}
#[tokio::test]
async fn relay_first_pass_persists_relay_inactive_unknown_as_authoritative() -> Result<()> {
const DID: &str = "did:web:relay-inactive.example";
const HEAD: &str = "bafyreidf747c4x3lps3k4n357l3a3r57k3k465743k573k465743k5";
let app = Router::new().route(
"/xrpc/com.atproto.sync.listRepos",
get(|| async {
Json(json!({
"repos": [{
"did": DID,
"head": HEAD,
"rev": "3j6s746574657",
"active": false,
"status": "takendown"
}]
}))
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.into_diagnostic()?;
let address = listener.local_addr().into_diagnostic()?;
let server = tokio::spawn(async move { axum::serve(listener, app).await });
let tmp = tempdir().into_diagnostic()?;
let config = Config {
database_path: tmp.path().to_path_buf(),
full_network: true,
enable_crawler: Some(true),
..Default::default()
};
let state = Arc::new(AppState::new(&config)?);
let stats = CrawlerStats::new(state.clone(), Vec::new(), state.throttler.clone());
let (worker, tx) = CrawlerWorker::new(state.clone(), 100, 50, stats.clone());
let worker = tokio::spawn(worker.run());
let pass_key = keys::relay_first_pass_key("http://relay.example/");
let producer = tokio::spawn(
ListReposProducer {
url: Url::parse(&format!("http://{address}/")).into_diagnostic()?,
checker: SignalChecker {
http: crate::net::PublicHttpClient::test(reqwest::Client::new()),
state: state.clone(),
throttler: state.throttler.clone(),
},
in_flight: InFlight::new(),
tx,
enabled: state.crawler_enabled.subscribe(),
stats,
completion_key: Some(pass_key.clone()),
last_completion: None,
}
.run(),
);
tokio::time::timeout(Duration::from_secs(5), async {
loop {
if state.db.cursors.get(&pass_key).ok().flatten().is_some() {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.into_diagnostic()?;
let did = Did::new_static(DID)?;
let bytes = state
.db
.repos
.get(keys::repo_key(&did))
.into_diagnostic()?
.expect("relay-inactive repo should be persisted");
let repo = crate::db::deser_repo_state(&bytes)?;
assert!(!repo.active);
assert_eq!(repo.status, crate::types::RepoStatus::Takendown);
assert_eq!(state.db.indexer.pending_iter().count(), 0);
state.db.indexer.assert_pending_dids_match();
producer.abort();
worker.abort();
server.abort();
Ok(())
}
}
pub(crate) async fn cursor_display(state: &AppState, relay_host: &Url) -> SmolStr {
let key = crawler_cursor_key(relay_host.as_str());
let cursor_bytes = match state
.db
.run(move |db| db.cursors.get(key).into_diagnostic())
.await
{
Ok(b) => b,
Err(e) => return e.to_smolstr(),
};
match cursor_bytes
.as_deref()
.map(rmp_serde::from_slice::<SmolStr>)
{
None | Some(Err(_)) => "none".to_smolstr(),
Some(Ok(c)) => c,
}
}