[db] refactor how repo counts are handled
1 unmodified line23456789103 unmodified lines1415161717181920212223242525262728293055 unmodified lines86878889899091929335 unmodified lines1291301311321331341351361379 unmodified lines1471481491501511521531541551561571581591601611621631647 unmodified lines1721731741751761771781791804 unmodified lines185186187188189190191192193194195196197198199200201202257 unmodified linesExpand all1 unmodified lineuse fjall::Slice;
#[cfg(feature = "backlinks")]use jacquard_common::Data;use jacquard_common::IntoStatic;use jacquard_common::types::did::Did;use miette::{Context, IntoDiagnostic, Result};use std::collections::HashMap;use tracing::debug;3 unmodified linesuse crate::filter::FilterConfig;use crate::ingest::stream::Commit;use crate::ingest::validation::ValidatedCommit;use crate::state::AppState;use crate::types::{RepoState, RepoStatus, ResyncState};use crate::types::{GaugeState, RepoState, RepoStatus, ResyncErrorKind, ResyncState};
#[cfg(feature = "indexer_stream")]use { crate::types::{ AccountEvt, BroadcastEvent, IdentityEvt, LiveRecordEvent, MarshallableEvt, StoredData, StoredEvent, }, jacquard_common::{CowStr, IntoStatic}, jacquard_common::CowStr, std::sync::atomic::Ordering,};
#[cfg(feature = "jetstream")]55 unmodified lines let metadata = db::deser_repo_meta(&metadata_bytes)?; batch.remove(&db.pending, keys::pending_key(metadata.index_id)); }
// 1. delete from resync, and metadata // we don't delete from repos, relay uses it as a tombstone // todo: we should still delete it after some time batch.remove(&db.resync, &repo_key); batch.remove(&db.repo_metadata, &metadata_key);35 unmodified lines
pub fn transition_repo<'s>( batch: &mut OwnedWriteBatch, db: &Db, lifecycle_transitions: &mut Vec<(Did<'static>, GaugeState)>, did: &Did, mut repo_state: RepoState<'s>, new_status: RepoStatus,) -> Result<RepoState<'s>> {9 unmodified lines
// manage queues match &new_status { RepoStatus::Synced => { lifecycle_transitions.push((did.clone().into_static(), GaugeState::Synced)); batch.remove(&db.pending, pending_key.as_slice()); // we dont have to remove from resync here because it has to transition resync -> pending first } RepoStatus::Error(msg) => { tracing::warn!("transitioning to error: {msg}"); lifecycle_transitions.push(( did.clone().into_static(), GaugeState::Resync(Some(ResyncErrorKind::Generic)), )); batch.remove(&db.pending, pending_key.as_slice()); // TODO: we need to make errors have kind instead of "message" in repo status // and then pass it to resync error kind let resync_state = crate::types::ResyncState::Error {7 unmodified lines rmp_serde::to_vec(&resync_state).into_diagnostic()?, ); } RepoStatus::Deactivated | RepoStatus::Takendown | RepoStatus::Suspended => { lifecycle_transitions.push((did.clone().into_static(), GaugeState::Resync(None))); // this shouldnt be needed since a repo wont be in a pending state when it gets to any of these states // batch.remove(&db.pending, &pending_key); let resync_state = ResyncState::Gone { status: new_status.clone(),4 unmodified lines rmp_serde::to_vec(&resync_state).into_diagnostic()?, ); } RepoStatus::Deleted => { lifecycle_transitions.push((did.clone().into_static(), GaugeState::Synced)); // terminal state: remove from queues, no resync entry needed batch.remove(&db.pending, pending_key.as_slice()); batch.remove(&db.resync, &repo_key); } RepoStatus::Desynchronized | RepoStatus::Throttled => { lifecycle_transitions.push(( did.clone().into_static(), GaugeState::Resync(Some(ResyncErrorKind::Generic)), )); // like an error: remove from pending and schedule a resync attempt batch.remove(&db.pending, pending_key.as_slice()); let resync_state = crate::types::ResyncState::Error { kind: crate::types::ResyncErrorKind::Generic,257 unmodified linesExpand all12234563 unmodified lines10111213141415161718192041 unmodified lines6263646566666768697071727374757576777877787980818 unmodified lines909192939594959697989910011 unmodified lines11211311411511711811911611711811912019 unmodified lines1401411421431471481441501451521531541551561461471481491501511521531641541551561571581591601661671611621631641651661671681699 unmodified lines179180181182185183184185186187188189191192190191192193194use crate::db::types::TrimmedDid;use crate::db::{self, CountDeltas, keys};use crate::db::{self, keys};use crate::state::AppState;use crate::types::{GaugeState, ResyncState};use miette::{IntoDiagnostic, Result};
3 unmodified lines
pub fn queue_gone_backfills(state: &Arc<AppState>) -> Result<()> { debug!("scanning for deactivated/takendown repos to retry..."); let mut transitions = 0usize; let mut count_deltas = CountDeltas::default();
let mut batch = state.db.inner.batch(); let mut lifecycle_counts = state.db.lifecycle_counts();
for guard in state.db.resync.iter() { let (key, val) = guard.into_inner().into_diagnostic()?; let did = match TrimmedDid::try_from(key.as_ref()) {41 unmodified lines &metadata_key, crate::db::ser_repo_meta(&metadata)?, );
count_deltas.add_gauge_diff(&GaugeState::Resync(None), &GaugeState::Pending); lifecycle_counts.transition(&mut batch, &did, GaugeState::Pending)?; transitions += 1; } }
if transitions == 0 { return Ok(()); }
let reservation = state.db.stage_count_deltas(&mut batch, &count_deltas); let lifecycle_reservation = lifecycle_counts.stage(&mut batch); batch.commit().into_diagnostic()?; state.db.apply_count_deltas(&count_deltas); drop(reservation); state.db.apply_lifecycle_counts(lifecycle_reservation);
state.notify_backfill();
info!(count = transitions, "queued gone backfills");8 unmodified lines std::thread::sleep(Duration::from_secs(60));
let now = chrono::Utc::now().timestamp(); let mut transitions = 0usize; let mut count_deltas = CountDeltas::default();
let mut batch = state.db.inner.batch(); let mut lifecycle_counts = state.db.lifecycle_counts();
for guard in db.resync.iter() { let (key, value) = match guard.into_inner() { Ok(t) => t,11 unmodified lines } };
match rmp_serde::from_slice::<ResyncState>(&value) { Ok(ResyncState::Error { kind, next_retry, .. }) => { Ok(ResyncState::Error { next_retry, .. }) => { if next_retry <= now { debug!(did = %did, "retrying backfill");
let metadata_key = keys::repo_metadata_key(&did);19 unmodified lines continue; } };
// move from resync back into pending batch.remove(&state.db.resync, key.clone()); let old_pending = keys::pending_key(metadata.index_id); batch.remove(&state.db.pending, old_pending); metadata.index_id = rand::random::<u64>(); batch.insert( &state.db.pending, keys::pending_key(metadata.index_id), key.clone(), ); let new_pending = keys::pending_key(metadata.index_id); let serialized_metadata = match crate::db::ser_repo_meta(&metadata) { Ok(s) => s, Err(e) => { error!(did = %did, err = %e, "failed to serialize repo metadata"); continue; } }; batch.insert(&state.db.repo_metadata, &metadata_key, serialized_metadata); if let Err(e) = lifecycle_counts.transition(&mut batch, &did, GaugeState::Pending) { error!(did = %did, err = %e, "failed to stage lifecycle transition"); continue; }
count_deltas .add_gauge_diff(&GaugeState::Resync(Some(kind)), &GaugeState::Pending); // move from resync back into pending batch.remove(&state.db.resync, key.clone()); batch.remove(&state.db.pending, old_pending); batch.insert(&state.db.pending, new_pending, key.clone()); batch.insert(&state.db.repo_metadata, &metadata_key, serialized_metadata); transitions += 1; } } Ok(_) => {9 unmodified lines if transitions == 0 { continue; }
let reservation = state.db.stage_count_deltas(&mut batch, &count_deltas); let lifecycle_reservation = lifecycle_counts.stage(&mut batch); if let Err(e) = batch.commit() { error!(err = %e, "failed to commit batch"); db::check_poisoned(&e); drop(lifecycle_reservation); continue; } state.db.apply_count_deltas(&count_deltas); drop(reservation); state.db.apply_lifecycle_counts(lifecycle_reservation); state.notify_backfill(); info!(count = transitions, "queued retries"); }}205 unmodified linesExpand all20620720820921021021121221321421521621721821922022122222322422522621222722822923023123223323423523623723823924024124224324424524624724821421521621721821922022122222324925025125222522622722822923025325425525614 unmodified lines27127227327425325425525625725825926026126227527627727827928028128228328428528628728828929029129229329429529629729829930030123 unmodified lines3253263273282942952962973293303312993323333343353363 unmodified lines34034134234331131231331431531631734434534634734832334935035135235335411 unmodified lines366367368369343370371372373374375376377378379345346380381382383348349350351352384385386387388389390391251 unmodified linesExpand all643644645646612647648649650318 unmodified linesExpand all9699709719729399739749759429769779789799805 unmodified lines986987988989956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991990991992993994995996997998114 unmodified linesExpand all111311141115111611171118111911201121112211231124111811191120112511261127112811291122112311301131112511261132113311341135113615 unmodified lines115211531154115511561157115811591160116111621163116411651166116711681169117011711150115111521153115411551156115711581172117311741175117611771178117911801181118211831184118511861187118811891160116111901191119211931194317 unmodified linesExpand all151215131514151514871516151715181490151915201521152215234 unmodified lines152815291530153115031504150515061507150815091510153215331534151415351536153715381539154015411519152015211522152315421543154415451546154715481549155015511552155315541555155615571558155915601561156219 unmodified lines205 unmodified linesExpand all strategy: BackfillStrategy,) -> Result<(), BackfillError> { let db = &state.db;
match process_did(state, &http, did, verify_signatures, strategy).await { match process_did( state, &http, did, pending_key.clone(), verify_signatures, strategy, ) .await { Ok(Some(_repo_state)) => { let applied = tokio::task::spawn_blocking({ let state = state.clone(); let did = did.clone(); let pending_key = pending_key.clone(); move || { let db = &state.db; let did_key = keys::repo_key(did); let did_key = keys::repo_key(&did); let mut batch = db.inner.batch(); let mut lifecycle_counts = db.lifecycle_counts(); let applied = lifecycle_counts.transition_pending_key( &mut batch, &did, pending_key.as_ref(), GaugeState::Synced, )?; batch.remove(&db.pending, pending_key.clone()); if applied { batch.remove(&db.resync, &did_key); } let lifecycle_reservation = lifecycle_counts.stage(&mut batch); batch.commit().into_diagnostic()?; db.apply_lifecycle_counts(lifecycle_reservation); Ok::<_, miette::Report>(applied) } }) .await .into_diagnostic()??;
// determine old gauge state // if it was error/suspended etc, we need to know which error kind it was to decrement correctly. let mut batch = db.inner.batch(); let mut count_deltas = CountDeltas::default(); // unconditionally remove from pending batch.remove(&db.pending, pending_key); // remove from resync, just in case batch.remove(&db.resync, &did_key); count_deltas.add_gauge_diff(&GaugeState::Pending, &GaugeState::Synced); let reservation = db.stage_count_deltas(&mut batch, &count_deltas); if !applied { return Ok(()); }
tokio::task::spawn_blocking(move || batch.commit().into_diagnostic()) .await .into_diagnostic()??; db.apply_count_deltas(&count_deltas); drop(reservation);
let state = state.clone(); tokio::task::spawn_blocking(move || { state .db14 unmodified lines } Ok(None) => Ok(()), Err(BackfillError::Deleted) => { warn!("orphaned pending entry, cleaning up"); let mut batch = db.inner.batch(); let mut count_deltas = CountDeltas::default(); batch.remove(&db.pending, pending_key); count_deltas.add("pending", -1); let reservation = db.stage_count_deltas(&mut batch, &count_deltas); tokio::task::spawn_blocking(move || batch.commit().into_diagnostic()) .await .into_diagnostic()??; db.apply_count_deltas(&count_deltas); drop(reservation); tokio::task::spawn_blocking({ let state = state.clone(); let did = did.clone(); let pending_key = pending_key.clone(); move || { let db = &state.db; let mut batch = db.inner.batch(); let mut lifecycle_counts = db.lifecycle_counts(); lifecycle_counts.transition_pending_key( &mut batch, &did, pending_key.as_ref(), GaugeState::Synced, )?; batch.remove(&db.pending, pending_key); let lifecycle_reservation = lifecycle_counts.stage(&mut batch); batch.commit().into_diagnostic()?; db.apply_lifecycle_counts(lifecycle_reservation); Ok::<_, miette::Report>(()) } }) .await .into_diagnostic()??; Ok(()) } Err(e) => { match &e {23 unmodified lines b.map(|b| rmp_serde::from_slice::<ResyncState>(&b).into_diagnostic()) .transpose() })?;
let (mut retry_count, prev_kind) = match existing_state { Some(ResyncState::Error { kind, retry_count, .. }) => (retry_count, Some(kind)), let mut retry_count = match existing_state { Some(ResyncState::Error { retry_count, .. }) => retry_count, Some(ResyncState::Gone { .. }) => return Ok(()), // should handle gone? original code didn't really? None => (0, None), None => 0, };
// Calculate new stats retry_count += 1;3 unmodified lines kind: error_kind, retry_count, next_retry, }; let old_gauge = prev_kind .map(|k| GaugeState::Resync(Some(k))) .unwrap_or(GaugeState::Pending); let new_gauge = GaugeState::Resync(Some(error_kind)); let mut count_deltas = CountDeltas::default(); count_deltas.add_gauge_diff(&old_gauge, &new_gauge);
let error_string = e.to_string();
tokio::task::spawn_blocking({ let state = state.clone(); let did_key = did_key.into_static(); let count_deltas = count_deltas.clone(); let did = did.clone(); let pending_key = pending_key.clone(); move || { // 3. save to resync let serialized_resync_state = rmp_serde::to_vec(&resync_state).into_diagnostic()?;11 unmodified lines None };
let mut batch = state.db.inner.batch(); batch.insert(&state.db.resync, &did_key, serialized_resync_state); let mut lifecycle_counts = state.db.lifecycle_counts(); let applied = lifecycle_counts.transition_pending_key( &mut batch, &did, pending_key.as_ref(), GaugeState::Resync(Some(error_kind)), )?; batch.remove(&state.db.pending, pending_key.clone()); if applied { batch.insert(&state.db.resync, &did_key, serialized_resync_state); if let Some(state_bytes) = serialized_repo_state { batch.insert(&state.db.repos, &did_key, state_bytes); if let Some(state_bytes) = serialized_repo_state { batch.insert(&state.db.repos, &did_key, state_bytes); } } let reservation = state.db.stage_count_deltas(&mut batch, &count_deltas); batch.commit().into_diagnostic().inspect(|_| { state.db.apply_count_deltas(&count_deltas); drop(reservation); }) let lifecycle_reservation = lifecycle_counts.stage(&mut batch); batch.commit().into_diagnostic()?; state.db.apply_lifecycle_counts(lifecycle_reservation); Ok::<_, miette::Report>(()) } }) .await .into_diagnostic()??;251 unmodified linesExpand all ) .await?;
let Some((records, state)) = result else { cleanup_discarded_repo(app_state, did).await?; return Ok(SparseBackfillResult::Discarded); };
Ok(SparseBackfillResult::Imported(SparseBackfillSuccess {318 unmodified linesExpand all }
let mut count_deltas = CountDeltas::default(); if delta != 0 { count_deltas.add("records", delta); count_deltas.add_records(delta); } if added_blocks > 0 { count_deltas.add("blocks", added_blocks); count_deltas.add_blocks(added_blocks); } let reservation = app_state.db.stage_count_deltas(&mut batch, &count_deltas); batch.commit().into_diagnostic()?; app_state.db.apply_count_deltas(&count_deltas);5 unmodified lines .into_diagnostic()? .map_err(BackfillError::from)}
async fn cleanup_discarded_repo( app_state: &Arc<AppState>, did: &Did<'static>,) -> Result<(), BackfillError> { let metadata_key = keys::repo_metadata_key(did); let metadata_bytes = app_state .db .repo_metadata .get(&metadata_key) .into_diagnostic()? .ok_or_else(|| miette::miette!("repo metadata not found for {}", did))?; let metadata = crate::db::deser_repo_meta(metadata_bytes.as_ref())?; let did_key = keys::repo_key(did); let backfill_pending_key = keys::pending_key(metadata.index_id); let app_state = app_state.clone();
tokio::task::spawn_blocking(move || { let mut batch = app_state.db.inner.batch(); let mut count_deltas = CountDeltas::default(); batch.remove(&app_state.db.repos, &did_key); batch.remove(&app_state.db.repo_metadata, &metadata_key); batch.remove(&app_state.db.pending, backfill_pending_key); count_deltas.add("repos", -1); count_deltas.add("pending", -1); let reservation = app_state.db.stage_count_deltas(&mut batch, &count_deltas); batch.commit().into_diagnostic().inspect(|_| { app_state.db.apply_count_deltas(&count_deltas); drop(reservation); }) }) .await .into_diagnostic()??;
Ok(())}
async fn process_did( app_state: &Arc<AppState>, http: &reqwest::Client, did: &Did<'static>, pending_key: Slice, verify_signatures: bool, strategy: BackfillStrategy,) -> Result<Option<RepoState<'static>>, BackfillError> { debug!("starting...");114 unmodified linesExpand all Err(XrpcError::Xrpc(e)) => { if matches!(e, GetRepoError::RepoNotFound(_)) { warn!("repo not found, deleting"); let mut batch = db.inner.batch(); let mut lifecycle_counts = db.lifecycle_counts(); let applied = lifecycle_counts.transition_pending_key( &mut batch, did, pending_key.as_ref(), GaugeState::Synced, )?; batch.remove(&db.pending, pending_key.clone()); let mut count_deltas = CountDeltas::default(); if let Err(e) = crate::ops::delete_repo(&mut batch, db, did, &state) { error!(err = %e, "failed to wipe repo during backfill"); if applied { if let Err(e) = crate::ops::delete_repo(&mut batch, db, did, &state) { error!(err = %e, "failed to wipe repo during backfill"); } } count_deltas.add("pending", -1); let reservation = db.stage_count_deltas(&mut batch, &count_deltas); let lifecycle_reservation = lifecycle_counts.stage(&mut batch); batch.commit().into_diagnostic()?; db.apply_count_deltas(&count_deltas); drop(reservation); db.apply_lifecycle_counts(lifecycle_reservation); // return None so did_task skips sending BackfillFinished (nothing to drain for a deleted repo) return Ok(None); }
15 unmodified lines }; let resync_bytes = rmp_serde::to_vec(&resync_state).into_diagnostic()?;
let app_state_clone = app_state.clone(); let did = did.clone(); let pending_key = pending_key.clone(); tokio::task::spawn_blocking(move || { let db = &app_state_clone.db; let mut batch = db.inner.batch(); let mut lifecycle_counts = db.lifecycle_counts(); let applied = lifecycle_counts.transition_pending_key( &mut batch, &did, pending_key.as_ref(), GaugeState::Resync(None), )?; batch.remove(&db.pending, pending_key.clone()); if applied { Db::update_repo_state( &mut batch, app_state .db .update_repo_state_async(did, move |state, (key, batch)| { state.active = false; state.status = status; batch.insert(&app_state_clone.db.resync, key, resync_bytes); Ok((true, ())) }) .await?; &db.repos, &did, move |state, (key, batch)| { state.active = false; state.status = status; batch.insert(&db.resync, key, resync_bytes); Ok((true, ())) }, )?; } let lifecycle_reservation = lifecycle_counts.stage(&mut batch); batch.commit().into_diagnostic()?; db.apply_lifecycle_counts(lifecycle_reservation); Ok::<_, miette::Report>(()) }) .await .into_diagnostic()??;
// return success so wrapper stops retrying return Ok(Some(previous_state)); return Ok(None); }
Err(e).into_diagnostic()? }317 unmodified linesExpand all }
let mut count_deltas = CountDeltas::default(); if delta != 0 { count_deltas.add("records", delta); count_deltas.add_records(delta); } if added_blocks > 0 { count_deltas.add("blocks", added_blocks); count_deltas.add_blocks(added_blocks); } let reservation = app_state.db.stage_count_deltas(&mut batch, &count_deltas); batch.commit().into_diagnostic()?; app_state.db.apply_count_deltas(&count_deltas);4 unmodified lines .await .into_diagnostic()?? };
let metadata_key = keys::repo_metadata_key(did); let metadata_bytes = db .repo_metadata .get(&metadata_key) .into_diagnostic()? .ok_or_else(|| miette::miette!("repo metadata not found for {}", did))?; let metadata = crate::db::deser_repo_meta(metadata_bytes.as_ref())?;
let Some(count) = result else { // signal mode: no signal-matching records found, clean up the optimistically-added repo let did_key = keys::repo_key(did); let backfill_pending_key = keys::pending_key(metadata.index_id); let metadata_key = keys::repo_metadata_key(did); let backfill_pending_key = pending_key.clone(); let app_state = app_state.clone(); let did = did.clone(); tokio::task::spawn_blocking(move || { let mut batch = app_state.db.inner.batch(); let mut count_deltas = CountDeltas::default(); batch.remove(&app_state.db.repos, &did_key); batch.remove(&app_state.db.repo_metadata, &metadata_key); batch.remove(&app_state.db.pending, backfill_pending_key); count_deltas.add("repos", -1); count_deltas.add("pending", -1); let mut lifecycle_counts = app_state.db.lifecycle_counts(); let applied = lifecycle_counts.transition_pending_key( &mut batch, &did, backfill_pending_key.as_ref(), GaugeState::Synced, )?; batch.remove(&app_state.db.pending, backfill_pending_key.clone()); if applied { batch.remove(&app_state.db.repos, &did_key); batch.remove(&app_state.db.repo_metadata, &metadata_key); count_deltas.add_repos(-1); } let lifecycle_reservation = lifecycle_counts.stage(&mut batch); let reservation = app_state.db.stage_count_deltas(&mut batch, &count_deltas); batch.commit().into_diagnostic().inspect(|_| { app_state.db.apply_lifecycle_counts(lifecycle_reservation); app_state.db.apply_count_deltas(&count_deltas); drop(reservation); }) })19 unmodified lines143 unmodified linesExpand all14414514614714814915015115220 unmodified lines17317417517617617717817817918018118218318418518618718818919019119219319419519619738 unmodified lines143 unmodified linesExpand all let mut rng: SmallRng = rand::make_rng(); let mut batch = app_state.db.inner.batch(); let mut surviving = Vec::new(); let mut count_deltas = CountDeltas::default(); let mut lifecycle_counts = app_state.db.lifecycle_counts(); for guard in guards { let did_key = keys::repo_key(&guard); let metadata_key = keys::repo_metadata_key(&guard); if app_state20 unmodified lines ); // clear any stale retry entry, this DID is confirmed and being enqueued batch.remove(&app_state.db.crawler, keys::crawler_retry_key(&guard)); trace!(did = %*guard, "enqueuing repo"); count_deltas.add("repos", 1); count_deltas.add_repos(1); #[cfg(feature = "indexer")] count_deltas.add("pending", 1); lifecycle_counts.transition( &mut batch, &guard, crate::types::GaugeState::Pending, )?; surviving.push(guard); } if let Some(cursor) = cursor_update { batch.insert(&app_state.db.cursors, cursor.key, cursor.value); } let lifecycle_reservation = lifecycle_counts.stage(&mut batch); let reservation = app_state.db.stage_count_deltas(&mut batch, &count_deltas); // todo: repo state overwrites here are acceptable? batch.commit().into_diagnostic()?; app_state.db.apply_lifecycle_counts(lifecycle_reservation); app_state.db.apply_count_deltas(&count_deltas); drop(reservation); Ok(surviving) }),38 unmodified lines1123423 unmodified lines2829303133343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798032333435195 unmodified linesExpand alluse crate::types::{GaugeState, RepoStatus, ResyncState};use fjall::{Keyspace, OwnedWriteBatch};use jacquard_common::IntoStatic;use jacquard_common::types::string::Did;use miette::{IntoDiagnostic, Result, WrapErr};23 unmodified lines } else { Ok(None) } }
pub(crate) async fn update_repo_state_async<F, T>( &self, did: &Did<'_>, f: F, ) -> Result<Option<(RepoState<'static>, T)>> where F: FnOnce(&mut RepoState, (&[u8], &mut fjall::OwnedWriteBatch)) -> Result<(bool, T)> + Send + 'static, T: Send + 'static, { let mut batch = self.inner.batch(); let repos = self.repos.clone(); let did = did.clone().into_static();
tokio::task::spawn_blocking(move || { let Some((state, t)) = Self::update_repo_state(&mut batch, &repos, &did, f)? else { return Ok(None); }; batch.commit().into_diagnostic()?; Ok(Some((state, t))) }) .await .into_diagnostic()? }
pub(crate) fn repo_gauge_state( repo_state: &RepoState, resync_bytes: Option<&[u8]>, ) -> GaugeState { match repo_state.status { RepoStatus::Synced => GaugeState::Synced, RepoStatus::Error(_) | RepoStatus::Deactivated | RepoStatus::Takendown | RepoStatus::Suspended | RepoStatus::Deleted | RepoStatus::Desynchronized | RepoStatus::Throttled => resync_bytes .and_then(|b| rmp_serde::from_slice::<ResyncState>(b).ok()) .and_then(|s| match s { ResyncState::Error { kind, .. } => Some(GaugeState::Resync(Some(kind))), _ => None, }) .unwrap_or(GaugeState::Resync(None)), } }}
pub fn set_record_count( batch: &mut OwnedWriteBatch,195 unmodified linesExpand all123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319use std::collections::BTreeMap;use std::sync::MutexGuard;
use fjall::OwnedWriteBatch;use jacquard_common::types::string::Did;use miette::{IntoDiagnostic, Result, WrapErr};
use crate::types::{GaugeState, ResyncErrorKind, ResyncState};
use super::{CountDeltaReservation, CountDeltas, Db, deser_repo_meta, keys};
pub(crate) struct LifecycleCountBatch<'a> { db: &'a Db, _lock: MutexGuard<'a, ()>, deltas: CountDeltas, staged: BTreeMap<Vec<u8>, GaugeState>,}
pub(crate) struct LifecycleCountReservation<'a> { _lock: MutexGuard<'a, ()>, _count_reservation: Option<CountDeltaReservation>, deltas: CountDeltas,}
impl Db { pub(crate) fn lifecycle_counts(&self) -> LifecycleCountBatch<'_> { LifecycleCountBatch { db: self, _lock: self .lifecycle_count_lock .lock() .expect("lifecycle count lock poisoned"), deltas: CountDeltas::default(), staged: BTreeMap::new(), } }
pub(crate) fn apply_lifecycle_counts(&self, reservation: LifecycleCountReservation<'_>) { self.apply_count_deltas(&reservation.deltas); drop(reservation); }}
impl<'a> LifecycleCountBatch<'a> { pub(crate) fn transition( &mut self, batch: &mut OwnedWriteBatch, did: &Did<'_>, new: GaugeState, ) -> Result<bool> { let did_key = keys::repo_key(did); let (old, had_projection) = self.current_gauge(did, &did_key)?;
if old != new { self.deltas.add_gauge_diff(&old, &new); self.stage_membership(batch, did, &did_key, new); return Ok(true); }
if !had_projection { self.stage_membership(batch, did, &did_key, new); }
Ok(false) }
pub(crate) fn transition_pending_key( &mut self, batch: &mut OwnedWriteBatch, did: &Did<'_>, pending_key: &[u8], new: GaugeState, ) -> Result<bool> { if !self.pending_key_matches(did, pending_key)? { return Ok(false); }
self.transition(batch, did, new)?; Ok(true) }
pub(crate) fn stage(self, batch: &mut OwnedWriteBatch) -> LifecycleCountReservation<'a> { let count_reservation = self.db.stage_count_deltas(batch, &self.deltas); LifecycleCountReservation { _lock: self._lock, _count_reservation: count_reservation, deltas: self.deltas, } }
fn current_gauge(&self, did: &Did<'_>, did_key: &[u8]) -> Result<(GaugeState, bool)> { if let Some(staged) = self.staged.get(did_key) { return Ok((*staged, true)); }
let projection_key = keys::count_gauge_key(did); if let Some(bytes) = self.db.counts.get(&projection_key).into_diagnostic()? { return decode_gauge(bytes.as_ref()).map(|gauge| (gauge, true)); }
if self.db.repos.get(did_key).into_diagnostic()?.is_some() { miette::bail!("missing lifecycle count projection for existing repo {did}"); }
Ok((GaugeState::Synced, false)) }
fn pending_key_matches(&self, did: &Did<'_>, pending_key: &[u8]) -> Result<bool> { let Some(current) = current_pending_key(self.db, did)? else { return Ok(false); };
Ok(current.as_slice() == pending_key && self.db.pending.get(current).into_diagnostic()?.is_some()) }
fn stage_membership( &mut self, batch: &mut OwnedWriteBatch, did: &Did<'_>, did_key: &[u8], new: GaugeState, ) { let projection_key = keys::count_gauge_key(did); batch.insert(&self.db.counts, projection_key, encode_gauge(new)); self.staged.insert(did_key.to_vec(), new); }}
fn current_pending_key(db: &Db, did: &Did<'_>) -> Result<Option<[u8; 8]>> { db.repo_metadata .get(keys::repo_metadata_key(did)) .into_diagnostic()? .map(|bytes| { deser_repo_meta(bytes.as_ref()) .map(|metadata| keys::pending_key(metadata.index_id)) .wrap_err_with(|| format!("invalid repo metadata for {did}")) }) .transpose()}
pub(super) fn gauge_from_resync(bytes: &[u8]) -> GaugeState { rmp_serde::from_slice::<ResyncState>(bytes) .map(|state| match state { ResyncState::Error { kind, .. } => GaugeState::Resync(Some(kind)), ResyncState::Gone { .. } => GaugeState::Resync(None), }) .unwrap_or(GaugeState::Resync(None))}
pub(super) fn encode_gauge(gauge: GaugeState) -> [u8; 1] { match gauge { GaugeState::Synced => [0], GaugeState::Pending => [1], GaugeState::Resync(None) => [2], GaugeState::Resync(Some(ResyncErrorKind::Ratelimited)) => [3], GaugeState::Resync(Some(ResyncErrorKind::Transport)) => [4], GaugeState::Resync(Some(ResyncErrorKind::Generic)) => [5], }}
fn decode_gauge(bytes: &[u8]) -> Result<GaugeState> { match bytes { [0] => Ok(GaugeState::Synced), [1] => Ok(GaugeState::Pending), [2] => Ok(GaugeState::Resync(None)), [3] => Ok(GaugeState::Resync(Some(ResyncErrorKind::Ratelimited))), [4] => Ok(GaugeState::Resync(Some(ResyncErrorKind::Transport))), [5] => Ok(GaugeState::Resync(Some(ResyncErrorKind::Generic))), _ => miette::bail!("invalid lifecycle gauge projection bytes: {bytes:?}"), }}
#[cfg(test)]mod tests { use super::*; use crate::config::Config; use crate::db::{ser_repo_meta, ser_repo_state}; use crate::types::{RepoMetadata, RepoState}; use jacquard_common::IntoStatic;
fn test_config(path: &std::path::Path) -> Config { Config { database_path: path.to_path_buf(), ..Default::default() } }
fn test_did() -> Result<Did<'static>> { Did::new("did:plc:yk4q3id7id6p5z3bypvshc64") .into_diagnostic() .map(IntoStatic::into_static) }
fn enqueue_pending(db: &Db, did: &Did<'_>, index_id: u64) -> Result<[u8; 8]> { let did_key = keys::repo_key(did); let metadata_key = keys::repo_metadata_key(did); let pending_key = keys::pending_key(index_id); let state = RepoState::backfilling(); let metadata = RepoMetadata::backfilling(index_id);
let mut batch = db.inner.batch(); batch.insert(&db.repos, &did_key, ser_repo_state(&state)?); batch.insert(&db.repo_metadata, metadata_key, ser_repo_meta(&metadata)?); batch.insert(&db.pending, pending_key, did_key); let mut lifecycle = db.lifecycle_counts(); lifecycle.transition(&mut batch, did, GaugeState::Pending)?; let lifecycle_reservation = lifecycle.stage(&mut batch); batch.commit().into_diagnostic()?; db.apply_lifecycle_counts(lifecycle_reservation); Ok(pending_key) }
fn complete_pending(db: &Db, did: &Did<'_>, pending_key: [u8; 8]) -> Result<bool> { let mut batch = db.inner.batch(); let mut lifecycle = db.lifecycle_counts(); let applied = lifecycle.transition_pending_key( &mut batch, did, pending_key.as_slice(), GaugeState::Synced, )?; if applied { batch.remove(&db.pending, pending_key); } let lifecycle_reservation = lifecycle.stage(&mut batch); batch.commit().into_diagnostic()?; db.apply_lifecycle_counts(lifecycle_reservation); Ok(applied) }
#[test] fn duplicate_pending_completion_is_idempotent() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let db = Db::open(&test_config(tmp.path()))?; let did = test_did()?; let pending_key = enqueue_pending(&db, &did, 1)?;
assert_eq!(db.get_count_sync("pending"), 1); assert!(complete_pending(&db, &did, pending_key)?); assert_eq!(db.get_count_sync("pending"), 0);
assert!(!complete_pending(&db, &did, pending_key)?); assert_eq!(db.get_count_sync("pending"), 0);
Ok(()) }
#[test] fn stale_pending_key_cannot_complete_current_pending_state() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let db = Db::open(&test_config(tmp.path()))?; let did = test_did()?; let stale_pending = enqueue_pending(&db, &did, 1)?; let current_pending = keys::pending_key(2); let did_key = keys::repo_key(&did); let metadata_key = keys::repo_metadata_key(&did);
let mut metadata = RepoMetadata::backfilling(2); metadata.tracked = true; let mut batch = db.inner.batch(); batch.remove(&db.pending, stale_pending); batch.insert(&db.pending, current_pending, &did_key); batch.insert(&db.repo_metadata, metadata_key, ser_repo_meta(&metadata)?); let mut lifecycle = db.lifecycle_counts(); lifecycle.transition(&mut batch, &did, GaugeState::Pending)?; let lifecycle_reservation = lifecycle.stage(&mut batch); batch.commit().into_diagnostic()?; db.apply_lifecycle_counts(lifecycle_reservation);
assert_eq!(db.get_count_sync("pending"), 1); assert!(!complete_pending(&db, &did, stale_pending)?); assert_eq!(db.get_count_sync("pending"), 1); assert!(db.pending.get(current_pending).into_diagnostic()?.is_some());
Ok(()) }
#[test] fn lifecycle_delta_replays_after_crash_before_apply() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let cfg = test_config(tmp.path()); let did = test_did()?;
{ let db = Db::open(&cfg)?; let pending_key = enqueue_pending(&db, &did, 1)?; db.persist()?;
let mut batch = db.inner.batch(); let mut lifecycle = db.lifecycle_counts(); let applied = lifecycle.transition_pending_key( &mut batch, &did, pending_key.as_slice(), GaugeState::Synced, )?; assert!(applied); batch.remove(&db.pending, pending_key); let lifecycle_reservation = lifecycle.stage(&mut batch); batch.commit().into_diagnostic()?; db.persist()?; drop(lifecycle_reservation); }
{ let db = Db::open(&cfg)?; assert_eq!(db.get_count_sync("pending"), 0); assert!( db.pending .get(keys::pending_key(1)) .into_diagnostic()? .is_none() ); }
Ok(()) }}