use std::sync::Arc;
use std::time::Duration;
use futures::StreamExt;
use miette::{Context, IntoDiagnostic};
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::{ListHostsOutput, ListedHost as Host, canonical_pds_host, list_hosts_url};
use crate::state::AppState;
const MAX_CONCURRENT_SEEDS: usize = 4;
const LIST_HOSTS_BODY_PREVIEW_BYTES: usize = 512;
/// seed firehose pds sources by calling `com.atproto.sync.listHosts` on each seed URL.
/// banned pds' are not added, everything else is (including offline)
pub(crate) async fn seed_from_list_hosts(
seed_urls: &[Url],
firehose: &FirehoseHandle,
state: &Arc<AppState>,
) {
info!("will seed urls...");
let http = seed_http_client();
let mut futs = futures::stream::iter(seed_urls.iter().cloned())
.map(|seed_url| {
let firehose = Some(firehose.clone());
let state = state.clone();
let http = http.clone();
async move { seed_one(&seed_url, firehose, &state, &http).await }
})
.buffer_unordered(MAX_CONCURRENT_SEEDS);
while futs.next().await.is_some() {}
}
/// refresh seed relay host status snapshots without adding sources.
///
/// this runs before persisted sources are spawned so existing PDS tasks start
/// with updated host status metadata while preserving local firehose cursors.
/// hosts without a cursor get the seed relay's seq for them.
pub(crate) async fn refresh_seed_snapshots(seed_urls: &[Url], state: &Arc<AppState>) {
info!("will refresh seed snapshots...");
let http = seed_http_client();
let mut futs = futures::stream::iter(seed_urls.iter().cloned())
.map(|seed_url| {
let state = state.clone();
let http = http.clone();
async move { seed_one(&seed_url, None, &state, &http).await }
})
.buffer_unordered(MAX_CONCURRENT_SEEDS);
while futs.next().await.is_some() {}
}
fn seed_http_client() -> reqwest::Client {
crate::net::trusted_http().clone()
}
#[tracing::instrument(skip_all, fields(seed_url = %seed_url))]
async fn seed_one(
seed_url: &Url,
firehose: Option<FirehoseHandle>,
state: &Arc<AppState>,
http: &reqwest::Client,
) {
// always start from the beginning. `listHosts` is small, and resuming from
// an old completed-run cursor hides earlier hosts after a restart.
let mut cursor: Option<String> = 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;
let mut page_count = 0;
loop {
page_count += 1;
if page_count > 100 {
warn!("listHosts pagination limit (100) reached, stopping");
break;
}
let url = list_hosts_url(seed_url, cursor.as_deref());
let resp = match http
.get(url.clone())
.timeout(Duration::from_secs(10))
.send()
.await
{
Ok(r) => r,
Err(e) => {
warn!(url = %url, err = %e, "failed to fetch listHosts, stopping");
break;
}
};
let status = resp.status();
let bytes = match read_limited_response(resp, 10 * 1024 * 1024).await {
Ok(b) => b,
Err(e) => {
warn!(url = %url, err = %e, "failed to read listHosts response, stopping");
break;
}
};
if !status.is_success() {
let body = body_preview(&bytes);
warn!(
url = %url,
status = %status,
body = %body,
"listHosts returned error status, stopping"
);
break;
}
let body: ListHostsOutput = match serde_json::from_slice(&bytes) {
Ok(b) => b,
Err(e) => {
warn!(
url = %url,
err = %e,
body = %body_preview(&bytes),
"failed to parse listHosts response, stopping"
);
break;
}
};
let next_cursor = body.cursor.as_deref().map(str::to_owned);
total += body.hosts.len();
let page_hosts = body.hosts.len();
if let Err(e) = apply_seed_snapshot(state, &body.hosts).await {
warn!(err = %e, "failed to apply listHosts seed snapshot");
}
let mut banned = 0usize;
let mut malformed = 0usize;
let mut already_known = 0usize;
let mut rejected = 0usize;
let mut queued = 0usize;
let mut page_added = 0usize;
if let Some(firehose) = firehose.as_ref() {
let mut seed_sources = Vec::with_capacity(body.hosts.len());
for host in &body.hosts {
// A relay-declared ban is still a terminal seed status.
if host.status.as_deref() == Some("banned") {
banned += 1;
continue;
}
let wss_url = match Url::parse(&format!("wss://{}/", host.hostname)) {
Ok(url) => url,
Err(error) => {
malformed += 1;
warn!(hostname = %host.hostname, %error, "invalid hostname in listHosts response, skipping");
continue;
}
};
queued += 1;
seed_sources.push(wss_url);
}
match firehose.add_seeded_sources(seed_sources, queue_crawl).await {
Ok(adds) => {
page_added = adds.added;
added += adds.added;
already_known = adds.known;
rejected = adds.rejected;
}
Err(error) => {
admitted_every_page = false;
warn!(%error, "failed to add seeded firehose sources");
}
}
}
info!(
page_hosts,
banned,
malformed,
rejected,
already_known,
queued,
added = page_added,
adding_sources = firehose.is_some(),
next_cursor = ?next_cursor,
"processed listHosts page"
);
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"
);
}
fn body_preview(bytes: &[u8]) -> String {
let len = bytes.len().min(LIST_HOSTS_BODY_PREVIEW_BYTES);
let mut body = String::from_utf8_lossy(&bytes[..len]).into_owned();
if bytes.len() > len {
body.push_str("...");
}
body
}
async fn apply_seed_snapshot(state: &Arc<AppState>, hosts: &[Host]) -> miette::Result<()> {
let _admission = state.pds_admission.lock().await;
let _status_write = state.pds_status_write.lock();
let mut batch = state.db.inner.batch();
let mut status_updates = Vec::with_capacity(hosts.len());
let mut cursor_seeds = 0usize;
for host in hosts {
let hostname = match canonical_pds_host(&host.hostname) {
Ok(hostname) => hostname,
Err(error) => {
debug!(hostname = %host.hostname, %error, "invalid hostname in listHosts snapshot, skipping");
continue;
}
};
let status = map_seed_status(host.status.as_deref());
crate::db::pds_meta::set_status(&mut batch, &state.db.filter, hostname.as_str(), status)?;
status_updates.push((hostname.as_str().to_string(), status));
if seed_absent_cursor(state, &mut batch, host)? {
cursor_seeds += 1;
}
}
batch.commit().into_diagnostic()?;
state.db.persist()?;
debug!(
hosts = hosts.len(),
status_updates = status_updates.len(),
cursor_seeds,
"applied listHosts seed snapshot"
);
state.pds_meta.rcu(|meta| {
let mut next = (**meta).clone();
for (hostname, status) in &status_updates {
next.update_host_entry(hostname, |entry| entry.status = *status);
}
next
});
Ok(())
}
/// stage the seed relay's seq as the cursor of a host we have no cursor for.
///
/// at startup this runs before the crawler, so a new host's stream starts from
/// before any of its repos are backfilled instead of from its jittered first
/// connect. a cursor we already hold, including the `-1` live tail marker,
/// always wins, and so does a running ingestor's in-memory one.
fn seed_absent_cursor(
state: &AppState,
batch: &mut fjall::OwnedWriteBatch,
host: &Host,
) -> miette::Result<bool> {
let Some(seq) = host
.seq
.and_then(|seq| i64::try_from(seq).ok())
.filter(|seq| *seq > 0)
else {
return Ok(false);
};
let Ok(url) = Url::parse(&format!("wss://{}/", host.hostname)) else {
return Ok(false);
};
if state.firehose_cursors.contains(&url) {
return Ok(false);
}
let key = keys::firehose_cursor_key_from_url(&url);
if state.db.cursors.contains_key(&key).into_diagnostic()? {
return Ok(false);
}
batch.insert(&state.db.cursors, key, seq.to_be_bytes());
Ok(true)
}
fn map_seed_status(status: Option<&str>) -> crate::pds_meta::HostStatus {
match status {
Some("active") | None => crate::pds_meta::HostStatus::Active,
Some("idle") => crate::pds_meta::HostStatus::Idle,
Some("offline") => crate::pds_meta::HostStatus::Offline,
Some("throttled") => crate::pds_meta::HostStatus::Throttled,
Some("banned") => crate::pds_meta::HostStatus::Banned,
Some(_) => crate::pds_meta::HostStatus::Active,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
use crate::control::Hydrant;
use jacquard_common::CowStr;
use std::sync::atomic::{AtomicI64, Ordering};
use tempfile::tempdir;
fn persisted_cursor(state: &AppState, hostname: &str) -> miette::Result<Option<i64>> {
state
.db
.cursors
.get(keys::firehose_cursor_key(hostname))
.into_diagnostic()?
.map(|bytes| {
bytes
.as_ref()
.try_into()
.into_diagnostic()
.map(i64::from_be_bytes)
})
.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<Vec<&'static str>>,
fail_second_page: std::sync::atomic::AtomicBool,
}
#[cfg(feature = "indexer")]
async fn mock_relay() -> miette::Result<(Url, Arc<MockRelay>)> {
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<Arc<MockRelay>>,
Query(query): Query<HashMap<String, String>>,
) -> 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<bool> {
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()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&cfg)?);
let hosts = vec![
Host {
hostname: "active.example".into(),
account_count: Some(42),
seq: Some(100),
status: Some("active".into()),
},
Host {
hostname: "offline.example".into(),
account_count: Some(7),
seq: Some(5),
status: Some("offline".into()),
},
];
apply_seed_snapshot(&state, &hosts).await?;
assert_eq!(
state
.db
.get_count_sync(&keys::pds_account_count_key("active.example")),
0
);
assert_eq!(
state
.db
.get_count_sync(&keys::pds_account_count_key("offline.example")),
0
);
let meta = state.pds_meta.load();
assert_eq!(
meta.status("active.example"),
crate::pds_meta::HostStatus::Active
);
assert_eq!(
meta.status("offline.example"),
crate::pds_meta::HostStatus::Offline
);
assert!(meta.hosts.contains_key("active.example"));
assert!(meta.hosts.contains_key("offline.example"));
assert_eq!(persisted_cursor(&state, "active.example")?, Some(100));
assert_eq!(persisted_cursor(&state, "offline.example")?, Some(5));
let active_url = Url::parse("wss://active.example/").into_diagnostic()?;
let in_memory = state
.firehose_cursors
.peek_with(&active_url, |_, cursor| cursor.load(Ordering::SeqCst));
assert_eq!(in_memory, None);
Ok(())
}
#[tokio::test]
async fn apply_seed_snapshot_preserves_existing_cursor() -> miette::Result<()> {
let tmp = tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&cfg)?);
let url = Url::parse("wss://active.example/").into_diagnostic()?;
crate::db::set_firehose_cursor(&state.db, &url, 150)?;
let _ = state
.firehose_cursors
.insert_sync(url.clone(), AtomicI64::new(150));
let hosts = vec![Host {
hostname: "active.example".into(),
account_count: Some(42),
seq: Some(500),
status: Some("active".into()),
}];
apply_seed_snapshot(&state, &hosts).await?;
assert_eq!(persisted_cursor(&state, "active.example")?, Some(150));
let in_memory = state
.firehose_cursors
.peek_with(&url, |_, cursor| cursor.load(Ordering::SeqCst));
assert_eq!(in_memory, Some(150));
Ok(())
}
#[tokio::test]
async fn apply_seed_snapshot_keeps_live_tail_marker() -> miette::Result<()> {
let tmp = tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&cfg)?);
let url = Url::parse("wss://active.example/").into_diagnostic()?;
crate::db::set_firehose_cursor(&state.db, &url, -1)?;
let hosts = vec![Host {
hostname: "active.example".into(),
account_count: Some(42),
seq: Some(500),
status: Some("active".into()),
}];
apply_seed_snapshot(&state, &hosts).await?;
assert_eq!(persisted_cursor(&state, "active.example")?, Some(-1));
Ok(())
}
#[tokio::test]
async fn apply_seed_snapshot_leaves_running_ingestor_cursor() -> miette::Result<()> {
let tmp = tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&cfg)?);
let url = Url::parse("wss://active.example/").into_diagnostic()?;
// a live tailing ingestor that has not checkpointed an event yet
let _ = state
.firehose_cursors
.insert_sync(url.clone(), AtomicI64::new(0));
let hosts = vec![Host {
hostname: "active.example".into(),
account_count: Some(42),
seq: Some(500),
status: Some("active".into()),
}];
apply_seed_snapshot(&state, &hosts).await?;
assert_eq!(persisted_cursor(&state, "active.example")?, None);
Ok(())
}
#[tokio::test]
async fn relay_banned_snapshot_blocks_admission_until_active() -> miette::Result<()> {
let tmp = tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let hydrant = Hydrant::new(cfg).await?;
hydrant.firehose.startable_for_test();
let banned = Host {
hostname: "relay-banned.example".into(),
account_count: None,
seq: None,
status: Some("banned".into()),
};
apply_seed_snapshot(&hydrant.state, &[banned]).await?;
assert_eq!(
hydrant.state.pds_meta.load().status("relay-banned.example"),
crate::pds_meta::HostStatus::Banned
);
assert!(hydrant.pds.is_banned("relay-banned.example"));
let active = Host {
hostname: "relay-banned.example".into(),
account_count: None,
seq: None,
status: Some("active".into()),
};
apply_seed_snapshot(&hydrant.state, &[active]).await?;
assert_eq!(
hydrant.state.pds_meta.load().status("relay-banned.example"),
crate::pds_meta::HostStatus::Active
);
assert!(!hydrant.pds.is_banned("relay-banned.example"));
assert_eq!(
hydrant
.firehose
.add_source_validated(Url::parse("wss://relay-banned.example/").into_diagnostic()?)
.await?,
Ok(crate::control::firehose::PdsAdmission::StartedNew)
);
hydrant
.firehose
.tasks
.remove_async(&Url::parse("wss://relay-banned.example/").into_diagnostic()?)
.await;
Ok(())
}
#[tokio::test]
async fn seed_snapshot_cannot_clear_a_committed_ban_after_reopen() -> miette::Result<()> {
let tmp = tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let hydrant = Hydrant::new(cfg.clone()).await?;
let hosts = vec![Host {
hostname: "BLOCKED.EXAMPLE.".into(),
account_count: None,
seq: None,
status: Some("active".into()),
}];
// the listHosts page was fetched before the local ban committed.
hydrant.pds.ban("blocked.example").await?;
apply_seed_snapshot(&hydrant.state, &hosts).await?;
assert_eq!(
hydrant.state.pds_meta.load().status("blocked.example"),
crate::pds_meta::HostStatus::Active
);
assert!(hydrant.pds.is_banned("blocked.example"));
drop(hydrant);
let reopened = AppState::new(&cfg)?;
assert_eq!(
reopened.pds_meta.load().status("blocked.example"),
crate::pds_meta::HostStatus::Active
);
assert!(reopened.pds_meta.load().is_banned("blocked.example"));
Ok(())
}
}
async fn read_limited_response(
resp: reqwest::Response,
limit: usize,
) -> miette::Result<bytes::Bytes> {
use bytes::BytesMut;
use futures::StreamExt as _;
let mut stream = resp.bytes_stream();
let mut buf = BytesMut::new();
while let Some(chunk) = stream.next().await {
let chunk = chunk
.into_diagnostic()
.context("failed to read response chunk")?;
if buf.len() + chunk.len() > limit {
miette::bail!("response body too large (exceeds {limit} bytes)");
}
buf.extend_from_slice(&chunk);
}
Ok(buf.freeze())
}