liminal/durability/channel/

storage.rs

//! Durable channel admission retains and reconciles uncertain writes, and recovers its
//! partition sequence from a definite append conflict, before reusing a sequence.
use std::fmt;
use std::sync::Arc;

use super::batch::PendingBatch;
use super::{MessageEnvelope, UnresolvedAppend};
use crate::durability::{
    DurabilityConfig, DurabilityError, DurabilityMode, DurableStore, HEAD_CHECKPOINT_EVERY,
    advance_head,
};

/// Caller-provided function used to route envelopes to partitions.
#[derive(Clone)]
pub struct PartitionKey {
    function: Arc<dyn Fn(&MessageEnvelope) -> u64 + Send + Sync + 'static>,
}

impl PartitionKey {
    /// Wraps a partition key closure.
    #[must_use]
    pub fn new<F>(function: F) -> Self
    where
        F: Fn(&MessageEnvelope) -> u64 + Send + Sync + 'static,
    {
        Self {
            function: Arc::new(function),
        }
    }

    fn apply(&self, envelope: &MessageEnvelope) -> u64 {
        (self.function)(envelope)
    }
}

impl fmt::Debug for PartitionKey {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("PartitionKey(<closure>)")
    }
}

/// Haematite-backed durable channel state.
#[derive(Clone)]
pub struct DurableChannel {
    pub(super) channel_id: String,
    partition_count: usize,
    partition_key: Option<PartitionKey>,
    next_sequences: Vec<u64>,
    pub(super) pending: Option<UnresolvedAppend>,
    pub(super) pending_batch: Option<PendingBatch>,
    pub(super) store: Arc<dyn DurableStore>,
}

impl DurableChannel {
    /// Creates a new durable channel with all partition sequences initialized to zero.
    ///
    /// # Errors
    ///
    /// Returns [`DurabilityError::ConfigError`] when `partition_count` is zero.
    pub fn new(
        channel_id: impl Into<String>,
        partition_count: usize,
        store: Arc<dyn DurableStore>,
    ) -> Result<Self, DurabilityError> {
        Self::from_parts(channel_id.into(), partition_count, store, None, None)
    }

    /// Creates a durable channel with a caller-provided partition key function.
    ///
    /// # Errors
    ///
    /// Returns [`DurabilityError::ConfigError`] when `partition_count` is zero.
    pub fn with_partition_key<F>(
        channel_id: impl Into<String>,
        partition_count: usize,
        store: Arc<dyn DurableStore>,
        partition_key: F,
    ) -> Result<Self, DurabilityError>
    where
        F: Fn(&MessageEnvelope) -> u64 + Send + Sync + 'static,
    {
        Self::from_parts(
            channel_id.into(),
            partition_count,
            store,
            Some(PartitionKey::new(partition_key)),
            None,
        )
    }

    /// Creates a durable channel from explicit recovered next sequence counters.
    ///
    /// # Errors
    ///
    /// Returns [`DurabilityError::ConfigError`] when `partition_count` is zero
    /// or the recovered sequence vector length differs from `partition_count`.
    pub fn from_recovered_sequences(
        channel_id: impl Into<String>,
        partition_count: usize,
        store: Arc<dyn DurableStore>,
        next_sequences: Vec<u64>,
    ) -> Result<Self, DurabilityError> {
        Self::from_parts(
            channel_id.into(),
            partition_count,
            store,
            None,
            Some(next_sequences),
        )
    }

    /// Creates a durable channel from validated durability configuration.
    ///
    /// # Errors
    ///
    /// Returns [`DurabilityError::ConfigError`] when `config` is ephemeral.
    pub fn from_config(
        channel_id: impl Into<String>,
        config: DurabilityConfig,
        store: Arc<dyn DurableStore>,
    ) -> Result<Self, DurabilityError> {
        if config.mode() == DurabilityMode::Ephemeral {
            return Err(DurabilityError::ConfigError(
                "durable channel requires a durable durability mode".to_owned(),
            ));
        }
        Self::new(channel_id, config.partition_count(), store)
    }

    /// Returns the durable channel identifier.
    #[must_use]
    pub fn channel_id(&self) -> &str {
        &self.channel_id
    }

    /// Returns the number of partitions owned by this channel.
    #[must_use]
    pub const fn partition_count(&self) -> usize {
        self.partition_count
    }

    /// Returns the next expected sequence for a partition, if it exists.
    #[must_use]
    pub fn next_expected_sequence(&self, partition_index: usize) -> Option<u64> {
        self.next_sequences.get(partition_index).copied()
    }

    /// Returns all next expected sequence counters.
    #[must_use]
    pub fn next_sequences(&self) -> &[u64] {
        &self.next_sequences
    }

    /// Computes the partition index for an envelope without touching storage.
    #[must_use]
    pub fn partition_for(&self, envelope: &MessageEnvelope) -> usize {
        route_partition(self.partition_count, self.partition_key.as_ref(), envelope)
    }

    /// Formats the haematite stream key for a channel partition.
    #[must_use]
    pub fn stream_key_for(&self, partition_index: usize) -> String {
        format!("{}:{partition_index}", self.channel_id)
    }

    /// Persists an envelope, retaining exact evidence before the cancellable append.
    ///
    /// # Errors
    ///
    /// Returns serialization and store errors. An earlier unresolved append is
    /// reconciled by one read before this one is sent, and an unreadable one
    /// refuses this publish by name. A definite sequence conflict reads the
    /// current tail and retries this refused append once. An uncertain append
    /// error keeps its evidence and is returned without resubmitting its payload.
    pub async fn publish(&mut self, envelope: &MessageEnvelope) -> Result<u64, DurabilityError> {
        self.publish_recorded(envelope)
            .await
            .map(|append| append.sequence())
    }

    /// Carries the exact append evidence into live delivery without encoding again.
    pub(crate) async fn publish_recorded(
        &mut self,
        envelope: &MessageEnvelope,
    ) -> Result<UnresolvedAppend, DurabilityError> {
        let payload = envelope.serialize()?;
        self.resolve_pending_batch().await?;
        self.resolve_pending().await?;
        let partition = self.partition_for(envelope);
        let stream_key = self.stream_key_for(partition);
        match self
            .append_recorded(partition, &stream_key, payload, envelope)
            .await
        {
            Err(DurabilityError::SequenceConflict { expected, actual }) => {
                self.refresh_sequence(partition, &stream_key, expected, actual)
                    .await?;
                match self
                    .append_recorded(partition, &stream_key, envelope.serialize()?, envelope)
                    .await
                {
                    Err(DurabilityError::SequenceConflict { expected, actual }) => {
                        self.refresh_sequence(partition, &stream_key, expected, actual)
                            .await?;
                        Err(DurabilityError::SequenceConflict { expected, actual })
                    }
                    outcome => outcome,
                }
            }
            outcome => outcome,
        }
    }

    /// Records the pending append, then sends it. Only a definite outcome
    /// clears the record: a success, or a conflict that stored nothing.
    async fn append_recorded(
        &mut self,
        partition: usize,
        stream_key: &str,
        payload: Vec<u8>,
        envelope: &MessageEnvelope,
    ) -> Result<UnresolvedAppend, DurabilityError> {
        let expected = self.sequence_for_append(partition)?;
        let next = expected.checked_add(1).ok_or_else(|| {
            DurabilityError::ConfigError(format!("channel {} sequence overflow", self.channel_id))
        })?;
        self.pending = Some(UnresolvedAppend::new(
            self.channel_id.clone(),
            partition,
            expected,
            payload.clone(),
            envelope.idempotency_key.clone(),
        ));
        match self.store.append(stream_key, payload, expected).await {
            Ok(assigned) if assigned == expected => {
                let append = self.pending.take().ok_or_else(|| {
                    DurabilityError::ConfigError(format!(
                        "channel {} lost its committed append evidence",
                        self.channel_id
                    ))
                })?;
                self.advance_partition(partition, stream_key, next).await?;
                Ok(append)
            }
            Ok(assigned) => self.append_unresolved(format!(
                "store answered sequence {assigned} for reserved sequence {expected}"
            )),
            Err(conflict @ DurabilityError::SequenceConflict { .. }) => {
                self.pending = None;
                Err(conflict)
            }
            Err(original) => self.append_unresolved(original.to_string()),
        }
    }

    fn append_unresolved(&mut self, cause: String) -> Result<UnresolvedAppend, DurabilityError> {
        let pending = self.pending.as_mut().ok_or_else(|| {
            DurabilityError::ConfigError(format!(
                "channel {} lost its pending append evidence",
                self.channel_id
            ))
        })?;
        pending.set_cause(cause.clone());
        Err(DurabilityError::AppendUnresolved {
            stream_key: format!("{}:{}", pending.channel(), pending.partition()),
            sequence: pending.sequence(),
            cause,
        })
    }

    /// Moves a partition past a row the store holds, and records the head once
    /// every [`HEAD_CHECKPOINT_EVERY`] appends.
    pub(super) async fn advance_partition(
        &mut self,
        partition: usize,
        stream_key: &str,
        next: u64,
    ) -> Result<(), DurabilityError> {
        self.set_next_sequence(partition, next)?;
        if next % HEAD_CHECKPOINT_EVERY == 0 {
            // The row is already durable, so a failed head record must not fail the
            // publish: the old record stays valid, the next start reads a longer tail,
            // and the next count records the head again.
            if let Err(error) = advance_head(self.store.as_ref(), stream_key, next).await {
                tracing::warn!(
                    stream = %stream_key,
                    head = next,
                    %error,
                    "durable channel head could not be recorded; the next start reads the \
                     rows since the previous record"
                );
            }
        }
        Ok(())
    }

    pub(super) async fn refresh_sequence(
        &mut self,
        partition_index: usize,
        stream_key: &str,
        expected: u64,
        mut next_sequence: u64,
    ) -> Result<(), DurabilityError> {
        // The conflict reports the store's next slot. Read only rows that could
        // have arrived after that observation, never the already settled prefix.
        if let Some(entry) = self.store.read_at(stream_key, next_sequence).await? {
            if entry.sequence < next_sequence {
                return Err(DurabilityError::ConfigError(format!(
                    "stream {stream_key} read returned sequence {} before {next_sequence}",
                    entry.sequence
                )));
            }
            next_sequence = entry.sequence.checked_add(1).ok_or_else(|| {
                DurabilityError::ConfigError(format!(
                    "stream {stream_key} sequence overflow during conflict recovery"
                ))
            })?;
        }
        tracing::warn!(
            channel = %self.channel_id,
            expected,
            recorded = next_sequence,
            "durable channel sequence conflict"
        );
        self.set_next_sequence(partition_index, next_sequence)
    }

    /// Flushes the backing store so every appended message is durably persisted.
    ///
    /// # Errors
    ///
    /// Propagates any [`DurabilityError`] returned by [`DurableStore::flush`].
    pub async fn flush_store(&self) -> Result<(), DurabilityError> {
        self.store.flush().await
    }

    fn from_parts(
        channel_id: String,
        partition_count: usize,
        store: Arc<dyn DurableStore>,
        partition_key: Option<PartitionKey>,
        next_sequences: Option<Vec<u64>>,
    ) -> Result<Self, DurabilityError> {
        validate_partition_count(partition_count)?;
        let next_sequences = next_sequences.unwrap_or_else(|| vec![0; partition_count]);
        if next_sequences.len() != partition_count {
            return Err(DurabilityError::ConfigError(
                "recovered sequence count must match partition_count".to_owned(),
            ));
        }
        Ok(Self {
            channel_id,
            partition_count,
            partition_key,
            next_sequences,
            pending: None,
            pending_batch: None,
            store,
        })
    }

    pub(super) fn sequence_for_append(
        &self,
        partition_index: usize,
    ) -> Result<u64, DurabilityError> {
        self.next_sequences
            .get(partition_index)
            .copied()
            .ok_or_else(|| {
                DurabilityError::ConfigError("partition sequence state missing".to_owned())
            })
    }

    fn set_next_sequence(
        &mut self,
        partition_index: usize,
        next_sequence: u64,
    ) -> Result<(), DurabilityError> {
        let Some(sequence) = self.next_sequences.get_mut(partition_index) else {
            return Err(DurabilityError::ConfigError(
                "partition sequence state missing".to_owned(),
            ));
        };
        *sequence = next_sequence;
        Ok(())
    }
}

impl fmt::Debug for DurableChannel {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("DurableChannel")
            .field("channel_id", &self.channel_id)
            .field("partition_count", &self.partition_count)
            .field("partition_key_configured", &self.partition_key.is_some())
            .field("next_sequences", &self.next_sequences)
            .field("pending", &self.pending)
            .field("pending_batch", &self.pending_batch)
            .field("store", &self.store)
            .finish()
    }
}

/// Ephemeral channel state with no durable store reference.
#[derive(Clone)]
pub struct EphemeralChannel {
    channel_id: String,
    partition_count: usize,
    partition_key: Option<PartitionKey>,
}

impl EphemeralChannel {
    /// Creates an ephemeral channel without requiring any durable store.
    ///
    /// # Errors
    ///
    /// Returns [`DurabilityError::ConfigError`] when `partition_count` is zero.
    pub fn new(
        channel_id: impl Into<String>,
        partition_count: usize,
    ) -> Result<Self, DurabilityError> {
        Self::from_parts(channel_id.into(), partition_count, None)
    }

    /// Creates an ephemeral channel with a caller-provided partition key function.
    ///
    /// # Errors
    ///
    /// Returns [`DurabilityError::ConfigError`] when `partition_count` is zero.
    pub fn with_partition_key<F>(
        channel_id: impl Into<String>,
        partition_count: usize,
        partition_key: F,
    ) -> Result<Self, DurabilityError>
    where
        F: Fn(&MessageEnvelope) -> u64 + Send + Sync + 'static,
    {
        Self::from_parts(
            channel_id.into(),
            partition_count,
            Some(PartitionKey::new(partition_key)),
        )
    }

    /// Creates an ephemeral channel from validated durability configuration.
    ///
    /// # Errors
    ///
    /// Returns [`DurabilityError::ConfigError`] when `config` is not ephemeral.
    pub fn from_config(
        channel_id: impl Into<String>,
        config: DurabilityConfig,
    ) -> Result<Self, DurabilityError> {
        if config.mode() != DurabilityMode::Ephemeral {
            return Err(DurabilityError::ConfigError(
                "ephemeral channel requires Ephemeral durability mode".to_owned(),
            ));
        }
        Self::new(channel_id, config.partition_count())
    }

    /// Returns the channel identifier.
    #[must_use]
    pub fn channel_id(&self) -> &str {
        &self.channel_id
    }

    /// Returns the number of configured partitions.
    #[must_use]
    pub const fn partition_count(&self) -> usize {
        self.partition_count
    }

    /// Computes the partition index for an envelope without touching storage.
    #[must_use]
    pub fn partition_for(&self, envelope: &MessageEnvelope) -> usize {
        route_partition(self.partition_count, self.partition_key.as_ref(), envelope)
    }

    fn from_parts(
        channel_id: String,
        partition_count: usize,
        partition_key: Option<PartitionKey>,
    ) -> Result<Self, DurabilityError> {
        validate_partition_count(partition_count)?;
        Ok(Self {
            channel_id,
            partition_count,
            partition_key,
        })
    }
}

impl fmt::Debug for EphemeralChannel {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("EphemeralChannel")
            .field("channel_id", &self.channel_id)
            .field("partition_count", &self.partition_count)
            .field("partition_key_configured", &self.partition_key.is_some())
            .finish()
    }
}

fn route_partition(
    partition_count: usize,
    partition_key: Option<&PartitionKey>,
    envelope: &MessageEnvelope,
) -> usize {
    if partition_count == 1 {
        return 0;
    }

    let Some(partition_key) = partition_key else {
        return 0;
    };

    let Ok(partition_count_u64) = u64::try_from(partition_count) else {
        return 0;
    };
    let routed = partition_key.apply(envelope) % partition_count_u64;
    usize::try_from(routed).unwrap_or_else(|_| partition_count.saturating_sub(1))
}

fn validate_partition_count(partition_count: usize) -> Result<(), DurabilityError> {
    if partition_count == 0 {
        return Err(DurabilityError::ConfigError(
            "partition_count must be at least 1".to_owned(),
        ));
    }
    Ok(())
}