use std::collections::BTreeMap;
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};
/// one txn's moves between the lifecycle gauges. each repo's projection is a read, a change
/// and a write back with no lock, since only the repo's owner shard moves it, and the gauge
/// totals are deltas, which add up in any order
pub(super) struct LifecycleCountBatch<'a> {
db: &'a Db,
deltas: CountDeltas,
staged: BTreeMap<Vec<u8>, GaugeState>,
}
pub(super) struct LifecycleCountReservation {
_count_reservation: Option<CountDeltaReservation>,
deltas: CountDeltas,
}
impl Db {
pub(super) fn lifecycle_counts(&self) -> LifecycleCountBatch<'_> {
LifecycleCountBatch {
db: self,
deltas: CountDeltas::default(),
staged: BTreeMap::new(),
}
}
pub(super) fn apply_lifecycle_counts(&self, reservation: LifecycleCountReservation) {
self.apply_count_deltas(&reservation.deltas);
drop(reservation);
}
}
impl<'a> LifecycleCountBatch<'a> {
pub(super) 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)
}
/// takes a repo out of the gauges for good, for a repo being erased rather
/// than moved to another state
pub(super) fn forget(&mut self, batch: &mut OwnedWriteBatch, did: &Did) -> Result<()> {
let did_key = keys::repo_key(did);
let (old, _) = self.current_gauge(did, &did_key)?;
self.deltas.add_gauge_diff(&old, &GaugeState::Synced);
batch.remove(&self.db.counts, keys::count_gauge_key(did));
self.staged.insert(did_key, GaugeState::Synced);
Ok(())
}
pub(super) fn transition_pending_key(
&mut self,
batch: &mut OwnedWriteBatch,
did: &Did,
pending_key: &[u8],
new: GaugeState,
) -> Result<bool> {
if !names_pending_key(self.db, did, pending_key)? {
return Ok(false);
}
self.transition(batch, did, new)?;
Ok(true)
}
pub(super) fn stage(self, batch: &mut OwnedWriteBatch) -> LifecycleCountReservation {
let count_reservation = self.db.stage_count_deltas(batch, &self.deltas);
LifecycleCountReservation {
_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() {
let metadata_key = keys::repo_metadata_key(did);
let index_id = self
.db
.repo_metadata
.get(&metadata_key)
.into_diagnostic()?
.and_then(|bytes| deser_repo_meta(bytes.as_ref()).ok())
.map(|meta| meta.index_id);
let is_pending = index_id
.map(|index_id| {
self.db
.indexer
.is_pending(&keys::pending_key(index_id))
.into_diagnostic()
})
.transpose()?
.unwrap_or(false);
let gauge = if is_pending {
GaugeState::Pending
} else if let Some(resync_bytes) =
self.db.indexer.resync.get(did_key).into_diagnostic()?
{
gauge_from_resync(&resync_bytes)
} else {
GaugeState::Synced
};
return Ok((gauge, false));
}
Ok((GaugeState::Synced, false))
}
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);
}
}
/// whether `pending_key` is queued and is the key the repo's metadata names, so not an orphan
/// left queued next to it
pub(super) fn names_pending_key(db: &Db, did: &Did, pending_key: &[u8]) -> Result<bool> {
let Some(current) = current_pending_key(db, did)? else {
return Ok(false);
};
Ok(current.as_slice() == pending_key && db.indexer.is_pending(¤t).into_diagnostic()?)
}
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> {
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)?);
db.indexer
.stage_pending_insert(&mut batch, did, &pending_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 {
db.indexer
.stage_pending_remove(&mut batch, did, &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 metadata_key = keys::repo_metadata_key(&did);
let mut metadata = RepoMetadata::backfilling(2);
metadata.tracked = true;
let mut batch = db.inner.batch();
db.indexer
.stage_pending_remove(&mut batch, &did, &stale_pending);
db.indexer
.stage_pending_insert(&mut batch, &did, ¤t_pending);
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.indexer.is_pending(¤t_pending).into_diagnostic()?);
db.indexer.assert_pending_dids_match();
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);
db.indexer
.stage_pending_remove(&mut batch, &did, &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.indexer
.is_pending(&keys::pending_key(1))
.into_diagnostic()?
);
db.indexer.assert_pending_dids_match();
}
Ok(())
}
}