use crate::{
journal::Error,
metadata::{Config as MetadataConfig, Metadata},
Context,
};
use commonware_utils::sequence::VecU64;
const PRUNING_BOUNDARY_KEY: u64 = 1;
const CLEAR_TARGET_KEY: u64 = 2;
const RECOVERY_WATERMARK_KEY: u64 = 3;
pub(super) struct Checkpoint<E: Context> {
metadata: Metadata<E, u64, VecU64>,
}
impl<E: Context> Checkpoint<E> {
pub(super) async fn open(context: E, partition_prefix: &str) -> Result<Self, Error> {
let metadata = Metadata::<_, u64, VecU64>::init(
context,
MetadataConfig {
partition: format!("{partition_prefix}-metadata"),
codec_config: (),
},
)
.await?;
Ok(Self { metadata })
}
fn get(&self, key: u64) -> Option<u64> {
self.metadata.get(&key).copied().map(u64::from)
}
pub(super) fn watermark(&self) -> Option<u64> {
self.get(RECOVERY_WATERMARK_KEY)
}
pub(super) fn boundary_hint(&self) -> Option<u64> {
self.get(PRUNING_BOUNDARY_KEY)
}
pub(super) fn clear_target(&self) -> Option<u64> {
self.get(CLEAR_TARGET_KEY)
}
pub(super) async fn persist(
&mut self,
items_per_blob: u64,
boundary: u64,
watermark: u64,
) -> Result<(), Error> {
if boundary.is_multiple_of(items_per_blob) {
if self.boundary_hint().is_some() {
self.metadata.remove(&PRUNING_BOUNDARY_KEY);
}
} else if self.boundary_hint() != Some(boundary) {
self.metadata.put(PRUNING_BOUNDARY_KEY, boundary.into());
}
if self.watermark() != Some(watermark) {
self.metadata.put(RECOVERY_WATERMARK_KEY, watermark.into());
}
self.sync().await
}
pub(super) fn lower_watermark(&mut self, limit: u64) -> bool {
match self.watermark() {
Some(current) if current > limit => {
self.metadata.put(RECOVERY_WATERMARK_KEY, limit.into());
true
}
_ => false,
}
}
pub(super) async fn stage_clear(&mut self, target: u64) -> Result<(), Error> {
self.lower_watermark(target);
self.metadata.put(CLEAR_TARGET_KEY, target.into());
self.sync().await
}
pub(super) async fn finish_clear(
&mut self,
items_per_blob: u64,
target: u64,
) -> Result<(), Error> {
self.metadata.remove(&CLEAR_TARGET_KEY);
self.persist(items_per_blob, target, target).await
}
pub(super) async fn sync(&mut self) -> Result<(), Error> {
self.metadata.sync().await?;
Ok(())
}
pub(super) async fn destroy(self) -> Result<(), Error> {
self.metadata.destroy().await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use commonware_macros::test_traced;
use commonware_runtime::{deterministic, Runner as _, Supervisor as _};
impl<E: Context> Checkpoint<E> {
pub(crate) fn set_watermark(&mut self, watermark: Option<u64>) {
match watermark {
Some(watermark) => {
self.metadata.put(RECOVERY_WATERMARK_KEY, watermark.into());
}
None => {
self.metadata.remove(&RECOVERY_WATERMARK_KEY);
}
}
}
pub(crate) fn set_boundary_hint(&mut self, boundary: u64) {
self.metadata.put(PRUNING_BOUNDARY_KEY, boundary.into());
}
pub(crate) fn set_clear_target(&mut self, target: u64) {
self.metadata.put(CLEAR_TARGET_KEY, target.into());
}
pub(crate) fn clear(&mut self) {
self.metadata.clear();
}
}
#[test_traced]
fn test_lower_watermark_only_lowers() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
let mut checkpoint = Checkpoint::open(context, "lw").await.unwrap();
assert!(!checkpoint.lower_watermark(5));
assert_eq!(checkpoint.watermark(), None);
checkpoint.persist(10, 0, 8).await.unwrap();
assert!(!checkpoint.lower_watermark(8));
assert!(!checkpoint.lower_watermark(12));
assert_eq!(checkpoint.watermark(), Some(8));
assert!(checkpoint.lower_watermark(3));
assert_eq!(checkpoint.watermark(), Some(3));
});
}
#[test_traced]
fn test_persist_round_trips_across_reopen() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
{
let mut checkpoint = Checkpoint::open(context.child("a"), "rt").await.unwrap();
checkpoint.persist(10, 13, 25).await.unwrap();
}
let checkpoint = Checkpoint::open(context.child("b"), "rt").await.unwrap();
assert_eq!(checkpoint.boundary_hint(), Some(13));
assert_eq!(checkpoint.watermark(), Some(25));
assert_eq!(checkpoint.clear_target(), None);
});
}
#[test_traced]
fn test_persist_boundary_hint_tracks_alignment() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
let mut checkpoint = Checkpoint::open(context, "align").await.unwrap();
checkpoint.persist(10, 13, 0).await.unwrap();
assert_eq!(checkpoint.boundary_hint(), Some(13));
checkpoint.persist(10, 20, 0).await.unwrap();
assert_eq!(checkpoint.boundary_hint(), None);
});
}
#[test_traced]
fn test_clear_lifecycle_survives_crash() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
{
let mut checkpoint = Checkpoint::open(context.child("a"), "clear").await.unwrap();
checkpoint.persist(10, 0, 30).await.unwrap();
checkpoint.stage_clear(20).await.unwrap();
assert_eq!(checkpoint.clear_target(), Some(20));
assert_eq!(checkpoint.watermark(), Some(20));
}
{
let mut checkpoint = Checkpoint::open(context.child("b"), "clear").await.unwrap();
assert_eq!(checkpoint.clear_target(), Some(20));
checkpoint.finish_clear(10, 20).await.unwrap();
}
let checkpoint = Checkpoint::open(context.child("c"), "clear").await.unwrap();
assert_eq!(checkpoint.clear_target(), None);
assert_eq!(checkpoint.watermark(), Some(20));
assert_eq!(checkpoint.boundary_hint(), None);
});
}
#[test_traced]
fn test_stage_clear_does_not_raise_watermark() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
let mut checkpoint = Checkpoint::open(context, "sc").await.unwrap();
checkpoint.persist(10, 0, 5).await.unwrap();
checkpoint.stage_clear(9).await.unwrap();
assert_eq!(checkpoint.clear_target(), Some(9));
assert_eq!(checkpoint.watermark(), Some(5));
});
}
}