use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use crate::pds_discovery::{ListHostsOutput, ListedHost as Host};
use jacquard_common::types::string::Did;
use miette::{Context, IntoDiagnostic, Result};
use serde::{Deserialize, Serialize};
use smol_str::SmolStr;
use tokio::sync::{mpsc, watch};
use tracing::{debug, error, info, trace, warn};
use url::Url;
use super::list_repos::ListReposProducer;
use super::pds_repos::{KnownRepos, PdsRepoLister};
use super::worker::{CrawlerBatch, CursorUpdate};
use super::{CrawlerStats, InFlight, SignalChecker, base_url};
use crate::db::keys;
use crate::net::{PublicHost, PublicHttpClient, public_http, trusted_http};
use crate::pds_discovery;
use crate::state::AppState;
use crate::util::WatchEnabledExt;
const MAX_CONCURRENT_PDS_CRAWLS: usize = 4;
const MAX_LIST_RESPONSE_BYTES: usize = 10 * 1024 * 1024;
const PASS_POLL_INTERVAL: Duration = Duration::from_secs(1);
const RETRY_DELAY: Duration = Duration::from_secs(30);
const MAX_LIST_HOSTS_PAGES: usize = 100;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
struct PassCursor {
pass: u64,
cursor: Option<SmolStr>,
complete: bool,
}
impl PassCursor {
fn active(pass: u64, cursor: Option<SmolStr>) -> Self {
Self {
pass,
cursor,
complete: false,
}
}
fn complete(pass: u64) -> Self {
Self {
pass,
cursor: None,
complete: true,
}
}
}
pub(crate) struct RelayFirstProducer {
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,
}
impl RelayFirstProducer {
pub(crate) async fn run(self) -> Result<()> {
let pass_key = keys::relay_first_pass_key(self.url.as_str());
let relay = ListReposProducer {
url: self.url.clone(),
checker: self.checker.clone(),
in_flight: self.in_flight.clone(),
tx: self.tx.clone(),
enabled: self.enabled.clone(),
stats: self.stats.clone(),
completion_key: Some(pass_key),
last_completion: None,
};
let deep = DeepDiscoveryProducer {
relay_url: self.url,
lister: PdsRepoLister {
checker: self.checker,
in_flight: self.in_flight,
tx: self.tx,
enabled: self.enabled,
stats: self.stats,
component: "relay-first deep crawler",
},
};
tokio::select! {
result = relay.run() => result,
result = deep.run() => result,
}
}
}
#[derive(Clone)]
struct DeepDiscoveryProducer {
relay_url: Url,
lister: PdsRepoLister,
}
impl DeepDiscoveryProducer {
async fn run(mut self) -> Result<()> {
loop {
self.lister
.enabled
.wait_enabled(self.lister.component)
.await;
let pass = match self.wait_for_unhandled_pass().await {
Ok(pass) => pass,
Err(err) => {
error!(err = ?err, "failed to read deep discovery pass state, retrying");
tokio::time::sleep(RETRY_DELAY).await;
continue;
}
};
info!(pass, "starting relay-first deep discovery pass");
if let Err(err) = self.crawl_pass(pass).await {
error!(pass, err = ?err, "deep discovery pass failed, retrying");
tokio::time::sleep(RETRY_DELAY).await;
continue;
}
info!(pass, "relay-first deep discovery pass complete");
}
}
async fn wait_for_unhandled_pass(&mut self) -> Result<u64> {
let pass_key = keys::relay_first_pass_key(self.relay_url.as_str());
let handled_key = keys::relay_first_handled_key(self.relay_url.as_str());
loop {
self.lister
.enabled
.wait_enabled(self.lister.component)
.await;
let state = self.lister.checker.state.clone();
let pass_key = pass_key.clone();
let handled_key = handled_key.clone();
let (pass, handled) = state
.db
.run(move |db| {
Ok((
read_u64(&db.cursors, &pass_key)?,
read_u64(&db.cursors, &handled_key)?,
))
})
.await?;
if let Some(pass) = pass
&& handled.is_none_or(|handled| pass > handled)
{
return Ok(pass);
}
tokio::select! {
_ = tokio::time::sleep(PASS_POLL_INTERVAL) => {}
changed = self.lister.enabled.changed() => {
if changed.is_err() {
miette::bail!("crawler enable channel closed");
}
}
}
}
}
async fn crawl_pass(&mut self, pass: u64) -> Result<()> {
let cursor_key = keys::relay_first_hosts_cursor_key(self.relay_url.as_str());
let mut cursor = self
.load_cursor(&cursor_key, pass)
.await?
.and_then(|state| (!state.complete).then_some(state.cursor).flatten());
let mut page_count = 0usize;
loop {
page_count += 1;
if page_count > MAX_LIST_HOSTS_PAGES {
miette::bail!("listHosts exceeded {MAX_LIST_HOSTS_PAGES} pages");
}
self.lister
.enabled
.wait_enabled(self.lister.component)
.await;
let url = pds_discovery::list_hosts_url(&base_url(&self.relay_url)?, cursor.as_deref());
let response = trusted_http()
.get(url.clone())
.send()
.await
.into_diagnostic()
.wrap_err_with(|| format!("failed to fetch {url}"))?
.error_for_status()
.into_diagnostic()
.wrap_err_with(|| format!("listHosts returned an error for {url}"))?;
let bytes = crate::util::read_limited_bytes(response, MAX_LIST_RESPONSE_BYTES).await?;
let output = serde_json::from_slice::<ListHostsOutput>(&bytes)
.into_diagnostic()
.wrap_err("failed to parse listHosts response")?;
let next_cursor = output.cursor;
let hosts: Vec<Host> = output.hosts;
let mut crawls = tokio::task::JoinSet::new();
for host in hosts {
if crawls.len() >= MAX_CONCURRENT_PDS_CRAWLS
&& let Some(result) = crawls.join_next().await
{
log_pds_crawl_result(result);
}
let producer = self.clone();
crawls.spawn(async move {
let hostname = host.hostname.to_string();
(hostname, producer.crawl_host(pass, host).await)
});
}
while let Some(result) = crawls.join_next().await {
log_pds_crawl_result(result);
}
if next_cursor.is_some() && next_cursor == cursor {
miette::bail!("listHosts returned the same cursor twice");
}
let cursor_updates = if let Some(next) = &next_cursor {
vec![CursorUpdate {
key: cursor_key.clone(),
value: Some(serialize_cursor(&PassCursor::active(
pass,
Some(next.clone()),
))?),
}]
} else {
vec![
CursorUpdate {
key: keys::relay_first_handled_key(self.relay_url.as_str()),
value: Some(pass.to_be_bytes().to_vec()),
},
CursorUpdate {
key: cursor_key.clone(),
value: None,
},
]
};
self.lister.send_cursor_updates(cursor_updates).await?;
let Some(next) = next_cursor else {
return Ok(());
};
cursor = Some(next);
}
}
async fn crawl_host(mut self, pass: u64, host: Host) -> Result<()> {
let hostname = host.hostname.as_str();
if matches!(
host.status.as_deref(),
Some("banned" | "offline" | "throttled")
) {
debug!(%hostname, status = ?host.status, "skipping unavailable deep-discovery PDS");
return Ok(());
}
let public_host = match PublicHost::parse(hostname) {
Ok(host) => host,
Err(error) => {
warn!(%hostname, %error, "unsafe/private hostname in listHosts response, skipping deep discovery");
return Ok(());
}
};
let url_host = public_host.as_str();
let pds = public_host.https_endpoint().url().clone();
let http = public_http();
if self
.lister
.checker
.state
.pds_meta
.load()
.is_banned(url_host)
{
debug!(%hostname, "locally banned PDS, skipping deep discovery");
return Ok(());
}
let account_count = host
.account_count
.and_then(|count| u64::try_from(count).ok())
.unwrap_or(1)
.max(1);
self.crawl_pds(pass, &pds, account_count, &http).await
}
async fn crawl_pds(
&mut self,
pass: u64,
pds: &Url,
account_count: u64,
http: &PublicHttpClient,
) -> Result<()> {
let hostname = pds
.host_str()
.ok_or_else(|| miette::miette!("PDS URL has no hostname"))?;
let cursor_key = keys::relay_first_pds_cursor_key(self.relay_url.as_str(), hostname);
let saved = self.load_cursor(&cursor_key, pass).await?;
if saved.as_ref().is_some_and(|cursor| cursor.complete) {
trace!(%hostname, pass, "deep PDS crawl already complete");
return Ok(());
}
let start = saved.and_then(|state| state.cursor);
self.lister
.list(
pds,
account_count,
http,
start,
KnownRepos::Skip,
None,
|next| {
let state = next.map_or_else(
|| PassCursor::complete(pass),
|next| PassCursor::active(pass, Some(next.clone())),
);
Ok(vec![CursorUpdate {
key: cursor_key.clone(),
value: Some(serialize_cursor(&state)?),
}])
},
)
.await?;
Ok(())
}
async fn load_cursor(&self, key: &[u8], pass: u64) -> Result<Option<PassCursor>> {
let key = key.to_vec();
self.lister
.checker
.state
.db
.run(move |db| {
db.cursors
.get(key)
.into_diagnostic()?
.map(|value| rmp_serde::from_slice::<PassCursor>(&value).into_diagnostic())
.transpose()
.map(|cursor| cursor.filter(|cursor| cursor.pass == pass))
})
.await
}
}
fn log_pds_crawl_result(result: std::result::Result<(String, Result<()>), tokio::task::JoinError>) {
match result {
Ok((hostname, Err(err))) => {
warn!(%hostname, err = ?err, "deep PDS crawl failed; retrying on the next relay pass");
}
Err(err) => error!(err = ?err, "deep PDS crawl task failed"),
Ok((_, Ok(()))) => {}
}
}
fn serialize_cursor(cursor: &PassCursor) -> Result<Vec<u8>> {
rmp_serde::to_vec(cursor).into_diagnostic()
}
fn read_u64(keyspace: &fjall::Keyspace, key: &[u8]) -> Result<Option<u64>> {
keyspace
.get(key)
.into_diagnostic()?
.map(|value| {
value
.as_ref()
.try_into()
.into_diagnostic()
.map(u64::from_be_bytes)
})
.transpose()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
use crate::crawler::{CrawlerWorker, InFlight};
use axum::{Json, Router, extract::Query, routing::get};
use serde_json::json;
use std::sync::atomic::{AtomicUsize, Ordering};
use tempfile::tempdir;
const HEAD: &str = "bafyreidf747c4x3lps3k4n357l3a3r57k3k465743k573k465743k5";
#[test]
fn pass_cursor_round_trips() -> Result<()> {
let cursor = PassCursor::active(7, Some("stinkpot".into()));
let encoded = serialize_cursor(&cursor)?;
assert_eq!(
crate::util::hex::encode(&encoded),
"9307a87374696e6b706f74c2"
);
assert_eq!(
rmp_serde::from_slice::<PassCursor>(&encoded).into_diagnostic()?,
cursor
);
let complete = PassCursor::complete(7);
let encoded = serialize_cursor(&complete)?;
assert_eq!(crate::util::hex::encode(&encoded), "9307c0c3");
assert_eq!(
rmp_serde::from_slice::<PassCursor>(&encoded).into_diagnostic()?,
complete
);
Ok(())
}
#[test]
fn pass_markers_are_read_big_endian() -> Result<()> {
let tmp = tempdir().into_diagnostic()?;
let db = crate::db::Db::open(&Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})?;
let key = keys::relay_first_pass_key("https://relay.example/");
assert_eq!(read_u64(&db.cursors, &key)?, None);
db.cursors
.insert(&key, [0, 0, 0, 0, 0, 0, 1, 2])
.into_diagnostic()?;
assert_eq!(read_u64(&db.cursors, &key)?, Some(258));
db.cursors.insert(&key, [0; 7]).into_diagnostic()?;
assert!(read_u64(&db.cursors, &key).is_err());
Ok(())
}
#[tokio::test]
async fn pds_pages_resume_and_commit_only_after_terminal_page() -> Result<()> {
let requests = Arc::new(AtomicUsize::new(0));
let app = Router::new().route(
"/xrpc/com.atproto.sync.listRepos",
get({
let requests = requests.clone();
move |Query(query): Query<HashMap<String, String>>| {
let requests = requests.clone();
async move {
requests.fetch_add(1, Ordering::Relaxed);
let body = if query.get("cursor").is_some_and(|value| value == "next") {
json!({
"repos": [
{"did": "did:web:second-unknown.example", "head": HEAD, "rev": "3j6s746574657"}
]
})
} else {
json!({
"cursor": "next",
"repos": [
{"did": "did:web:known.example", "head": HEAD, "rev": "3j6s746574657", "active": false, "status": "takendown"},
{"did": "did:web:first-unknown.example", "head": HEAD, "rev": "3j6s746574657", "active": false, "status": "deactivated"}
]
})
};
Json(body)
}
}
}),
);
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 known: Did = Did::new_static("did:web:known.example")?;
state
.db
.repos
.insert(keys::repo_key(&known), b"relay-authority")
.into_diagnostic()?;
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());
// Keep the URL public while pinning the test client to the loopback
// listener. This exercises the endpoint wrapper without live DNS.
let pds = Url::parse("http://pds.example/").into_diagnostic()?;
let mut producer = DeepDiscoveryProducer {
relay_url: Url::parse("https://relay.example").into_diagnostic()?,
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: "relay-first deep crawler",
},
};
let http = PublicHttpClient::test(
reqwest::Client::builder()
.resolve("pds.example", address)
.build()
.into_diagnostic()?,
);
producer.crawl_pds(7, &pds, 1, &http).await?;
let cursor_key =
keys::relay_first_pds_cursor_key(producer.relay_url.as_str(), pds.host_str().unwrap());
tokio::time::timeout(Duration::from_secs(5), async {
loop {
let complete = state
.db
.cursors
.get(&cursor_key)
.ok()
.flatten()
.and_then(|value| rmp_serde::from_slice::<PassCursor>(&value).ok())
.is_some_and(|cursor| cursor == PassCursor::complete(7));
if complete {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.into_diagnostic()?;
assert_eq!(requests.load(Ordering::Relaxed), 2);
let first: Did = Did::new_static("did:web:first-unknown.example")?;
let second: Did = Did::new_static("did:web:second-unknown.example")?;
assert!(
state
.db
.repos
.contains_key(keys::repo_key(&first))
.into_diagnostic()?
);
assert!(
state
.db
.repos
.contains_key(keys::repo_key(&second))
.into_diagnostic()?
);
assert_eq!(
state
.db
.repos
.get(keys::repo_key(&known))
.into_diagnostic()?
.as_deref(),
Some(b"relay-authority".as_slice())
);
producer.crawl_pds(7, &pds, 1, &http).await?;
assert_eq!(requests.load(Ordering::Relaxed), 2);
state
.db
.cursors
.insert(
&cursor_key,
serialize_cursor(&PassCursor::active(8, Some("next".into())))?,
)
.into_diagnostic()?;
producer.crawl_pds(8, &pds, 1, &http).await?;
assert_eq!(requests.load(Ordering::Relaxed), 3);
state.pds_meta.rcu(|meta| {
let mut next = (**meta).clone();
next.update_host_entry(pds.host_str().unwrap(), |entry| {
entry.status = crate::pds_meta::HostStatus::Banned;
});
next
});
producer.crawl_pds(9, &pds, 1, &http).await?;
assert_eq!(requests.load(Ordering::Relaxed), 3);
worker.abort();
server.abort();
Ok(())
}
#[test]
fn private_pds_addresses_are_rejected() {
assert!(PublicHost::parse("127.0.0.1").is_err());
assert!(PublicHost::parse("::1").is_err());
}
#[tokio::test]
async fn deep_discovery_waits_for_a_new_relay_pass_marker() -> Result<()> {
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 (tx, _rx) = mpsc::channel(1);
let mut producer = DeepDiscoveryProducer {
relay_url: Url::parse("https://relay.example").into_diagnostic()?,
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: "relay-first deep crawler",
},
};
assert!(
tokio::time::timeout(
Duration::from_millis(20),
producer.wait_for_unhandled_pass()
)
.await
.is_err()
);
state
.db
.cursors
.insert(
keys::relay_first_pass_key(producer.relay_url.as_str()),
7u64.to_be_bytes(),
)
.into_diagnostic()?;
state
.db
.cursors
.insert(
keys::relay_first_handled_key(producer.relay_url.as_str()),
7u64.to_be_bytes(),
)
.into_diagnostic()?;
assert!(
tokio::time::timeout(
Duration::from_millis(20),
producer.wait_for_unhandled_pass()
)
.await
.is_err()
);
state
.db
.cursors
.insert(
keys::relay_first_pass_key(producer.relay_url.as_str()),
8u64.to_be_bytes(),
)
.into_diagnostic()?;
assert_eq!(producer.wait_for_unhandled_pass().await?, 8);
Ok(())
}
}