Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558use crate::db::{Db, keys};use fjall::OwnedWriteBatch;use miette::{Context, IntoDiagnostic, Result};use smol_str::SmolStr;use std::collections::{BTreeMap, BTreeSet};use std::sync::Arc;use std::sync::Mutex;use std::sync::atomic::Ordering;use tracing::error;
#[derive(Debug, Clone, Default)]pub struct CountDeltas { deltas: BTreeMap<SmolStr, i64>,}
impl CountDeltas { fn add(&mut self, key: &str, delta: i64) { if delta == 0 { return; }
let entry = self.deltas.entry(SmolStr::new(key)).or_insert(0); *entry += delta; if *entry == 0 { self.deltas.remove(key); } }
pub(crate) fn add_repos(&mut self, delta: i64) { self.add("repos", delta); }
#[cfg(feature = "indexer")] pub(crate) fn add_records(&mut self, delta: i64) { self.add("records", delta); }
#[cfg(feature = "indexer")] pub(crate) fn add_history(&mut self, delta: i64) { self.add("history", delta); }
pub(crate) fn add_pds_account(&mut self, host: &str, delta: i64) { self.add(&keys::pds_account_count_key(host), delta); }
pub(crate) fn projected_pds_account_count(&self, db: &Db, host: &str) -> u64 { self.projected_count(db, &keys::pds_account_count_key(host)) }
#[cfg(feature = "indexer")] pub(crate) fn add_gauge_diff( &mut self, old: &crate::types::GaugeState, new: &crate::types::GaugeState, ) { use crate::types::GaugeState;
if old == new { return; }
match (old, new) { (GaugeState::Pending, GaugeState::Pending) => {} (GaugeState::Pending, _) => self.add("pending", -1), (_, GaugeState::Pending) => self.add("pending", 1), _ => {} }
match (old.is_resync(), new.is_resync()) { (true, false) => self.add("resync", -1), (false, true) => self.add("resync", 1), _ => {} }
if let GaugeState::Resync(Some(kind)) = old { self.add( match kind { crate::types::ResyncErrorKind::Ratelimited => "error_ratelimited", crate::types::ResyncErrorKind::Transport => "error_transport", crate::types::ResyncErrorKind::Generic => "error_generic", }, -1, ); }
if let GaugeState::Resync(Some(kind)) = new { self.add( match kind { crate::types::ResyncErrorKind::Ratelimited => "error_ratelimited", crate::types::ResyncErrorKind::Transport => "error_transport", crate::types::ResyncErrorKind::Generic => "error_generic", }, 1, ); } }
pub(crate) fn is_empty(&self) -> bool { self.deltas.is_empty() }
pub(crate) fn len(&self) -> usize { self.deltas.len() }
pub(crate) fn get(&self, key: &str) -> i64 { self.deltas.get(key).copied().unwrap_or(0) }
fn projected_count(&self, db: &Db, key: &str) -> u64 { apply_count_delta(db.get_count_sync(key), self.get(key)) }
pub(crate) fn iter(&self) -> impl Iterator<Item = (&SmolStr, &i64)> { self.deltas.iter() }}
pub(crate) struct CountDeltaReservation { in_flight: Arc<Mutex<BTreeSet<u64>>>, start_id: u64,}
impl Drop for CountDeltaReservation { fn drop(&mut self) { let Ok(mut in_flight) = self.in_flight.lock() else { error!( start_id = self.start_id, "count delta reservations poisoned" ); return; }; in_flight.remove(&self.start_id); }}
pub(crate) fn apply_count_delta(current: u64, delta: i64) -> u64 { if delta >= 0 { current.saturating_add(delta as u64) } else { current.saturating_sub(delta.unsigned_abs()) }}
pub fn set_ks_count(batch: &mut OwnedWriteBatch, db: &Db, name: &str, count: u64) { let key = keys::count_keyspace_key(name); batch.insert(&db.counts, key, count.to_be_bytes());}
pub fn set_count_delta_watermark(batch: &mut OwnedWriteBatch, db: &Db, watermark: u64) { batch.insert( &db.counts, keys::COUNT_DELTA_WATERMARK_KEY, watermark.to_be_bytes(), );}
pub fn load_count_delta_watermark(db: &Db) -> Result<u64> { db.counts .get(keys::COUNT_DELTA_WATERMARK_KEY) .into_diagnostic()? .map(|value| read_u64_counter(&value)) .transpose() .map(|watermark| watermark.unwrap_or(0))}
pub(crate) fn read_u64_counter(value: &[u8]) -> Result<u64> { value .try_into() .into_diagnostic() .wrap_err("counter value must be 8 bytes") .map(u64::from_be_bytes)}
fn read_i64_counter_delta(value: &[u8]) -> Result<i64> { value .try_into() .into_diagnostic() .wrap_err("counter delta must be 8 bytes") .map(i64::from_be_bytes)}
pub(crate) fn replay_count_deltas(db: &Db, watermark: u64) -> Result<()> { let start = watermark.saturating_add(1); for (name, delta) in load_count_delta_range(db, start, u64::MAX)? { db.update_count(&name, delta); } Ok(())}
fn load_count_delta_range(db: &Db, start: u64, end: u64) -> Result<BTreeMap<SmolStr, i64>> { let mut aggregated = BTreeMap::new(); for guard in db.counts.range(keys::count_delta_start_key(start)..) { let (key, value) = guard.into_inner().into_diagnostic()?; if !key.starts_with(keys::COUNT_DELTA_PREFIX) { break; }
let (id, name) = keys::parse_count_delta_key(&key)?; if id > end { break; }
*aggregated.entry(SmolStr::new(name)).or_insert(0) += read_i64_counter_delta(&value)?; } Ok(aggregated)}
fn get_persisted_ks_count(db: &Db, name: &str) -> Result<u64> { db.counts .get(keys::count_keyspace_key(name)) .into_diagnostic()? .map(|value| read_u64_counter(&value)) .transpose() .map(|count| count.unwrap_or(0))}
impl Db { pub(super) fn stage_count_deltas( &self, batch: &mut OwnedWriteBatch, deltas: &CountDeltas, ) -> Option<CountDeltaReservation> { if deltas.is_empty() { return None; }
let start_id = self .next_count_delta_id .fetch_add(deltas.len() as u64, Ordering::SeqCst); self.count_delta_in_flight .lock() .expect("count delta reservations poisoned") .insert(start_id);
for (offset, (key, delta)) in deltas.iter().enumerate() { batch.insert( &self.counts, keys::count_delta_key(start_id + offset as u64, key), delta.to_be_bytes(), ); }
Some(CountDeltaReservation { in_flight: self.count_delta_in_flight.clone(), start_id, }) }
pub(super) fn apply_count_deltas(&self, deltas: &CountDeltas) { for (key, delta) in deltas.iter() { self.update_count(key, *delta); } }
pub fn checkpoint_count_deltas(&self) -> Result<Option<u64>> { let start = self .count_delta_checkpoint_watermark .load(Ordering::SeqCst) .saturating_add(1);
let mut end = self .next_count_delta_id .load(Ordering::SeqCst) .saturating_sub(1);
let lowest_in_flight = self .count_delta_in_flight .lock() .expect("count delta reservations poisoned") .first() .copied(); if let Some(lowest_in_flight) = lowest_in_flight { end = end.min(lowest_in_flight.saturating_sub(1)); }
if end < start { return Ok(None); }
let aggregated = load_count_delta_range(self, start, end)?; let mut batch = self.inner.batch();
for (name, delta) in aggregated { let current = get_persisted_ks_count(self, &name)?; set_ks_count(&mut batch, self, &name, apply_count_delta(current, delta)); }
set_count_delta_watermark(&mut batch, self, end); batch.commit().into_diagnostic()?; self.count_delta_checkpoint_watermark .store(end, Ordering::SeqCst);
Ok(Some(end)) }
pub fn mark_count_checkpoint_persisted(&self, watermark: u64) { self.count_delta_gc_watermark .store(watermark, Ordering::SeqCst); }
pub fn update_count(&self, key: &str, delta: i64) -> u64 { let mut entry = self.counts_map.entry_sync(SmolStr::new(key)).or_insert(0); if delta < 0 && *entry < delta.unsigned_abs() { error!( key, current = *entry, decrement = delta.unsigned_abs(), "count underflow !!! this is a bug" ); *entry = 0; } else { *entry = apply_count_delta(*entry, delta); } *entry }
pub async fn update_count_async(&self, key: &str, delta: i64) { let mut entry = self .counts_map .entry_async(SmolStr::new(key)) .await .or_insert(0); if delta < 0 && *entry < delta.unsigned_abs() { error!( key, current = *entry, decrement = delta.unsigned_abs(), "count underflow !!! this is a bug" ); *entry = 0; } else { *entry = apply_count_delta(*entry, delta); } }
pub async fn get_count(&self, key: &str) -> u64 { self.counts_map .read_async(key, |_, v| *v) .await .unwrap_or(0) }
pub fn get_count_sync(&self, key: &str) -> u64 { self.counts_map.read_sync(key, |_, v| *v).unwrap_or(0) }}
#[cfg(test)]mod tests { use super::*;
fn test_config(path: &std::path::Path) -> crate::config::Config { crate::config::Config { database_path: path.to_path_buf(), ..Default::default() } }
#[test] fn count_deltas_replay_and_checkpoint_across_restart() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let cfg = test_config(tmp.path());
{ let db = Db::open(&cfg)?; let mut batch = db.inner.batch(); set_ks_count(&mut batch, &db, "repos", 10); batch.commit().into_diagnostic()?; let mut batch = db.inner.batch(); let mut deltas = CountDeltas::default(); deltas.add_repos(2); #[cfg(feature = "indexer")] deltas.add_records(1); let reservation = db .stage_count_deltas(&mut batch, &deltas) .expect("count deltas should reserve ids"); batch.commit().into_diagnostic()?; db.apply_count_deltas(&deltas); drop(reservation); db.persist()?; }
{ let db = Db::open(&cfg)?; assert_eq!(db.get_count_sync("repos"), 12); #[cfg(feature = "indexer")] assert_eq!(db.get_count_sync("records"), 1); let checkpointed_watermark = db .checkpoint_count_deltas()? .expect("checkpoint should fold pending deltas"); db.persist()?; db.mark_count_checkpoint_persisted(checkpointed_watermark); assert_eq!(load_count_delta_watermark(&db)?, checkpointed_watermark); }
{ let db = Db::open(&cfg)?; assert_eq!(db.get_count_sync("repos"), 12); #[cfg(feature = "indexer")] assert_eq!(db.get_count_sync("records"), 1); assert!(load_count_delta_watermark(&db)? >= 1); }
Ok(()) }
#[test] fn checkpoint_skips_inflight_deltas() -> Result<()> { 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 deltas = CountDeltas::default(); deltas.add_repos(1); let reservation = db .stage_count_deltas(&mut batch, &deltas) .expect("count deltas should reserve ids");
assert!(db.checkpoint_count_deltas()?.is_none());
batch.commit().into_diagnostic()?; db.apply_count_deltas(&deltas); drop(reservation);
assert!(db.checkpoint_count_deltas()?.is_some());
Ok(()) }
#[test] fn checkpointed_count_deltas_are_gcable() -> Result<()> { 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 deltas = CountDeltas::default(); deltas.add_repos(1); let reservation = db .stage_count_deltas(&mut batch, &deltas) .expect("count deltas should reserve ids"); batch.commit().into_diagnostic()?; db.apply_count_deltas(&deltas); drop(reservation);
let watermark = db .checkpoint_count_deltas()? .expect("checkpoint should fold pending deltas"); db.persist()?; db.mark_count_checkpoint_persisted(watermark);
db.counts.rotate_memtable_and_wait().into_diagnostic()?; db.counts.major_compact().into_diagnostic()?;
assert_eq!(db.counts.prefix(keys::COUNT_DELTA_PREFIX).count(), 0);
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_owned(did_str).into_diagnostic()?; let state = crate::types::v4::RepoState { active, status: crate::types::v4::RepoStatus::Synced, root: None, last_message_time: None, last_updated_at: 0, signing_key: None, pds: Some(CowStr::Borrowed(pds)), handle: None, }; let bytes = rmp_serde::to_vec(&state).into_diagnostic()?; batch.insert(&db.repos, keys::repo_key(&did), bytes); 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.commit().into_diagnostic()?; crate::db::migration::rewind_version_for_test(&db, 5)?; 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(()) }}