diff --git a/src/control/firehose.rs b/src/control/firehose.rs index 85871b9..6ca0cbf 100644 --- a/src/control/firehose.rs +++ b/src/control/firehose.rs @@ -411,30 +411,65 @@ impl FirehoseHandle { /// add a direct PDS source through the public discovery policy. #[cfg_attr(not(feature = "relay"), allow(dead_code))] pub(crate) async fn add_source_validated(&self, url: Url) -> Result { - self.admit_validated_source(url, false).await + self.admit_validated_source(url).await } - async fn admit_validated_source( - &self, - url: Url, - delay_startup: bool, - ) -> Result { + async fn admit_validated_source(&self, url: Url) -> Result { + let (host, url) = match canonical_pds_url(&url) { + Ok(source) => source, + Err(error) => return Ok(Err(PdsAdmissionRejection::InvalidHost(error))), + }; + let _admission = self.state.pds_admission.lock().await; + + if let Some(early) = self + .check_admission_preconditions(&url, host.as_str(), false) + .await? + { + return Ok(early); + } + + let Some(reservation) = self.state.pds_daily_limit.reserve() else { + return Ok(Err(PdsAdmissionRejection::DailyLimitReached)); + }; + + self.spawn_admitted_source(&url, host.as_str(), false, reservation) + .await + } + + async fn admit_seeded_source(&self, url: Url) -> Result { let (host, url) = match canonical_pds_url(&url) { Ok(source) => source, Err(error) => return Ok(Err(PdsAdmissionRejection::InvalidHost(error))), }; let _admission = self.state.pds_admission.lock().await; - if self.state.pds_meta.load().is_banned(host.as_str()) { - return Ok(Err(PdsAdmissionRejection::HostBanned)); + if let Some(early) = self + .check_admission_preconditions(&url, host.as_str(), true) + .await? + { + return Ok(early); + } + + self.spawn_admitted_source(&url, host.as_str(), true, PdsDailyReservation::Unlimited) + .await + } + + async fn check_admission_preconditions( + &self, + url: &Url, + host: &str, + delay_startup: bool, + ) -> Result> { + if self.state.pds_meta.load().is_banned(host) { + return Ok(Some(Err(PdsAdmissionRejection::HostBanned))); } #[cfg(test)] self.pause_validated_admission_after_ban_check_if_requested() .await; - if self.is_source_running(&url) { - return Ok(Ok(PdsAdmission::AlreadyRunning)); + if self.is_source_running(url) { + return Ok(Some(Ok(PdsAdmission::AlreadyRunning))); } - if self.is_source_known(&url) { + if self.is_source_known(url) { // Do not route this through admin add_source: that trusted override // intentionally clears public provenance. A validated request stays // public across restart and persistence. @@ -444,26 +479,32 @@ impl FirehoseHandle { .ok_or_else(|| miette::miette!("firehose worker not started"))?; let is_pds = self .known_sources - .read_async(&url, |_, &is_pds| is_pds) + .read_async(url, |_, &is_pds| is_pds) .await .unwrap_or(true); let source = FirehoseSource { url: url.clone(), is_pds, }; - let throttle = self.state.throttler.get_handle(&url).await; + let throttle = self.state.throttler.get_handle(url).await; throttle.record_success(); let prepared = self - .prepare_firehose_ingestor(&source, shared, self.public_sources.contains_sync(&url)) + .prepare_firehose_ingestor(&source, shared, self.public_sources.contains_sync(url)) .await?; self.spawn_prepared_firehose_ingestor(&source, prepared, delay_startup) .await; - return Ok(Ok(PdsAdmission::RestartedKnown)); + return Ok(Some(Ok(PdsAdmission::RestartedKnown))); } + Ok(None) + } - let Some(reservation) = self.state.pds_daily_limit.reserve() else { - return Ok(Err(PdsAdmissionRejection::DailyLimitReached)); - }; + async fn spawn_admitted_source( + &self, + url: &Url, + host: &str, + delay_startup: bool, + reservation: PdsDailyReservation, + ) -> Result { let shared = match self.shared.get() { Some(shared) => shared, None => { @@ -471,7 +512,10 @@ impl FirehoseHandle { return Err(miette::miette!("firehose worker not started")); } }; - let source = FirehoseSource { url, is_pds: true }; + let source = FirehoseSource { + url: url.clone(), + is_pds: true, + }; // all fallible startup preparation happens before taking a quota slot's write. let prepared = match self.prepare_firehose_ingestor(&source, shared, true).await { Ok(prepared) => prepared, @@ -482,7 +526,7 @@ impl FirehoseHandle { }; match self - .persist_validated_source(&source, host.as_str(), reservation) + .persist_validated_source(&source, host, reservation) .await { Ok(true) => {} @@ -565,8 +609,8 @@ impl FirehoseHandle { /// add PDS sources discovered from a seed relay. /// - /// each accepted source persists its own quota reservation with its source - /// record. startup stays staggered so a page cannot stampede remote hosts. + /// seeded sources bypass the public crawl request daily limit. + /// startup stays staggered so a page cannot stampede remote hosts. pub(super) async fn add_seeded_sources(&self, urls: Vec) -> Result { let mut adds = SeededSourceAdds { added: 0, @@ -578,7 +622,7 @@ impl FirehoseHandle { adds.known += 1; continue; } - match self.admit_validated_source(url.clone(), true).await? { + match self.admit_seeded_source(url.clone()).await? { Ok(PdsAdmission::StartedNew) => adds.added += 1, Ok(PdsAdmission::RestartedKnown | PdsAdmission::AlreadyRunning) => { adds.known += 1; @@ -894,7 +938,7 @@ mod tests { } #[tokio::test] - async fn seeded_admission_stops_at_the_daily_limit() -> Result<()> { + async fn seeded_admission_does_not_spend_daily_limit() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?; startable(&hydrant); @@ -906,14 +950,29 @@ mod tests { .add_seeded_sources(vec![first.clone(), second.clone()]) .await?; - assert_eq!(adds.added, 1); - assert_eq!(adds.rejected, 1); - assert_eq!(persisted_sources(&hydrant).await?.len(), 1); - assert!(matches!( + assert_eq!(adds.added, 2); + assert_eq!(adds.rejected, 0); + assert_eq!(persisted_sources(&hydrant).await?.len(), 2); + assert_eq!( crate::db::load_pds_daily_adds(&hydrant.state.db)?, - Some((_, 1)) - )); + None + ); + + // public validated requests still enforce the daily limit independently + let third = Url::parse("wss://third-public.example/").into_diagnostic()?; + assert_eq!( + hydrant.firehose.add_source_validated(third.clone()).await?, + Ok(PdsAdmission::StartedNew) + ); + let fourth = Url::parse("wss://fourth-public.example/").into_diagnostic()?; + assert_eq!( + hydrant.firehose.add_source_validated(fourth.clone()).await?, + Err(PdsAdmissionRejection::DailyLimitReached) + ); + hydrant.firehose.tasks.remove_async(&first).await; + hydrant.firehose.tasks.remove_async(&second).await; + hydrant.firehose.tasks.remove_async(&third).await; Ok(()) }