diff --git a/src/config.rs b/src/config.rs index 787931d..195c974 100644 --- a/src/config.rs +++ b/src/config.rs @@ -380,6 +380,11 @@ pub struct Config { /// set via `HYDRANT_NEW_HOST_LIMIT`. pub new_host_limit: Option, + /// how often offline firehose sources are automatically retried. + /// set via `HYDRANT_OFFLINE_HOST_RETRY_INTERVAL` (humantime duration, e.g. `30min`). + /// set to `none` to disable automatic retries. + pub offline_host_retry_interval: Option, + /// base URL(s) of relay or aggregator services to seed firehose PDS sources from at startup. /// /// hydrant calls `com.atproto.sync.listHosts` on each URL and adds the returned PDSes @@ -501,6 +506,7 @@ impl Default for Config { enable_backlinks: false, only_index_links: false, new_host_limit: Some(50), + offline_host_retry_interval: Some(Duration::from_secs(30 * 60)), tier_rules: vec![], tier_policy: { let mut tiers = HashMap::new(); @@ -681,6 +687,18 @@ impl Config { .ok() .and_then(|s| s.parse().ok()); + let offline_retry_interval: Option = + match std::env::var("HYDRANT_OFFLINE_HOST_RETRY_INTERVAL") + .ok() + .as_deref() + { + None => defaults.offline_host_retry_interval, + Some("none") => None, + Some(s) => humantime::parse_duration(s) + .ok() + .or(defaults.offline_host_retry_interval), + }; + // start with built-in tier definitions, then layer in any env-defined overrides. // format: HYDRANT_RATE_TIERS=name:base/mul/hourly/daily,... let mut tiers = defaults.tier_policy.tiers.clone(); @@ -795,6 +813,7 @@ impl Config { enable_backlinks, only_index_links, new_host_limit: max_pds_added_per_day, + offline_host_retry_interval: offline_retry_interval, tier_policy, tier_rules, cache_size, @@ -928,6 +947,14 @@ impl fmt::Display for Config { if let Some(limit) = self.new_host_limit { config_line!(f, "max pds/day", limit)?; } + match self.offline_host_retry_interval { + Some(d) => config_line!( + f, + "offline retry interval", + format_args!("{}sec", d.as_secs()) + )?, + None => config_line!(f, "offline retry interval", "disabled")?, + } Ok(()) } } diff --git a/src/control/firehose.rs b/src/control/firehose.rs index af476d0..0b421cc 100644 --- a/src/control/firehose.rs +++ b/src/control/firehose.rs @@ -46,8 +46,8 @@ pub struct FirehoseHandle { pub(super) shared: Arc>, /// per-relay running tasks, keyed by url. pub(super) tasks: Arc>, - /// set of known source urls, includes API-added (db-persisted) and static config sources. - pub(super) known_sources: Arc>, + /// known source urls → is_pds flag; includes API-added (db-persisted) and static config sources. + pub(super) known_sources: Arc>, /// ids assigned to spawned tasks next_task_id: Arc, } @@ -58,7 +58,7 @@ impl FirehoseHandle { state, shared: Arc::new(std::sync::OnceLock::new()), tasks: Arc::new(scc::HashMap::new()), - known_sources: Arc::new(scc::HashSet::new()), + known_sources: Arc::new(scc::HashMap::new()), next_task_id: Arc::new(AtomicUsize::new(0)), } } @@ -205,7 +205,7 @@ impl FirehoseHandle { .await .into_diagnostic()??; - let _ = self.known_sources.insert_async(url.clone()).await; + let _ = self.known_sources.insert_async(url.clone(), is_pds).await; // reset failure state so the fresh task gets a clean slate. // if the previous task exited after max failures, the failure counter @@ -246,6 +246,21 @@ impl FirehoseHandle { Ok(self.tasks.remove_async(url).await.is_some()) } + /// restart an offline firehose source without touching the database or daily limits. + pub(super) async fn restart_source(&self, url: Url, is_pds: bool) -> Result<()> { + let shared = self + .shared + .get() + .ok_or_else(|| miette::miette!("firehose worker not started"))?; + + // clear the failure counter so the new task isn't immediately terminated + let throttle = self.state.throttler.get_handle(&url).await; + throttle.record_success(); + + self.spawn_firehose_ingestor(&FirehoseSource { url, is_pds }, shared, true) + .await + } + /// reset the stored firehose cursor for a given URL. pub async fn reset_cursor(&self, url: &str) -> Result<()> { let url = Url::parse(url).into_diagnostic()?; diff --git a/src/control/mod.rs b/src/control/mod.rs index 9648c62..20304b4 100644 --- a/src/control/mod.rs +++ b/src/control/mod.rs @@ -55,6 +55,7 @@ use firehose::FirehoseShared; use stream::event_stream_thread; #[cfg(feature = "relay")] use stream::relay_stream_thread; +use url::Url; #[derive(Debug, Clone)] /// infromation about a host hydrant is consuming from. @@ -416,7 +417,7 @@ impl Hydrant { for source in &relay_hosts { let _ = firehose .known_sources - .insert_async(source.url.clone()) + .insert_async(source.url.clone(), source.is_pds) .await; firehose .spawn_firehose_ingestor(source, fire_shared, true) @@ -434,7 +435,7 @@ impl Hydrant { for source in &persisted_sources { let _ = firehose .known_sources - .insert_async(source.url.clone()) + .insert_async(source.url.clone(), source.is_pds) .await; if firehose.tasks.contains_async(&source.url).await { continue; @@ -454,6 +455,41 @@ impl Hydrant { }); } + // 10d. periodic retry of offline firehose sources + if let Some(retry_interval) = config.offline_host_retry_interval { + tokio::spawn({ + let firehose = firehose.clone(); + async move { + loop { + tokio::time::sleep(retry_interval).await; + + let mut to_restart: Vec<(Url, bool)> = Vec::new(); + { + let meta = firehose.state.pds_meta.load(); + firehose + .known_sources + .iter_async(|url, &is_pds| { + if firehose.tasks.contains_sync(url) { + return true; + } + let host = url.host_str().unwrap_or(url.as_str()); + if meta.is_banned(host) { + return true; + } + to_restart.push((url.clone(), is_pds)); + true + }) + .await; + } + + for (url, is_pds) in to_restart { + let _ = firehose.restart_source(url, is_pds).await; + } + } + } + }); + } + // 11. spawn crawler infrastructure #[cfg(feature = "indexer")] {