diff --git a/src/control/seed.rs b/src/control/seed.rs --- a/src/control/seed.rs +++ b/src/control/seed.rs @@ -9,7 +9,7 @@ use url::Url; use super::firehose::FirehoseHandle; -use crate::db::{self, CountDeltas, keys}; +use crate::db::{self, keys}; use crate::state::AppState; const MAX_CONCURRENT_SEEDS: usize = 4; @@ -130,32 +130,6 @@ // skip sources that are already running if firehose.tasks.contains_async(&wss_url).await { continue; - } - - // initialise account count for hosts we haven't seen before - if let Some(count) = host.account_count.filter(|&c| c > 0) { - let count_key = keys::pds_account_count_key(host.hostname.as_ref()); - let current = state.db.get_count(&count_key).await; - if current == 0 { - let state = state.clone(); - let count_key = count_key.clone(); - let result = tokio::task::spawn_blocking(move || -> miette::Result<()> { - let mut batch = state.db.inner.batch(); - let mut count_deltas = CountDeltas::default(); - count_deltas.add(&count_key, count); - let reservation = state.db.stage_count_deltas(&mut batch, &count_deltas); - batch.commit().into_diagnostic()?; - state.db.apply_count_deltas(&count_deltas); - drop(reservation); - Ok(()) - }) - .await - .into_diagnostic() - .flatten(); - if let Err(e) = result { - warn!(hostname = %host.hostname, err = %e, "failed to seed host account count"); - } - } } match firehose.add_source(wss_url, true).await { diff --git a/src/db/mod.rs b/src/db/mod.rs --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -1226,4 +1226,91 @@ Ok(()) } + + #[test] + fn pds_account_count_migration_rebuilds_from_active_repos() -> Result<()> { + use jacquard_common::CowStr; + use jacquard_common::types::string::Did; + + let tmp = tempfile::tempdir().into_diagnostic()?; + let cfg = test_config(tmp.path()); + + { + let db = Db::open(&cfg)?; + let mut batch = db.inner.batch(); + + let mut insert_repo = |did_str: &str, pds: &'static str, active: bool| -> Result<()> { + let did = Did::new(did_str).into_diagnostic()?; + let mut state = RepoState::backfilling(); + state.active = active; + state.pds = Some(CowStr::Borrowed(pds)); + batch.insert(&db.repos, keys::repo_key(&did), ser_repo_state(&state)?); + Ok(()) + }; + + insert_repo("did:web:one.test", "https://pds.example/", true)?; + insert_repo("did:web:two.test", "https://pds.example/", true)?; + insert_repo("did:web:inactive.test", "https://pds.example/", false)?; + insert_repo("did:web:other.test", "https://other.example/", true)?; + + set_ks_count( + &mut batch, + &db, + &keys::pds_account_count_key("pds.example"), + 100, + ); + set_ks_count( + &mut batch, + &db, + &keys::pds_account_count_key("other.example"), + 42, + ); + set_ks_count( + &mut batch, + &db, + &keys::pds_account_count_key("orphan.example"), + 7, + ); + batch.insert( + &db.counts, + keys::count_delta_key(1, &keys::pds_account_count_key("pds.example")), + 5_i64.to_be_bytes(), + ); + batch.insert( + &db.counts, + keys::count_delta_key(2, "repos"), + 1_i64.to_be_bytes(), + ); + batch.insert(&db.counts, keys::VERSIONING_KEY, 5_u64.to_be_bytes()); + batch.commit().into_diagnostic()?; + db.persist()?; + } + + let db = Db::open(&cfg)?; + assert_eq!( + db.get_count_sync(&keys::pds_account_count_key("pds.example")), + 2 + ); + assert_eq!( + db.get_count_sync(&keys::pds_account_count_key("other.example")), + 1 + ); + assert_eq!( + db.get_count_sync(&keys::pds_account_count_key("orphan.example")), + 0 + ); + assert_eq!(db.get_count_sync("repos"), 1); + + let mut pds_delta_count = 0; + for guard in db.counts.prefix(keys::COUNT_DELTA_PREFIX) { + let key = guard.key().into_diagnostic()?; + let (_, name) = keys::parse_count_delta_key(&key)?; + if name.starts_with("p|") { + pds_delta_count += 1; + } + } + assert_eq!(pds_delta_count, 0); + + Ok(()) + } } diff --git a/src/ingest/relay.rs b/src/ingest/relay.rs --- a/src/ingest/relay.rs +++ b/src/ingest/relay.rs @@ -410,6 +410,8 @@ return Ok(()); } repo_state.advance_message_time(event_ms); + let was_active = repo_state.active; + let was_pds_host = Self::pds_host(repo_state.pds.as_deref()); #[cfg(feature = "indexer")] let (was_handle, was_signing_key) = ( @@ -436,6 +438,14 @@ if is_pds && repo_state.handle != identity.handle { identity.handle = None; } + + Self::update_pds_account_count( + ctx, + was_active, + was_pds_host.as_deref(), + repo_state.active, + Self::pds_host(repo_state.pds.as_deref()).as_deref(), + ); let repo_key = keys::repo_key(&identity.did); @@ -476,7 +486,7 @@ repo_state: &mut RepoState, firehose: &Url, #[allow(unused_mut)] mut account: Account<'static>, - is_pds: bool, + _is_pds: bool, ) -> Result<()> { let event_ms = account.time.0.timestamp_millis(); if repo_state.last_message_time.is_some_and(|t| event_ms <= t) { @@ -488,6 +498,7 @@ // always capture was_active for count tracking, not just in indexer mode let was_active = repo_state.active; + let was_pds_host = Self::pds_host(repo_state.pds.as_deref()); #[cfg(feature = "indexer")] let was_status = repo_state.status.clone(); @@ -517,24 +528,13 @@ }; } - // update per-PDS active account count on transitions - if is_pds && let Some(host) = firehose.host_str() { - let count_key = pds_account_count_key(host); - let delta = if !was_active && repo_state.active { - 1 - } else if was_active && !repo_state.active { - -1 - } else { - 0 - }; - - if delta != 0 { - ctx.count_deltas.add(&count_key, delta); - let count = ctx.count_deltas.projected_count(&ctx.state.db, &count_key); - ctx.state - .apply_host_limit_status(&mut ctx.batch, host, count); - } - } + Self::update_pds_account_count( + ctx, + was_active, + was_pds_host.as_deref(), + repo_state.active, + Self::pds_host(repo_state.pds.as_deref()).as_deref(), + ); let repo_key = keys::repo_key(&account.did); @@ -572,6 +572,39 @@ ); Ok(()) + } + + fn pds_host(pds: Option<&str>) -> Option { + pds.and_then(|pds| Url::parse(pds).ok()) + .and_then(|url| url.host_str().map(SmolStr::new)) + } + + fn update_pds_account_count( + ctx: &mut WorkerContext, + old_active: bool, + old_host: Option<&str>, + new_active: bool, + new_host: Option<&str>, + ) { + if old_active && old_host == new_host && new_active { + return; + } + + let mut update_host = |host: &str, delta| { + let count_key = pds_account_count_key(host); + ctx.count_deltas.add(&count_key, delta); + let count = ctx.count_deltas.projected_count(&ctx.state.db, &count_key); + ctx.state + .apply_host_limit_status(&mut ctx.batch, host, count); + }; + + if old_active && let Some(host) = old_host { + update_host(host, -1); + } + + if new_active && let Some(host) = new_host { + update_host(host, 1); + } } } @@ -870,6 +903,13 @@ .unwrap_or_else(RepoState::backfilling); repo_state.update_from_doc(doc); + RelayWorker::update_pds_account_count( + self, + false, + None, + repo_state.active, + RelayWorker::pds_host(repo_state.pds.as_deref()).as_deref(), + ); self.batch.insert( &db.repos, @@ -886,14 +926,6 @@ } self.count_deltas.add("repos", 1); - - // track initial active state for per-PDS rate limiting - if msg.is_pds - && repo_state.active - && let Some(host) = msg.firehose.host_str() - { - self.count_deltas.add(&pds_account_count_key(host), 1); - } Ok(Some(repo_state)) } diff --git a/src/db/migration/mod.rs b/src/db/migration/mod.rs --- a/src/db/migration/mod.rs +++ b/src/db/migration/mod.rs @@ -9,6 +9,7 @@ mod v3; mod v4; mod v5; +mod v6; type MigrationFn = fn(&Db, &mut OwnedWriteBatch) -> Result<()>; @@ -19,6 +20,7 @@ ("firehose_source_is_pds", v3::firehose_source_is_pds), ("repo_state_active", v4::repo_state_active), ("pds_meta_layout", v5::pds_meta_layout), + ("rebuild_pds_account_counts", v6::rebuild_pds_account_counts), ]; fn read_version(db: &Db) -> Result { diff --git a/src/db/migration/v6.rs b/src/db/migration/v6.rs new file mode 100644 --- /dev/null +++ b/src/db/migration/v6.rs @@ -0,0 +1,55 @@ +use std::collections::BTreeMap; + +use crate::db::keys::{self, COUNT_KS_PREFIX}; +use crate::db::{Db, deser_repo_state, set_ks_count}; +use fjall::OwnedWriteBatch; +use miette::{Context, IntoDiagnostic, Result}; +use smol_str::SmolStr; +use url::Url; + +pub(crate) fn rebuild_pds_account_counts(db: &Db, batch: &mut OwnedWriteBatch) -> Result<()> { + let mut counts: BTreeMap = BTreeMap::new(); + + for guard in db.repos.iter() { + let (_, value) = guard.into_inner().into_diagnostic()?; + let state = deser_repo_state(value.as_ref())?; + if !state.active { + continue; + } + + let Some(host) = state + .pds + .as_deref() + .and_then(|pds| Url::parse(pds).ok()) + .and_then(|url| url.host_str().map(SmolStr::new)) + else { + continue; + }; + + *counts.entry(host).or_insert(0) += 1; + } + + for guard in db.counts.prefix(COUNT_KS_PREFIX) { + let (key, _) = guard.into_inner().into_diagnostic()?; + let name = std::str::from_utf8(&key[COUNT_KS_PREFIX.len()..]) + .into_diagnostic() + .wrap_err("expected valid utf8 for count key")?; + if name.starts_with("p|") { + batch.remove(&db.counts, key); + } + } + + for guard in db.counts.prefix(keys::COUNT_DELTA_PREFIX) { + let key = guard.key().into_diagnostic()?; + let (_, name) = keys::parse_count_delta_key(&key)?; + if name.starts_with("p|") { + batch.remove(&db.counts, key); + } + } + + for (host, count) in counts { + set_ks_count(batch, db, &keys::pds_account_count_key(&host), count); + } + + Ok(()) +}