use super::floor::Cutoff;
use crate::types::{Epoch, Height, Round, View};
use commonware_codec::varint::UInt;
use commonware_storage::{
Context,
metadata::{self, Metadata},
};
use commonware_utils::sequence::U64;
const LATEST_KEY: U64 = U64::new(0xFF);
const CUTOFF_EPOCH_KEY: U64 = U64::new(0xFE);
const CUTOFF_VIEW_KEY: U64 = U64::new(0xFD);
#[derive(Clone, Copy)]
enum State {
Unprocessed,
Processed(Height),
}
impl State {
const fn new(processed_height: Option<Height>) -> Self {
match processed_height {
Some(height) => Self::Processed(height),
None => Self::Unprocessed,
}
}
const fn processed_height(self) -> Option<Height> {
match self {
Self::Unprocessed => None,
Self::Processed(height) => Some(height),
}
}
const fn next_height(self) -> Height {
match self {
Self::Unprocessed => Height::zero(),
Self::Processed(height) => height.next(),
}
}
const fn acknowledge(&mut self, height: Height) {
*self = Self::Processed(height);
}
}
pub(super) struct Stream<E: Context> {
metadata: Metadata<E, U64, UInt<u64>>,
state: State,
synced: bool,
}
impl<E: Context> Stream<E> {
pub(super) async fn new(context: E, application_metadata_partition: &str) -> Self {
let metadata: Metadata<E, U64, UInt<u64>> = Metadata::init(
context,
metadata::Config {
partition: application_metadata_partition.to_string(),
codec_config: (),
},
)
.await
.expect("failed to initialize application metadata");
let state = State::new(
metadata
.get(&LATEST_KEY)
.map(|height| Height::new(height.0)),
);
Self {
metadata,
state,
synced: true,
}
}
pub(super) const fn is_synced(&self) -> bool {
self.synced
}
pub(super) const fn processed_height(&self) -> Option<Height> {
self.state.processed_height()
}
pub(super) const fn next_height(&self) -> Height {
self.state.next_height()
}
pub(super) fn cutoff(&self) -> Cutoff {
match (
self.metadata.get(&CUTOFF_EPOCH_KEY),
self.metadata.get(&CUTOFF_VIEW_KEY),
) {
(Some(epoch), Some(view)) => {
Cutoff::Round(Round::new(Epoch::new(epoch.0), View::new(view.0)))
}
(Some(epoch), None) => Cutoff::Epoch(Epoch::new(epoch.0)),
(None, _) => Cutoff::zero(),
}
}
pub(super) fn acknowledge(&mut self, height: Height) {
self.state.acknowledge(height);
self.metadata.put(LATEST_KEY, UInt(height.get()));
self.synced = false;
}
pub(super) fn set_cutoff(&mut self, cutoff: Cutoff) {
match cutoff {
Cutoff::Round(round) => {
self.metadata
.put(CUTOFF_EPOCH_KEY, UInt(round.epoch().get()));
self.metadata.put(CUTOFF_VIEW_KEY, UInt(round.view().get()));
}
Cutoff::Epoch(epoch) => {
self.metadata.put(CUTOFF_EPOCH_KEY, UInt(epoch.get()));
self.metadata.remove(&CUTOFF_VIEW_KEY);
}
}
self.synced = false;
}
pub(super) async fn sync(mut self) -> Result<Self, metadata::Error> {
if self.synced {
return Ok(self);
}
self.metadata = self.metadata.sync().await?;
self.synced = true;
Ok(self)
}
}
#[cfg(test)]
mod tests {
use super::*;
use commonware_runtime::{Runner as _, Supervisor as _, deterministic};
const PARTITION: &str = "stream-metadata";
async fn seed_processed_height(context: deterministic::Context, height: Height) {
let mut metadata: Metadata<deterministic::Context, U64, Height> = Metadata::init(
context,
metadata::Config {
partition: PARTITION.to_string(),
codec_config: (),
},
)
.await
.expect("metadata initialization must succeed");
metadata.put(LATEST_KEY, height);
metadata.sync().await.expect("metadata sync must succeed");
}
#[test]
fn test_round_floor_round_trips_beside_processed_height() {
let runner = deterministic::Runner::default();
runner.start(|context| async move {
seed_processed_height(context.child("seed"), Height::new(7)).await;
let mut stream = Stream::new(context.child("first"), PARTITION).await;
assert_eq!(stream.processed_height(), Some(Height::new(7)));
assert_eq!(stream.cutoff(), Cutoff::zero());
let cutoff = Cutoff::Round(Round::new(Epoch::new(2), View::new(40)));
stream.set_cutoff(cutoff);
stream.acknowledge(Height::new(8));
let stream = stream.sync().await.expect("stream sync must succeed");
drop(stream);
let stream = Stream::new(context.child("second"), PARTITION).await;
assert_eq!(stream.processed_height(), Some(Height::new(8)));
assert_eq!(stream.cutoff(), cutoff);
});
}
#[test]
fn test_round_floor_kind_survives_reopen_across_transitions() {
let runner = deterministic::Runner::default();
runner.start(|context| async move {
let at_round = Cutoff::Round(Round::new(Epoch::new(2), View::new(40)));
let mut stream = Stream::new(context.child("first"), PARTITION).await;
stream.set_cutoff(at_round);
let stream = stream.sync().await.expect("stream sync must succeed");
drop(stream);
let mut stream = Stream::new(context.child("second"), PARTITION).await;
assert_eq!(stream.cutoff(), at_round);
let whole_epoch = Cutoff::Epoch(Epoch::new(2));
stream.set_cutoff(whole_epoch);
assert_eq!(stream.cutoff(), whole_epoch);
let stream = stream.sync().await.expect("stream sync must succeed");
drop(stream);
let mut stream = Stream::new(context.child("third"), PARTITION).await;
assert_eq!(stream.cutoff(), whole_epoch);
let later = Cutoff::Round(Round::new(Epoch::new(3), View::new(5)));
stream.set_cutoff(later);
let stream = stream.sync().await.expect("stream sync must succeed");
drop(stream);
let stream = Stream::new(context.child("fourth"), PARTITION).await;
assert_eq!(stream.cutoff(), later);
});
}
}