diff --git a/src/control/firehose.rs b/src/control/firehose.rs index 07f74ad..37d09c1 100644 --- a/src/control/firehose.rs +++ b/src/control/firehose.rs @@ -470,6 +470,7 @@ impl FirehoseHandle { rmp_serde::to_vec(&db::FirehoseSourceMeta { is_pds: true }).map_err(|error| { miette::miette!("failed to serialize firehose source meta: {error}") })?; + let source_url = source.url.clone(); self.state .db .run(move |db| { @@ -493,7 +494,13 @@ impl FirehoseHandle { } batch.insert(&db.crawler, key, value); batch.commit().into_diagnostic()?; - db.persist()?; + if let Err(error) = db.persist() { + tracing::error!( + source = %source_url, + err = %error, + "PDS source committed but durability persist failed" + ); + } Ok(true) }) .await @@ -764,6 +771,37 @@ mod tests { Ok(()) } + #[tokio::test] + async fn committed_admission_is_not_refunded_when_persist_fails() -> Result<()> { + let dir = tempfile::tempdir().into_diagnostic()?; + let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?; + startable(&hydrant); + let url = Url::parse("wss://post-commit-persist.example/").into_diagnostic()?; + hydrant.state.db.fail_next_persist_for_test(); + + assert_eq!( + hydrant.firehose.add_source_validated(url.clone()).await?, + Ok(PdsAdmission::StartedNew) + ); + assert!(hydrant.firehose.is_source_running(&url)); + assert!(matches!( + crate::db::load_pds_daily_adds(&hydrant.state.db)?, + Some((_, 1)) + )); + assert_eq!( + hydrant + .firehose + .add_source_validated( + Url::parse("wss://second-post-commit.example/").into_diagnostic()? + ) + .await?, + Err(PdsAdmissionRejection::DailyLimitReached) + ); + + hydrant.firehose.tasks.remove_async(&url).await; + Ok(()) + } + #[tokio::test] async fn running_validated_source_does_not_spend_another_slot() -> Result<()> { let dir = tempfile::tempdir().into_diagnostic()?; diff --git a/src/db/mod.rs b/src/db/mod.rs index e6ce3cc..d4b8f7e 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -8,6 +8,8 @@ use smol_str::SmolStr; use std::collections::BTreeSet; use std::sync::atomic::AtomicU64; +#[cfg(test)] +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use url::Url; @@ -77,6 +79,8 @@ pub struct Db { count_delta_gc_watermark: Arc, count_delta_in_flight: Arc>>, pub(crate) compaction_running: Arc, + #[cfg(test)] + pub(crate) persist_failures: Arc, /// 256 lock-sharded mutexes keyed by one byte of the trimmed DID. this is /// synchronization only, not storage partitioning: unrelated DIDs may /// conservatively share a lock. record mutations hold the selected lock @@ -111,6 +115,21 @@ impl Db { } pub fn persist(&self) -> Result<()> { + #[cfg(test)] + if self + .persist_failures + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| { + if remaining > 0 { + Some(remaining - 1) + } else { + None + } + }) + .is_ok() + { + return Err(miette::miette!("injected persist failure")); + } + #[cfg(not(feature = "__persist_sync_all"))] const MODE: PersistMode = PersistMode::Buffer; #[cfg(feature = "__persist_sync_all")] @@ -119,6 +138,11 @@ impl Db { Ok(()) } + #[cfg(test)] + pub(crate) fn fail_next_persist_for_test(&self) { + self.persist_failures.fetch_add(1, Ordering::Release); + } + pub async fn compact(&self) -> Result<()> { use std::sync::atomic::Ordering; if self diff --git a/src/db/open.rs b/src/db/open.rs index bddf6c1..20cc528 100644 --- a/src/db/open.rs +++ b/src/db/open.rs @@ -4,6 +4,8 @@ use miette::{Context, IntoDiagnostic, Result}; use scc::HashMap; use smol_str::SmolStr; use std::collections::{BTreeSet, HashMap as StdHashMap}; +#[cfg(test)] +use std::sync::atomic::AtomicUsize; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; @@ -193,6 +195,8 @@ impl Db { count_delta_gc_watermark, count_delta_in_flight: Arc::new(Mutex::new(BTreeSet::new())), compaction_running: Arc::new(std::sync::atomic::AtomicBool::new(false)), + #[cfg(test)] + persist_failures: Arc::new(AtomicUsize::new(0)), #[cfg(feature = "indexer")] record_write_locks: Arc::new(std::array::from_fn(|_| std::sync::Mutex::new(()))), })