liminal/durability/

head_checkpoint.rs

//! Recorded partition heads: where a durable channel's restart resumes reading its log.
//!
//! A durable channel partition's only folded state is its next append sequence. Without
//! a record of it, every construction walks the partition's whole log to find the head,
//! so a restart costs the channel's entire history. The head is recorded here, beside
//! the log, as the sequence it covers; a restart verifies the record against the log and
//! reads only the rows appended after it.
//!
//! The record is written by count: once every [`HEAD_CHECKPOINT_EVERY`] appends, and
//! at recovery whenever the tail read past the record reached that count or the record
//! was refused. Nothing here runs on a clock.
//!
//! A missing record is the ordinary first checkpoint, recorded after a walk at INFO.
//! A record claiming rows the log does not hold is refused by name at WARN, then
//! rebuilt from a full walk. A clean stop records every appended stream's final head.

use std::fmt;

use super::{DurabilityError, DurableStore};

/// How many rows a partition's log may grow past its recorded head before the head is
/// recorded again. This bounds the rows a restart reads per partition.
pub const HEAD_CHECKPOINT_EVERY: u64 = 256;

/// How many rows one recovery read asks the store for.
pub(crate) const READ_BATCH_SIZE: usize = 1_024;

/// Namespace of the recorded partition heads in the store's value keyspace.
const HEAD_CHECKPOINT_NAMESPACE: &str = "liminal:channel-head";

/// The value key under which a partition stream's recorded head is kept.
#[must_use]
pub fn head_checkpoint_key(stream_key: &str) -> String {
    format!("{HEAD_CHECKPOINT_NAMESPACE}:{stream_key}")
}

/// Why a recorded head was not trusted.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HeadCheckpointRefusal {
    /// The record claims `covered` rows and the row before that position is absent.
    BeyondLog {
        /// The next sequence the record claims.
        covered: u64,
    },
    /// The row the record names carries a different sequence than its position.
    WrongSequence {
        /// The next sequence the record claims.
        covered: u64,
        /// The sequence found at the position the record names.
        found: u64,
    },
}

impl fmt::Display for HeadCheckpointRefusal {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            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}"
            ),
        }
    }
}

/// What one walk of a partition log found.
struct Walk {
    /// Rows the walk read.
    rows: u64,
    /// The highest sequence the walk read, if it read any row.
    highest: Option<u64>,
}

impl Walk {
    /// The next append sequence after this walk, never below `floor`.
    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))
    }
}

/// Recovers one partition's next append sequence, reading only the rows appended after
/// its recorded head when that head verifies against the log.
///
/// # Errors
///
/// Propagates store read and write errors, and returns
/// [`DurabilityError::ConfigError`] on offset or sequence overflow, or when a row the
/// walk reads lies before the position it was asked to read from.
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)?;
        advance_head(store, stream_key, next).await?;
        tracing::info!(
            stream = stream_key,
            rows = walk.rows,
            "durable channel head recorded for a stream under its first checkpoint"
        );
        return Ok(next);
    };
    match verify(store, stream_key, covered).await? {
        None => {
            if store.read_stream_next_seq(stream_key).await? == Some(covered) {
                return Ok(covered);
            }
            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)
        }
    }
}

/// Records `next` as the partition's head when it is past the recorded one.
///
/// A concurrent writer that recorded a head in between is read back from the refusal,
/// and the record is retried only while it is still behind `next`, so the loop ends
/// once either this write lands or another writer has recorded at least as far.
///
/// # Errors
///
/// Propagates store read and write errors other than a lost race.
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(())
}

/// Replaces a refused record with the head a full walk found.
///
/// An empty log cannot be recorded, because the value store keeps no physical zero; the
/// refused record then stays, and the next start refuses it again by name.
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
}

/// Checks a recorded head against the log: the row just before `covered` must exist and
/// carry the sequence of its position. Answers the refusal, or `None` when it holds.
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,
    })
}

/// Reads a partition log from `offset` to its end, keeping only the row count and the
/// highest sequence.
///
/// Each read starts past the highest sequence the last one answered, never at a count
/// of rows: a pruned log answers a read below its first retained sequence from that
/// sequence, so a count-based position would ask for the same retained page again
/// until it had counted past every pruned row.
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 < position {
                return Err(DurabilityError::ConfigError(format!(
                    "stream {stream_key} answered sequence {} to a read from {position}",
                    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())
        })?;
        let Some(last) = batch.last() else {
            return Ok(walk);
        };
        position = last.sequence.checked_add(1).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;