use std::collections::{HashMap, HashSet};
use std::time::Duration;
use super::list_repos::producer::{ListReposOutput, ListedRepo};
use jacquard_common::{IntoStatic, types::string::Did};
use miette::{Context, IntoDiagnostic, Result};
use reqwest::StatusCode;
use smol_str::SmolStr;
use tokio::sync::{mpsc, watch};
use tracing::debug;
use url::Url;
use super::worker::{CrawlerBatch, CursorUpdate, RepoListing};
use super::{CrawlerStats, InFlight, SignalChecker};
use crate::control::new_pds_crawl::Budget;
use crate::db::counts::pds_host;
use crate::db::types::DbTid;
use crate::db::{Db, keys};
use crate::ingest::indexer::message::RepoOp;
use crate::ingest::indexer::repos::{on_owners, repo_ops};
use crate::net::{PublicEndpoint, PublicHttpClient};
use crate::state::AppState;
use crate::types::{RepoState, RepoStatus, ResyncState};
use crate::util::{WatchEnabledExt, parse_retry_after};
/// an honest page of 1000 repos is about 145 KiB
const MAX_LIST_RESPONSE_BYTES: usize = 512 * 1024;
const LIST_LIMIT: usize = 1000;
const BLOCKING_TASK_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_PAGES: usize = 100_000;
/// how a walk over one PDS ended.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum PdsListing {
Complete,
/// the host was banned locally partway through, so the rest was skipped.
Banned,
/// the walk's bytes are spent. it stopped after the last page it committed, `until` is the
/// unix time its window ends.
OutOfBudget {
until: i64,
},
}
/// what a walk does with the known repos it lists whose doc points at the walked host.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum KnownRepos {
Skip,
/// resync the ones listed past the rev we hold
ResyncNewer,
/// stuck repos are resynced even when listed at the rev we hold
ResyncNewerOrStuck,
}
/// walks one PDS's `listRepos` and hands the repos we don't know yet to the crawler worker.
#[derive(Clone)]
pub(crate) struct PdsRepoLister {
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) component: &'static str,
}
impl PdsRepoLister {
/// `progress` gets the cursor of the next page, or `None` once the walk is done, and
/// returns the cursor writes to commit together with the page that produced it.
pub(crate) async fn list(
&mut self,
pds: &Url,
account_count: u64,
http: &PublicHttpClient,
mut cursor: Option<SmolStr>,
known: KnownRepos,
mut budget: Option<&mut Budget>,
progress: impl Fn(Option<&SmolStr>) -> Result<Vec<CursorUpdate>>,
) -> Result<PdsListing> {
let hostname = pds
.host_str()
.ok_or_else(|| miette::miette!("PDS URL has no hostname"))?;
let mut page_count = 0usize;
let mut admitted = 0u64;
let mut seen: HashSet<u64> = cursor.iter().map(crate::util::hash).collect();
loop {
page_count += 1;
if page_count > MAX_PAGES {
miette::bail!("PDS listRepos exceeded {MAX_PAGES} pages");
}
self.enabled.wait_enabled(self.component).await;
if self.checker.state.pds_meta.load().is_banned(hostname) {
debug!(%hostname, "locally banned PDS, stopping listRepos");
return Ok(PdsListing::Banned);
}
if let Some(until) = budget
.as_deref()
.and_then(|budget| budget.spent_until(chrono::Utc::now().timestamp()))
{
debug!(%hostname, until, "listRepos bytes spent, stopping the walk");
return Ok(PdsListing::OutOfBudget { until });
}
let mut url = pds
.join("/xrpc/com.atproto.sync.listRepos")
.into_diagnostic()?;
url.query_pairs_mut()
.append_pair("limit", &LIST_LIMIT.to_string());
if let Some(value) = &cursor {
url.query_pairs_mut().append_pair("cursor", value);
}
let Some(bytes) = self
.fetch_page(pds, hostname, account_count, url, http)
.await?
else {
self.send_cursor_updates(progress(None)?).await?;
return Ok(PdsListing::Complete);
};
if let Some(budget) = budget.as_deref_mut() {
let (key, value) = budget.charge(chrono::Utc::now().timestamp(), bytes.len())?;
self.send_cursor_updates(vec![CursorUpdate {
key,
value: Some(value),
}])
.await?;
}
let page = self
.parse_page(bytes, known, SmolStr::new(hostname))
.await?;
self.stats.record_crawled(page.count);
// before the page's cursor moves. a walk cut short resyncs them next time
self.resync(page.stuck).await?;
let state = &self.checker.state;
let mut unknown = page.unknown_dids;
let mut skipped = 0;
// a host with no room left isn't worth a signal check per account
if room_left(state, hostname, admitted, 1) == 0 {
skipped = unknown.len();
unknown.clear();
}
let in_flight = self.in_flight.acquire(unknown).await;
let filter = self.checker.state.filter.load_full();
let mut confirmed = if filter.check_signals() && !in_flight.is_empty() {
let mut retry_batch = self.checker.state.db.inner.batch();
let confirmed = self
.checker
.check_signals_batch(in_flight, &filter, &mut retry_batch, &HashMap::new())
.await?;
self.checker
.state
.db
.run(move |_db| retry_batch.commit().into_diagnostic())
.await?;
confirmed
} else {
in_flight
};
// the limit counts the accounts that passed the signal check, since the ones that
// didn't never count for the host
skipped += keep_under_account_limit(state, hostname, admitted, &mut confirmed);
if skipped > 0 {
debug!(%hostname, skipped, "account limit reached for host, skipping crawled accounts");
self.stats.record_over_limit(skipped);
}
let next_cursor = page.cursor;
if let Some(next) = &next_cursor
&& !seen.insert(crate::util::hash(next))
{
miette::bail!("PDS listRepos returned a cursor it already returned");
}
admitted += confirmed.len() as u64;
self.tx
.send(CrawlerBatch {
guards: confirmed,
listings: page.newer,
cursor_updates: progress(next_cursor.as_ref())?,
applied: None,
})
.await
.into_diagnostic()
.wrap_err("crawler worker exited")?;
let Some(next) = next_cursor else {
return Ok(PdsListing::Complete);
};
cursor = Some(next);
}
}
async fn fetch_page(
&self,
pds: &Url,
hostname: &str,
account_count: u64,
url: Url,
http: &PublicHttpClient,
) -> Result<Option<bytes::Bytes>> {
let throttle = self.checker.throttler.get_handle(pds).await;
if throttle.is_throttled() {
miette::bail!("PDS is throttled until {}", throttle.throttled_until());
}
let _permit = throttle.acquire().await;
let tier = self.checker.state.resolve_pds_tier(hostname);
throttle.wait_for_allow(account_count, &tier).await;
if throttle.is_throttled() {
miette::bail!("PDS became throttled while waiting for pacing");
}
let endpoint = PublicEndpoint::parse_http(&url)
.into_diagnostic()
.wrap_err("unsafe PDS endpoint for listRepos")?;
let response = match http.get(&endpoint).timeout(throttle.timeout()).send().await {
Ok(response) => response,
Err(err) => {
throttle.record_failure_detail("list_repos", err.to_string());
return Err(err).into_diagnostic();
}
};
if response.status() == StatusCode::TOO_MANY_REQUESTS {
throttle.record_ratelimit(parse_retry_after(&response));
miette::bail!("PDS rate limited listRepos");
}
if matches!(
response.status(),
StatusCode::BAD_REQUEST | StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED
) || response.status().is_redirection()
{
debug!(status = %response.status(), "PDS does not support listRepos; skipping");
throttle.record_success();
return Ok(None);
}
let response = match response.error_for_status() {
Ok(response) => response,
Err(err) => {
throttle.record_failure_detail("list_repos", err.to_string());
return Err(err).into_diagnostic();
}
};
let bytes = crate::util::read_limited_bytes(response, MAX_LIST_RESPONSE_BYTES).await?;
throttle.record_success();
Ok(Some(bytes))
}
async fn parse_page(
&self,
bytes: bytes::Bytes,
known: KnownRepos,
host: SmolStr,
) -> Result<Page> {
let db = self.checker.state.db.clone();
let skew_secs = self.checker.state.rev_clock_skew_secs;
tokio::time::timeout(
BLOCKING_TASK_TIMEOUT,
tokio::task::spawn_blocking(move || {
let output = serde_json::from_slice::<ListReposOutput>(&bytes)
.into_diagnostic()
.wrap_err("failed to parse PDS listRepos response")?;
select_unknown_repos(&db, output, known, &host, skew_secs)
}),
)
.await
.map_err(|_| miette::miette!("PDS listRepos parse timed out"))?
.into_diagnostic()?
}
/// costs each repo a refetch, unlike an operator's resync
async fn resync(&self, dids: Vec<Did>) -> Result<()> {
if dids.is_empty() {
return Ok(());
}
let answers = on_owners(
&self.checker.state,
dids,
|did| did,
|state, shard, batch| repo_ops(state, shard.as_deref(), RepoOp::ResyncCrawled, &batch),
)
.await;
let mut resynced = 0;
for answer in answers {
resynced += answer??.len();
}
if resynced > 0 {
debug!(resynced, "resyncing stuck repos a walk listed");
self.checker.state.notify_backfill();
}
Ok(())
}
pub(super) async fn send_cursor_updates(
&self,
cursor_updates: Vec<CursorUpdate>,
) -> Result<()> {
self.tx
.send(CrawlerBatch {
guards: Vec::new(),
listings: Vec::new(),
cursor_updates,
applied: None,
})
.await
.into_diagnostic()
.wrap_err("crawler worker exited")
}
}
/// drops the accounts in `accounts` past `hostname`'s tier account limit, the one the firehose
/// holds a host to as well, and says how many it dropped
fn keep_under_account_limit<T>(
state: &AppState,
hostname: &str,
admitted: u64,
accounts: &mut Vec<T>,
) -> usize {
let room = room_left(state, hostname, admitted, accounts.len());
let skipped = accounts.len() - room;
accounts.truncate(room);
skipped
}
/// how many of `wanted` more accounts `hostname` has room for. a crawled repo only counts for
/// its host once its backfill fills in its pds, so `admitted` adds the ones this walk already
/// sent
fn room_left(state: &AppState, hostname: &str, admitted: u64, wanted: usize) -> usize {
let count = state
.db
.get_count_sync(&keys::pds_account_count_key(hostname))
.saturating_add(admitted);
(0..wanted as u64)
.take_while(|n| !state.is_over_account_limit(hostname, count + n))
.count()
}
struct Page {
unknown_dids: Vec<Did>,
newer: Vec<RepoListing>,
stuck: Vec<Did>,
cursor: Option<SmolStr>,
count: usize,
}
fn select_unknown_repos(
db: &Db,
output: ListReposOutput,
known: KnownRepos,
host: &str,
skew_secs: i64,
) -> Result<Page> {
let count = output.repos.len();
if count > LIST_LIMIT {
miette::bail!("PDS listRepos returned {count} repos for a limit of {LIST_LIMIT}");
}
let cursor = output.cursor;
let mut unknown_dids = Vec::new();
let mut newer = Vec::new();
let mut stuck = Vec::new();
for repo in output.repos {
let exclude_key = crate::db::filter::exclude_key(repo.did.as_str())?;
if db.filter.contains_key(exclude_key).into_diagnostic()? {
continue;
}
if let Some(stored) = db.repos.get(keys::repo_key(&repo.did)).into_diagnostic()? {
if known == KnownRepos::Skip {
continue;
}
let stored = crate::db::deser_repo_state(&stored)?;
if stored.pds.as_deref().and_then(pds_host).as_deref() != Some(host) {
continue;
}
if let Some(listing) = listing_if_newer(&repo, &stored, skew_secs) {
newer.push(listing);
} else if known == KnownRepos::ResyncNewerOrStuck && is_stuck(db, &repo.did, &stored)? {
stuck.push(repo.did.into_static());
}
continue;
}
let retry_key = keys::crawler_retry_key(&repo.did);
if !db.crawler.contains_key(retry_key).into_diagnostic()? {
unknown_dids.push(repo.did);
}
}
Ok(Page {
unknown_dids,
newer,
stuck,
cursor,
count,
})
}
/// a tracked repo left desynchronized or with a failed backfill waiting on its retry
fn is_stuck(db: &Db, did: &Did, stored: &RepoState<'_>) -> Result<bool> {
let metadata = db
.repo_metadata
.get(keys::repo_metadata_key(did))
.into_diagnostic()?
.map(|bytes| crate::db::deser_repo_meta(&bytes))
.transpose()?;
// the retry worker requeues a deferred refetch when its window ends, a walk leaves it alone
if !metadata.is_some_and(|metadata| metadata.tracked && !metadata.refetch_deferred) {
return Ok(false);
}
if matches!(
stored.status,
RepoStatus::Desynchronized | RepoStatus::Error(_)
) {
return Ok(true);
}
let failed = db
.indexer
.resync
.get(keys::repo_key(did))
.into_diagnostic()?
.map(|bytes| rmp_serde::from_slice::<ResyncState>(&bytes).into_diagnostic())
.transpose()?;
Ok(matches!(failed, Some(ResyncState::Error { .. })))
}
/// `repo` as a listing if it's listed past the rev we hold
fn listing_if_newer(
repo: &ListedRepo<'_>,
stored: &RepoState<'_>,
skew_secs: i64,
) -> Option<RepoListing> {
let listing = RepoListing {
did: repo.did.clone().into_static(),
rev: DbTid::from(&repo.rev),
active: repo.active.unwrap_or(true),
status: repo.status.clone().into_static(),
create_if_missing: false,
};
listing.claims_newer(stored, skew_secs).then_some(listing)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
use serde_json::json;
use tempfile::tempdir;
const HEAD: &str = "bafyreidf747c4x3lps3k4n357l3a3r57k3k465743k573k465743k5";
fn list_repos_output(value: serde_json::Value) -> Result<ListReposOutput<'static>> {
let bytes = serde_json::to_vec(&value).into_diagnostic()?;
let out = serde_json::from_slice::<ListReposOutput>(&bytes).into_diagnostic()?;
Ok(ListReposOutput {
cursor: out.cursor,
repos: out
.repos
.into_iter()
.map(|r| crate::crawler::list_repos::producer::ListedRepo {
active: r.active,
did: r.did,
rev: r.rev,
status: r.status.into_static(),
})
.collect(),
})
}
#[test]
fn a_page_with_more_repos_than_the_limit_is_refused() -> Result<()> {
let tmp = tempdir().into_diagnostic()?;
let state = AppState::new(&Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})?;
for (count, fits) in [(LIST_LIMIT, true), (LIST_LIMIT + 1, false)] {
let repos: Vec<_> = (0..count)
.map(|i| json!({"did": format!("did:web:r{i}.example"), "head": HEAD, "rev": "3j6s746574657"}))
.collect();
let page = select_unknown_repos(
&state.db,
list_repos_output(json!({ "repos": repos }))?,
KnownRepos::Skip,
"pds.example",
state.rev_clock_skew_secs,
);
assert_eq!(page.is_ok(), fits, "{count} repos");
}
Ok(())
}
#[test]
fn a_listed_rev_from_the_future_is_not_newer() -> Result<()> {
let tmp = tempdir().into_diagnostic()?;
let state = AppState::new(&Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})?;
let did = Did::new_static("did:web:alice.example").into_diagnostic()?;
let mut row = RepoState::synced();
row.pds = Some(jacquard_common::CowStr::Borrowed("https://pds.example/"));
let mut txn = crate::db::Txn::new(&state.db);
txn.put_repo(&did, &row)?;
txn.commit()?;
for (rev, listed) in [("3mz2j7sg2kh22", 1), ("jzzzzzzzzzzzz", 0)] {
let page = select_unknown_repos(
&state.db,
list_repos_output(json!({
"repos": [{"did": did.as_str(), "head": HEAD, "rev": rev}]
}))?,
KnownRepos::ResyncNewer,
"pds.example",
state.rev_clock_skew_secs,
)?;
assert_eq!(page.newer.len(), listed, "{rev}");
}
Ok(())
}
#[test]
fn a_repo_waiting_on_its_refetch_allowance_is_not_stuck() -> Result<()> {
use crate::ingest::indexer::refetch::tests as refetch;
let (_tmp, state) = refetch::state(Some(1));
let did = refetch::did();
let mut metadata = refetch::seed(&state, true, None);
let entry = ResyncState::Error {
kind: crate::types::ResyncErrorKind::Generic,
retry_count: 1,
next_retry: 0,
};
refetch::put_entry(&state, &did, &entry);
assert!(is_stuck(&state.db, &did, &refetch::row(&state))?);
metadata.refetch_deferred = true;
refetch::put_metadata(&state, &did, &metadata);
assert!(!is_stuck(&state.db, &did, &refetch::row(&state))?);
Ok(())
}
#[test]
fn pages_select_only_unknown_repos_without_trusting_status() -> Result<()> {
let tmp = tempdir().into_diagnostic()?;
let config = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = AppState::new(&config)?;
let known: Did = Did::new_static("did:web:known.example")?;
let excluded: Did = Did::new_static("did:web:excluded.example")?;
let retry: Did = Did::new_static("did:web:retry.example")?;
state
.db
.repos
.insert(keys::repo_key(&known), b"relay")
.into_diagnostic()?;
state
.db
.filter
.insert(crate::db::filter::exclude_key(excluded.as_str())?, b"")
.into_diagnostic()?;
state
.db
.crawler
.insert(keys::crawler_retry_key(&retry), b"")
.into_diagnostic()?;
let page = select_unknown_repos(
&state.db,
list_repos_output(json!({
"cursor": "next",
"repos": [
{"did": known, "head": HEAD, "rev": "3j6s746574657", "active": false, "status": "takendown"},
{"did": excluded, "head": HEAD, "rev": "3j6s746574657"},
{"did": retry, "head": HEAD, "rev": "3j6s746574657"},
{"did": "did:web:inactive-unknown.example", "head": HEAD, "rev": "3j6s746574657", "active": false, "status": "deactivated"}
]
}))?,
KnownRepos::Skip,
"pds.example",
state.rev_clock_skew_secs,
)?;
assert_eq!(page.count, 4);
assert_eq!(page.cursor.as_deref(), Some("next"));
assert_eq!(page.unknown_dids.len(), 1);
assert_eq!(
page.unknown_dids[0].as_str(),
"did:web:inactive-unknown.example"
);
assert_eq!(
state
.db
.repos
.get(keys::repo_key(&known))
.into_diagnostic()?
.as_deref(),
Some(b"relay".as_slice())
);
Ok(())
}
#[test]
fn a_host_at_its_account_limit_gets_no_more_crawled_accounts() -> Result<()> {
let tmp = tempdir().into_diagnostic()?;
let config = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = AppState::new(&config)?;
let limit = state
.resolve_pds_tier("pds.example")
.account_limit
.expect("the host's tier has an account limit");
state.db.update_count(
&keys::pds_account_count_key("pds.example"),
limit as i64 - 5,
);
let dids = |n: usize| -> Vec<Did> {
(0..n)
.map(|i| Did::new_owned(format!("did:web:repo{i}.example")).unwrap())
.collect()
};
let mut page = dids(8);
assert_eq!(
keep_under_account_limit(&state, "pds.example", 0, &mut page),
3
);
assert_eq!(page, dids(5));
// what the walk already sent counts too, since the host's count hasn't caught up yet
let mut page = dids(8);
assert_eq!(
keep_under_account_limit(&state, "pds.example", 5, &mut page),
8
);
assert!(page.is_empty());
// another host's limit is its own
let mut page = dids(8);
assert_eq!(
keep_under_account_limit(&state, "other.example", 5, &mut page),
0
);
Ok(())
}
}