use std::fmt;
use super::{DurabilityError, DurableStore};
pub const HEAD_CHECKPOINT_EVERY: u64 = 256;
pub(crate) const READ_BATCH_SIZE: usize = 1_024;
const HEAD_CHECKPOINT_NAMESPACE: &str = "liminal:channel-head";
#[must_use]
pub fn head_checkpoint_key(stream_key: &str) -> String {
format!("{HEAD_CHECKPOINT_NAMESPACE}:{stream_key}")
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HeadCheckpointRefusal {
Missing,
BeyondLog {
covered: u64,
},
WrongSequence {
covered: u64,
found: u64,
},
}
impl fmt::Display for HeadCheckpointRefusal {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Missing => formatter.write_str("no head is recorded and the log holds rows"),
Self::BeyondLog { covered } => write!(
formatter,
"the recorded head {covered} claims rows the log does not hold"
),
Self::WrongSequence { covered, found } => write!(
formatter,
"the recorded head {covered} names a row whose sequence is {found}"
),
}
}
}
struct Walk {
rows: u64,
highest: Option<u64>,
}
impl Walk {
fn next_sequence(&self, floor: u64) -> Result<u64, DurabilityError> {
let Some(highest) = self.highest else {
return Ok(floor);
};
let next = highest.checked_add(1).ok_or_else(|| {
DurabilityError::ConfigError(
"sequence number overflow after channel recovery".to_owned(),
)
})?;
Ok(next.max(floor))
}
}
pub async fn recover_partition_head(
store: &dyn DurableStore,
stream_key: &str,
) -> Result<u64, DurabilityError> {
let key = head_checkpoint_key(stream_key);
let recorded = store.read_value(&key).await?.filter(|covered| *covered > 0);
let Some(covered) = recorded else {
let walk = walk_from(store, stream_key, 0).await?;
if walk.rows == 0 {
return Ok(0);
}
let next = walk.next_sequence(0)?;
warn_refused(stream_key, HeadCheckpointRefusal::Missing, next);
advance_head(store, stream_key, next).await?;
return Ok(next);
};
match verify(store, stream_key, covered).await? {
None => {
let tail = walk_from(store, stream_key, covered).await?;
let next = tail.next_sequence(covered)?;
if tail.rows >= HEAD_CHECKPOINT_EVERY {
advance_head(store, stream_key, next).await?;
}
Ok(next)
}
Some(refusal) => {
let walk = walk_from(store, stream_key, 0).await?;
let next = walk.next_sequence(0)?;
warn_refused(stream_key, refusal, next);
replace_refused_head(store, stream_key, covered, next).await?;
Ok(next)
}
}
}
pub async fn advance_head(
store: &dyn DurableStore,
stream_key: &str,
next: u64,
) -> Result<(), DurabilityError> {
if next == 0 {
return Ok(());
}
let key = head_checkpoint_key(stream_key);
let mut stored = store.read_value(&key).await?.unwrap_or(0);
while stored < next {
match store.cas(&key, stored, next).await {
Ok(()) => return Ok(()),
Err(DurabilityError::CursorRegression {
stored: observed, ..
}) if observed != stored => stored = observed,
Err(error) => return Err(error),
}
}
Ok(())
}
async fn replace_refused_head(
store: &dyn DurableStore,
stream_key: &str,
refused: u64,
next: u64,
) -> Result<(), DurabilityError> {
if next == 0 {
tracing::warn!(
stream = stream_key,
refused,
"a refused durable channel head cannot be replaced while its log is empty; the \
next start refuses it again and walks the empty log"
);
return Ok(());
}
store
.cas(&head_checkpoint_key(stream_key), refused, next)
.await
}
async fn verify(
store: &dyn DurableStore,
stream_key: &str,
covered: u64,
) -> Result<Option<HeadCheckpointRefusal>, DurabilityError> {
let last = covered.checked_sub(1).ok_or_else(|| {
DurabilityError::ConfigError("a recorded head of zero has no row to verify".to_owned())
})?;
Ok(match store.read_at(stream_key, last).await? {
None => Some(HeadCheckpointRefusal::BeyondLog { covered }),
Some(entry) if entry.sequence != last => Some(HeadCheckpointRefusal::WrongSequence {
covered,
found: entry.sequence,
}),
Some(_) => None,
})
}
async fn walk_from(
store: &dyn DurableStore,
stream_key: &str,
offset: u64,
) -> Result<Walk, DurabilityError> {
let mut position = offset;
let mut walk = Walk {
rows: 0,
highest: None,
};
loop {
let batch = store
.read_from(stream_key, position, READ_BATCH_SIZE)
.await?;
let batch_len = batch.len();
for stored in &batch {
if stored.sequence < offset {
return Err(DurabilityError::ConfigError(format!(
"stream {stream_key} answered sequence {} to a read from {offset}",
stored.sequence
)));
}
walk.highest = Some(
walk.highest
.map_or(stored.sequence, |s| s.max(stored.sequence)),
);
}
let read = u64::try_from(batch_len).map_err(|error| {
DurabilityError::ConfigError(format!(
"channel recovery entry count cannot fit u64: {error}"
))
})?;
walk.rows = walk.rows.checked_add(read).ok_or_else(|| {
DurabilityError::ConfigError("channel recovery row count overflow".to_owned())
})?;
position = position.checked_add(read).ok_or_else(|| {
DurabilityError::ConfigError("channel recovery read offset overflow".to_owned())
})?;
if batch_len < READ_BATCH_SIZE {
return Ok(walk);
}
}
}
fn warn_refused(stream_key: &str, refusal: HeadCheckpointRefusal, next: u64) {
tracing::warn!(
stream = stream_key,
reason = %refusal,
rebuilt_head = next,
"durable channel head record refused; the whole log was walked and the head it \
found is recorded again"
);
}
#[cfg(test)]
mod tests;