diff --git a/src/control/pds.rs b/src/control/pds.rs index 9233983..01f4986 100644 --- a/src/control/pds.rs +++ b/src/control/pds.rs @@ -50,7 +50,7 @@ impl PdsControl { async fn update(&self, db_op: F, mem_op: G) -> Result<()> where F: FnOnce(&mut fjall::OwnedWriteBatch, &fjall::Keyspace) + Send + 'static, - G: FnOnce(&mut PdsMeta), + G: FnOnce(&mut PdsMeta) + Send + 'static, { let _admission = self.0.pds_admission.lock().await; self.update_locked(db_op, mem_op).await @@ -59,23 +59,25 @@ impl PdsControl { async fn update_locked(&self, db_op: F, mem_op: G) -> Result<()> where F: FnOnce(&mut fjall::OwnedWriteBatch, &fjall::Keyspace) + Send + 'static, - G: FnOnce(&mut PdsMeta), + G: FnOnce(&mut PdsMeta) + Send + 'static, { + let status_write = self.0.pds_status_write.clone(); + let pds_meta = self.0.pds_meta.clone(); self.0 .db .run(move |db| { + let _status_write = status_write.lock(); let mut batch = db.inner.batch(); db_op(&mut batch, &db.filter); batch.commit().into_diagnostic()?; - db.persist() - }) - .await?; - - let mut snapshot = (**self.0.pds_meta.load()).clone(); - mem_op(&mut snapshot); - self.0.pds_meta.store(Arc::new(snapshot)); + db.persist()?; - Ok(()) + let mut snapshot = (**pds_meta.load()).clone(); + mem_op(&mut snapshot); + pds_meta.store(Arc::new(snapshot)); + Ok(()) + }) + .await } /// list all current per-PDS tier assignments (explicit api-assigned overrides only). diff --git a/src/control/seed.rs b/src/control/seed.rs index 3c50dc2..5adc236 100644 --- a/src/control/seed.rs +++ b/src/control/seed.rs @@ -220,6 +220,7 @@ fn body_preview(bytes: &[u8]) -> String { 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()); @@ -232,22 +233,14 @@ async fn apply_seed_snapshot(state: &Arc, hosts: &[Host]) -> miette::R } }; let status = map_seed_status(host.status.as_ref()); - let current = state - .db - .filter - .get(crate::db::pds_meta::pds_status_key(hostname.as_str())) - .into_diagnostic()? - .map(|stored| { - rmp_serde::from_slice::(stored.as_ref()) - .into_diagnostic() - }) - .transpose()?; - if current == Some(crate::pds_meta::HostStatus::Banned) { - continue; + if crate::db::pds_meta::set_status_preserving_ban( + &mut batch, + &state.db.filter, + hostname.as_str(), + status, + )? { + status_updates.push((hostname.as_str().to_string(), status)); } - - 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()?; diff --git a/src/db/pds_meta.rs b/src/db/pds_meta.rs index 0ed060e..476ce2e 100644 --- a/src/db/pds_meta.rs +++ b/src/db/pds_meta.rs @@ -67,6 +67,29 @@ pub mod v5 { Ok(()) } + /// stage an operational status unless the durable status is an operator ban. + /// + /// callers must update their in-memory mirror only when this returns `true`; + /// the ban is the durable authority and can be cleared only by `unban`. + pub fn set_status_preserving_ban( + batch: &mut OwnedWriteBatch, + ks: &Keyspace, + host: &str, + status: HostStatus, + ) -> Result { + let current = ks + .get(pds_status_key(host)) + .into_diagnostic()? + .map(|stored| rmp_serde::from_slice::(stored.as_ref())) + .transpose() + .into_diagnostic()?; + if current == Some(HostStatus::Banned) { + return Ok(false); + } + set_status(batch, ks, host, status)?; + Ok(true) + } + pub fn remove_status(batch: &mut OwnedWriteBatch, ks: &Keyspace, host: &str) { batch.remove(ks, pds_status_key(host)); } diff --git a/src/ingest/firehose.rs b/src/ingest/firehose.rs index c7d9ed6..300f1aa 100644 --- a/src/ingest/firehose.rs +++ b/src/ingest/firehose.rs @@ -500,8 +500,17 @@ impl FirehoseIngestor { debug!(host = %host, status = ?status, "updating host status"); + let _status_write = self.state.pds_status_write.lock(); let mut batch = self.state.db.inner.batch(); - crate::db::pds_meta::set_status(&mut batch, &self.state.db.filter, host, status)?; + let applied = crate::db::pds_meta::set_status_preserving_ban( + &mut batch, + &self.state.db.filter, + host, + status, + )?; + if !applied { + return Ok(()); + } batch.commit().into_diagnostic()?; crate::pds_meta::PdsMeta::update_host(&self.state.pds_meta, host, |h| h.status = status); @@ -621,8 +630,60 @@ impl FirehoseIngestor { mod tests { use super::*; use crate::config::Config; + use crate::control::Hydrant; use std::sync::atomic::AtomicI64; + #[tokio::test] + async fn banned_host_keeps_status_through_operational_transitions() -> Result<()> { + let tmp = tempfile::tempdir().into_diagnostic()?; + let cfg = Config { + database_path: tmp.path().to_path_buf(), + ..Config::default() + }; + let hydrant = Hydrant::new(cfg.clone()).await?; + hydrant.firehose.startable_for_test(); + let host = "banned-transition.example"; + let relay_host = Url::parse("wss://banned-transition.example/").into_diagnostic()?; + + hydrant.pds.ban(host).await?; + let ingestor = FirehoseIngestor::new( + hydrant.state.clone(), + BufferTx::channel(1).0, + relay_host, + true, + hydrant.state.filter.clone(), + hydrant.state.firehose_enabled.subscribe(), + false, + 1, + ) + .await; + ingestor.set_host_status(HostStatus::Offline)?; + ingestor.set_host_status(HostStatus::Active)?; + + assert_eq!( + hydrant.state.pds_meta.load().status(host), + HostStatus::Banned + ); + drop(ingestor); + drop(hydrant); + + let reopened = Hydrant::new(cfg).await?; + assert_eq!( + reopened.state.pds_meta.load().status(host), + HostStatus::Banned + ); + assert_eq!( + reopened + .firehose + .add_source_validated( + Url::parse("wss://banned-transition.example/").into_diagnostic()? + ) + .await?, + Err(crate::control::firehose::PdsAdmissionRejection::HostBanned) + ); + Ok(()) + } + #[tokio::test] async fn clear_stale_cursor_resets_db_and_memory() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; diff --git a/src/state.rs b/src/state.rs index 1576089..92de251 100644 --- a/src/state.rs +++ b/src/state.rs @@ -42,6 +42,8 @@ pub struct AppState { pub(crate) pds_meta: PdsMetaHandle, /// Serializes validated PDS admission. pub(crate) pds_admission: Mutex<()>, + /// serializes durable host-status writes with their in-memory mirrors. + pub(crate) pds_status_write: std::sync::Arc>, pub(crate) pds_daily_limit: PdsDailyLimit, pub(crate) tier_policy: TierPolicy, pub firehose_cursors: scc::HashIndex, @@ -136,6 +138,7 @@ impl AppState { filter, pds_meta, pds_admission: Mutex::new(()), + pds_status_write: std::sync::Arc::new(parking_lot::Mutex::new(())), tier_policy: config.tier_policy.clone(), firehose_cursors: relay_cursors, firehose_stats: FirehoseStats::default(), @@ -287,6 +290,7 @@ impl AppState { use crate::pds_meta::PdsMeta; use tracing::{debug, error}; + let _status_write = self.pds_status_write.lock(); let (current_status, limit) = { let meta = self.pds_meta.load(); let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref()); @@ -302,12 +306,17 @@ impl AppState { debug!(%host, count, ?limit, ?new_status, "account count crossed limit, shifting status"); - if let Err(e) = db_pds::set_status(batch, &self.db.filter, host, new_status) { - error!(%host, err = %e, "failed to write host status"); - return; + let applied = + match db_pds::set_status_preserving_ban(batch, &self.db.filter, host, new_status) { + Ok(applied) => applied, + Err(e) => { + error!(%host, err = %e, "failed to write host status"); + return; + } + }; + if applied { + PdsMeta::update_host(&self.pds_meta, host, |h| h.status = new_status); } - - PdsMeta::update_host(&self.pds_meta, host, |h| h.status = new_status); } /// checks whether `host` is at or over its account limit at the given count. @@ -333,9 +342,48 @@ impl AppState { mod tests { use super::*; use crate::config::Config; + use crate::pds_meta::HostStatus; + use miette::IntoDiagnostic; use std::time::Duration; use tokio::time::timeout; + #[test] + fn banned_host_survives_an_account_limit_transition() -> Result<()> { + let tmp = tempfile::tempdir().into_diagnostic()?; + let cfg = Config { + database_path: tmp.path().to_path_buf(), + ..Config::default() + }; + let state = AppState::new(&cfg)?; + let host = "banned-limit.example"; + + let mut ban = state.db.inner.batch(); + crate::db::pds_meta::set_status(&mut ban, &state.db.filter, host, HostStatus::Banned)?; + ban.commit().into_diagnostic()?; + state.db.persist()?; + PdsMeta::update_host(&state.pds_meta, host, |entry| { + entry.status = HostStatus::Banned; + }); + + let mut batch = state.db.inner.batch(); + state.apply_host_limit_status(&mut batch, host, u64::MAX); + batch.commit().into_diagnostic()?; + state.db.persist()?; + + assert_eq!(state.pds_meta.load().status(host), HostStatus::Banned); + let stored = state + .db + .filter + .get(crate::db::pds_meta::pds_status_key(host)) + .into_diagnostic()? + .expect("banned status must remain durable"); + assert_eq!( + rmp_serde::from_slice::(stored.as_ref()).into_diagnostic()?, + HostStatus::Banned + ); + Ok(()) + } + #[tokio::test] async fn test_with_ingestion_paused_restore_on_cancel() -> Result<()> { let temp = tempfile::tempdir().unwrap();