Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499use std::sync::Arc;use std::time::Duration;
use futures::StreamExt;use jacquard_api::com_atproto::sync::HostStatus;use jacquard_api::com_atproto::sync::list_hosts::{Host, ListHostsOutput};use miette::{Context, IntoDiagnostic};use tracing::{debug, info, warn};use url::Url;
use super::firehose::FirehoseHandle;use crate::pds_discovery::{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.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");
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 matches!(host.status, Some(HostStatus::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).await { Ok(adds) => { page_added = adds.added; added += adds.added; already_known = adds.known; rejected = adds.rejected; } Err(error) => { 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() { break; } }
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());
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_ref()); crate::db::pds_meta::set_status(&mut batch, &state.db.filter, hostname.as_str(), status)?; status_updates.push((hostname.as_str().to_string(), status)); }
batch.commit().into_diagnostic()?; state.db.persist()?; debug!( hosts = hosts.len(), status_updates = status_updates.len(), "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(())}
fn map_seed_status(status: Option<&HostStatus>) -> crate::pds_meta::HostStatus { match status { Some(HostStatus::Active) | None => crate::pds_meta::HostStatus::Active, Some(HostStatus::Idle) => crate::pds_meta::HostStatus::Idle, Some(HostStatus::Offline) => crate::pds_meta::HostStatus::Offline, Some(HostStatus::Throttled) => crate::pds_meta::HostStatus::Throttled, Some(HostStatus::Banned) => crate::pds_meta::HostStatus::Banned, Some(HostStatus::Other(_)) => crate::pds_meta::HostStatus::Active, }}
#[cfg(test)]mod tests { use super::*; use crate::config::Config; use crate::control::Hydrant; use crate::db::keys; 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() }
#[tokio::test] async fn apply_seed_snapshot_persists_statuses_and_preserves_absent_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 hosts = vec![ Host { hostname: "active.example".into(), account_count: Some(42), seq: Some(100), status: Some(HostStatus::Active), extra_data: None, }, Host { hostname: "offline.example".into(), account_count: Some(7), seq: Some(5), status: Some(HostStatus::Offline), extra_data: None, }, ];
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")?, None); assert_eq!(persisted_cursor(&state, "offline.example")?, None);
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(HostStatus::Active), extra_data: None, }];
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 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(HostStatus::Banned), extra_data: None, }; 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(HostStatus::Active), extra_data: None, }; 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(HostStatus::Active), extra_data: None, }];
// 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())}