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, ) { 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) { 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, state: &Arc, 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 = 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, 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 { 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> { 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>, fail_second_page: std::sync::atomic::AtomicBool, } #[cfg(feature = "indexer")] async fn mock_relay() -> miette::Result<(Url, Arc)> { 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>, Query(query): Query>, ) -> 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 { 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 { 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()) }