use std::collections::BTreeMap;
use crate::db::counts::pds_host;
use crate::db::keys::{self, COUNT_KS_PREFIX};
use crate::db::{Db, set_ks_count};
use crate::types::v4;
use fjall::OwnedWriteBatch;
use miette::{Context, IntoDiagnostic, Result};
use smol_str::SmolStr;
pub(crate) fn rebuild_pds_account_counts(db: &Db, batch: &mut OwnedWriteBatch) -> Result<bool> {
recount_pds_accounts(db, batch, |value| {
let state: v4::RepoState = rmp_serde::from_slice(value)
.into_diagnostic()
.wrap_err("invalid v5 repo state")?;
Ok(counted_host(state.active, state.pds.as_deref()))
})
}
/// the host a repo row with this `active` and `pds` counts one account for
pub(super) fn counted_host(active: bool, pds: Option<&str>) -> Option<SmolStr> {
pds.filter(|_| active).and_then(pds_host)
}
/// replaces every per-pds account count, and its pending deltas, with what `host_of` says
/// the repo rows count for, so a host none of them count for anymore goes to zero
pub(super) fn recount_pds_accounts(
db: &Db,
batch: &mut OwnedWriteBatch,
host_of: impl Fn(&[u8]) -> Result<Option<SmolStr>>,
) -> Result<bool> {
let mut counts: BTreeMap<SmolStr, u64> = BTreeMap::new();
for guard in db.repos.iter() {
let (_, value) = guard.into_inner().into_diagnostic()?;
if let Some(host) = host_of(&value)? {
*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(true)
}