liminal-server 0.20.1

//! Conditional rows of one logical operation, published at one durability boundary.

use std::collections::BTreeMap;

use haematite::tree::Hash;

use liminal::durability::settle::apply_routed_batch_settled;
use liminal::durability::{DurabilityError, DurableStore, RoutedBatchItem, StoredEntry};

use super::restore_checkpoint::{
    CheckpointError, CheckpointLiveState, CheckpointPositions, RestoreCheckpoint,
};
use super::restore_checkpoint_locator::{
    CheckpointBodyKey, CheckpointLocator, CheckpointLocatorError, CheckpointSource,
    advance_prefix_digest, checkpoint_conversation_route,
};
use super::routed_stream::{RoutedStreamError, RoutedStreamRefusal, prepare_page, read_head};
use super::routed_stream_codec;

/// The number of covered operations between checkpoints.
pub(super) const CHECKPOINT_EVERY: u64 = 256;

/// Owns all staged rows until a single atomic commit succeeds.
#[derive(Debug)]
pub(super) struct OperationBatch {
    route_key: String,
    items: Vec<RoutedBatchItem>,
    keys: BTreeMap<Vec<u8>, Option<usize>>,
    streams: BTreeMap<String, StagedStream>,
    staging_refused: bool,
    prefix_rows: Vec<StagedPrefixRow>,
    /// The cut decided at this batch's checkpoint (DUR-018): its receipts
    /// are staged items; the cut itself runs only after the batch commits.
    prune: Option<super::prune::PrunePlan>,
}

/// Failure to stage a checkpoint without publishing any part of its operation.
#[derive(Debug, thiserror::Error)]
pub(super) enum CheckpointStageError {
    #[error(transparent)]
    Store(#[from] DurabilityError),
    #[error(transparent)]
    Checkpoint(#[from] CheckpointError),
    #[error(transparent)]
    Locator(#[from] CheckpointLocatorError),
    #[error("operation batch staging was already refused")]
    BatchRefused,
    #[error("checkpoint belongs to another operation route")]
    RouteMismatch,
    #[error("checkpoint positions regress behind the durable locator")]
    PositionRegression,
    #[error("content-addressed checkpoint body differs from existing bytes")]
    BodyMismatch,
}

#[derive(Debug)]
struct StagedPrefixRow {
    source: CheckpointSource,
    sequence: u64,
    item: usize,
}

#[derive(Debug)]
struct StagedStream {
    head_key: Vec<u8>,
    expected: Option<Hash>,
    next_sequence: u64,
}

impl OperationBatch {
    /// Starts an operation on the conversation's owned route.
    #[must_use]
    pub(super) fn new(route_key: &str) -> Self {
        Self {
            route_key: route_key.to_owned(),
            items: Vec::new(),
            keys: BTreeMap::new(),
            streams: BTreeMap::new(),
            staging_refused: false,
            prefix_rows: Vec::new(),
            prune: None,
        }
    }

    /// Stages a checkpoint's prune receipts and remembers the plan for after
    /// the commit (DUR-018 R3: the receipt is durable before any row goes).
    ///
    /// # Errors
    /// A receipt that does not prove its own cut, or a refused key.
    pub(super) fn stage_prune(
        &mut self,
        plan: super::prune::PrunePlan,
    ) -> Result<(), super::log::OperationLogError> {
        super::prune::stage(self, &plan)?;
        self.prune = Some(plan);
        Ok(())
    }

    /// Takes the plan staged by [`Self::stage_prune`], for the cut that
    /// follows a successful commit.
    pub(super) const fn take_prune(&mut self) -> Option<super::prune::PrunePlan> {
        self.prune.take()
    }

    /// Prevents submission after an enclosing encoder or preflight refused.
    pub(super) const fn refuse(&mut self) {
        self.staging_refused = true;
    }

    /// Whether this operation has any pending logical rows.
    pub(super) fn is_empty(&self) -> Result<bool, DurabilityError> {
        if self.staging_refused {
            return Err(DurabilityError::ConfigError(
                "operation batch staging was refused".into(),
            ));
        }
        Ok(self.items.is_empty() && self.streams.is_empty())
    }

    /// Stages one row with its original content-hash witness.
    ///
    /// # Errors
    /// Refuses an empty or repeated key before it can enter the batch.
    pub(super) fn stage(&mut self, item: RoutedBatchItem) -> Result<(), DurabilityError> {
        if self.staging_refused {
            return Err(DurabilityError::ConfigError(
                "operation batch staging was refused".into(),
            ));
        }
        if item.key.is_empty() {
            self.staging_refused = true;
            return Err(DurabilityError::ConfigError(
                "operation batch key is empty".into(),
            ));
        }
        if self.keys.contains_key(&item.key) {
            self.staging_refused = true;
            return Err(DurabilityError::ConfigError(
                "operation batch repeats a key".into(),
            ));
        }
        self.keys.insert(item.key.clone(), Some(self.items.len()));
        self.items.push(item);
        Ok(())
    }

    /// Stages contiguous rows while retaining the original head witness.
    ///
    /// The caller establishes creation or migration authority separately.
    /// Absence here is only a conditional-write witness.
    ///
    /// # Errors
    /// Refuses stale sequences, colliding keys and encoding or store failures.
    pub(super) async fn stage_stream(
        &mut self,
        store: &dyn DurableStore,
        stream: &str,
        payloads: Vec<Vec<u8>>,
        expected_sequence: u64,
        timestamp: u64,
    ) -> Result<u64, RoutedStreamError> {
        if self.staging_refused {
            return Err(RoutedStreamRefusal::BatchRefused.into());
        }
        let result = self
            .prepare_stream(store, stream, payloads, expected_sequence, timestamp)
            .await;
        if result.is_err() {
            self.staging_refused = true;
        }
        result
    }

    /// Records this operation's canonical sources in their actual staging order.
    ///
    /// # Errors
    /// Refuses any stage or reservation failure before submission.
    pub(super) async fn stage_source_stream(
        &mut self,
        store: &dyn DurableStore,
        stream: &str,
        source: CheckpointSource,
        entry: StoredEntry,
    ) -> Result<u64, RoutedStreamError> {
        let result = async {
            self.prefix_rows
                .try_reserve_exact(1)
                .map_err(RoutedStreamRefusal::Allocation)?;
            let first = self.items.len();
            let sequence = entry.sequence;
            let next = self
                .stage_stream(
                    store,
                    stream,
                    vec![entry.payload],
                    sequence,
                    entry.timestamp,
                )
                .await?;
            if self.items.len().checked_sub(first) != Some(1) {
                return Err(RoutedStreamRefusal::BatchRefused.into());
            }
            self.prefix_rows.push(StagedPrefixRow {
                source,
                sequence,
                item: first,
            });
            Ok(next)
        }
        .await;
        if result.is_err() {
            self.refuse();
        }
        result
    }

    /// Advances a validated prior witness over this operation alone.
    ///
    /// # Errors
    /// Refuses invalid staged framing without reading or copying stored history.
    pub(super) fn checkpoint_prefix(
        &mut self,
        prior: [u8; 32],
    ) -> Result<[u8; 32], RoutedStreamError> {
        if self.staging_refused {
            return Err(RoutedStreamRefusal::BatchRefused.into());
        }
        let result = (|| {
            let mut digest = prior;
            for row in &self.prefix_rows {
                let item = self
                    .items
                    .get(row.item)
                    .ok_or(RoutedStreamRefusal::BatchRefused)?;
                let (_, payload) =
                    routed_stream_codec::decode_row_payload(&item.value, row.sequence)?;
                digest = advance_prefix_digest(digest, row.source, row.sequence, payload)
                    .map_err(RoutedStreamRefusal::CheckpointLocator)?;
            }
            Ok(digest)
        })();
        if result.is_err() {
            self.refuse();
        }
        result
    }

    async fn prepare_stream(
        &mut self,
        store: &dyn DurableStore,
        stream: &str,
        payloads: Vec<Vec<u8>>,
        expected_sequence: u64,
        timestamp: u64,
    ) -> Result<u64, RoutedStreamError> {
        if payloads.is_empty() {
            return Err(RoutedStreamRefusal::EmptyAppend.into());
        }
        let head_key = routed_stream_codec::head_key(stream)?;
        let (actual, original_hash, is_new) = if let Some(pending) = self.streams.get(stream) {
            (pending.next_sequence, pending.expected, false)
        } else {
            if self.keys.contains_key(&head_key) {
                return Err(RoutedStreamRefusal::KeyCollision.into());
            }
            let (actual, original_hash) = match read_head(store, &self.route_key, stream).await? {
                Some(head) => (head.next_sequence, Some(head.expected)),
                // A routed stream with no head begins at the native head: the
                // rows below it are flat (a pre-pruning conversation, or a
                // migrated one), and the first routed commit continues them.
                None => (store.read_stream_next_seq(stream).await?.unwrap_or(0), None),
            };
            (actual, original_hash, true)
        };
        if actual != expected_sequence {
            return Err(RoutedStreamRefusal::HeadMismatch {
                expected: expected_sequence,
                actual,
            }
            .into());
        }
        let count =
            u64::try_from(payloads.len()).map_err(|_| RoutedStreamRefusal::SequenceOverflow)?;
        let next_sequence = actual
            .checked_add(count)
            .ok_or(RoutedStreamRefusal::SequenceOverflow)?;
        let mut prepared = Vec::new();
        prepared
            .try_reserve_exact(payloads.len())
            .map_err(RoutedStreamRefusal::Allocation)?;
        for (index, payload) in payloads.into_iter().enumerate() {
            let index = u64::try_from(index).map_err(|_| RoutedStreamRefusal::SequenceOverflow)?;
            let sequence = actual
                .checked_add(index)
                .ok_or(RoutedStreamRefusal::SequenceOverflow)?;
            let key = routed_stream_codec::row_key(stream, sequence)?;
            if self.keys.contains_key(&key) {
                return Err(RoutedStreamRefusal::KeyCollision.into());
            }
            let value = routed_stream_codec::encode_row(&StoredEntry {
                payload,
                sequence,
                timestamp,
            })?;
            prepared.push(RoutedBatchItem {
                key,
                expected: None,
                value,
                ttl: None,
            });
        }
        self.items
            .try_reserve_exact(prepared.len())
            .map_err(RoutedStreamRefusal::Allocation)?;
        if is_new {
            if self.keys.contains_key(&head_key) {
                return Err(RoutedStreamRefusal::KeyCollision.into());
            }
            self.keys.insert(head_key.clone(), None);
        }
        for item in prepared {
            self.keys.insert(item.key.clone(), Some(self.items.len()));
            self.items.push(item);
        }
        self.streams.insert(
            stream.to_owned(),
            StagedStream {
                head_key,
                expected: original_hash,
                next_sequence,
            },
        );
        Ok(next_sequence)
    }

    /// Establishes a stream boundary from the caller's validated replay or owner.
    ///
    /// # Errors
    /// Refuses a different existing head or a colliding key and poisons the batch.
    pub(super) async fn stage_bootstrap_stream(
        &mut self,
        store: &dyn DurableStore,
        stream: &str,
        next_sequence: u64,
    ) -> Result<(), RoutedStreamError> {
        if self.staging_refused {
            return Err(RoutedStreamRefusal::BatchRefused.into());
        }
        let result = self
            .prepare_bootstrap_stream(store, stream, next_sequence)
            .await;
        if result.is_err() {
            self.staging_refused = true;
        }
        result
    }

    async fn prepare_bootstrap_stream(
        &mut self,
        store: &dyn DurableStore,
        stream: &str,
        next_sequence: u64,
    ) -> Result<(), RoutedStreamError> {
        if let Some(pending) = self.streams.get(stream) {
            if pending.next_sequence != next_sequence {
                return Err(RoutedStreamRefusal::HeadMismatch {
                    expected: next_sequence,
                    actual: pending.next_sequence,
                }
                .into());
            }
            return Ok(());
        }
        let key = routed_stream_codec::head_key(stream)?;
        if self.keys.contains_key(&key) {
            return Err(RoutedStreamRefusal::KeyCollision.into());
        }
        let previous = read_head(store, &self.route_key, stream).await?;
        if let Some(previous) = previous.as_ref() {
            if previous.next_sequence != next_sequence {
                return Err(RoutedStreamRefusal::HeadMismatch {
                    expected: next_sequence,
                    actual: previous.next_sequence,
                }
                .into());
            }
        }
        self.keys.insert(key.clone(), None);
        self.streams.insert(
            stream.to_owned(),
            StagedStream {
                head_key: key,
                expected: previous.map(|head| head.expected),
                next_sequence,
            },
        );
        Ok(())
    }

    /// Reads the candidate stream without publishing any pending row.
    ///
    /// Absence remains explicit when the stream has no staged or durable head.
    /// It does not establish creation or recovery authority.
    ///
    /// # Errors
    /// Any read refusal prevents submission of the entire pending operation.
    pub(super) async fn read_stream_page(
        &mut self,
        store: &dyn DurableStore,
        stream: &str,
        offset: u64,
        limit: usize,
    ) -> Result<Option<Vec<StoredEntry>>, RoutedStreamError> {
        if self.staging_refused {
            return Err(RoutedStreamRefusal::BatchRefused.into());
        }
        let result = self.read_pending_page(store, stream, offset, limit).await;
        if result.is_err() {
            self.staging_refused = true;
        }
        result
    }

    /// Reads one authoritative source from the operation's uncommitted view.
    pub(super) async fn read_source(
        &self,
        store: &dyn DurableStore,
        stream: &str,
        sequence: u64,
    ) -> Result<Option<StoredEntry>, RoutedStreamError> {
        let rows = self.read_pending_page(store, stream, sequence, 1).await?;
        Ok(rows.and_then(|mut rows| rows.pop()))
    }

    async fn read_pending_page(
        &self,
        store: &dyn DurableStore,
        stream: &str,
        offset: u64,
        limit: usize,
    ) -> Result<Option<Vec<StoredEntry>>, RoutedStreamError> {
        if limit == 0 {
            return Err(RoutedStreamRefusal::InvalidLimit.into());
        }
        let next_sequence = if let Some(pending) = self.streams.get(stream) {
            pending.next_sequence
        } else if let Some(head) = read_head(store, &self.route_key, stream).await? {
            head.next_sequence
        } else {
            return Ok(None);
        };
        let (sequences, mut rows) = prepare_page(next_sequence, offset, limit)?;
        for sequence in sequences {
            let key = routed_stream_codec::row_key(stream, sequence)?;
            let row = match self.keys.get(&key) {
                Some(Some(index)) => {
                    let item = self
                        .items
                        .get(*index)
                        .ok_or(RoutedStreamRefusal::BatchRefused)?;
                    if item.key != key {
                        return Err(RoutedStreamRefusal::BatchRefused.into());
                    }
                    match routed_stream_codec::decode_key(&item.key)? {
                        routed_stream_codec::StreamKey::Row {
                            stream: actual,
                            sequence: position,
                        } if actual == stream && position == sequence => {}
                        _ => return Err(RoutedStreamRefusal::KeyCollision.into()),
                    }
                    routed_stream_codec::decode_row(&item.value, sequence)?
                }
                Some(None) => return Err(RoutedStreamRefusal::KeyCollision.into()),
                None => {
                    let bytes =
                        store
                            .get_routed(&self.route_key, &key)
                            .await?
                            .ok_or_else(|| RoutedStreamRefusal::MissingRow {
                                stream: stream.as_bytes().to_vec(),
                                sequence,
                            })?;
                    routed_stream_codec::decode_row(&bytes, sequence)?
                }
            };
            rows.push(row);
        }
        Ok(Some(rows))
    }

    /// Stages a fresh checkpoint after the caller loudly rebuilt its full history.
    ///
    /// # Errors
    /// Keeps the corrupt bytes' CAS witness and refuses any concurrent replacement.
    pub(super) fn stage_rebuilt_checkpoint<'a, L: CheckpointLiveState, const N: usize>(
        &'a mut self,
        store: &'a dyn DurableStore,
        checkpoint: &RestoreCheckpoint<L, N>,
        prefix_digest: [u8; 32],
        observed_locator: Option<&'a [u8]>,
    ) -> impl std::future::Future<Output = Result<(), CheckpointStageError>> + Send + use<'a, L, N>
    {
        let body = if self.staging_refused {
            Err(CheckpointStageError::BatchRefused)
        } else {
            checkpoint
                .encode_canonical()
                .map_err(CheckpointStageError::from)
        };
        self.finish_checkpoint(
            store,
            checkpoint.positions(),
            body,
            prefix_digest,
            true,
            observed_locator,
        )
    }

    /// Encodes the complete live owner before any store future is returned.
    ///
    /// # Errors
    /// Refuses invalid positions, root framing or joint live-owner validation.
    pub(super) fn stage_borrowed_checkpoint<'a, L: CheckpointLiveState, const N: usize>(
        &'a mut self,
        store: &'a dyn DurableStore,
        positions: CheckpointPositions,
        history_root: &[u8; N],
        live: &L,
        prefix_digest: [u8; 32],
    ) -> impl std::future::Future<Output = Result<(), CheckpointStageError>> + Send + use<'a, L, N>
    {
        let body = (|| {
            if self.staging_refused {
                return Err(CheckpointStageError::BatchRefused);
            }
            if N == 0 || u16::try_from(N).is_err() {
                return Err(CheckpointError::RootWidth {
                    expected: N,
                    actual: N,
                }
                .into());
            }
            positions.validate()?;
            live.validate_live(&positions, history_root)?;
            let live_bytes = live.encode_live()?;
            Ok(super::restore_checkpoint_codec::encode(
                positions,
                history_root,
                &live_bytes,
            )?)
        })();
        self.finish_checkpoint(store, positions, body, prefix_digest, false, None)
    }

    fn finish_checkpoint<'a>(
        &'a mut self,
        store: &'a dyn DurableStore,
        positions: CheckpointPositions,
        body: Result<Vec<u8>, CheckpointStageError>,
        prefix_digest: [u8; 32],
        rebuilt: bool,
        observed_locator: Option<&'a [u8]>,
    ) -> impl std::future::Future<Output = Result<(), CheckpointStageError>> + Send + use<'a> {
        if body.is_err() {
            self.refuse();
        }
        async move {
            let result = match body {
                Ok(body) => {
                    self.prepare_checkpoint(
                        store,
                        positions,
                        body,
                        prefix_digest,
                        rebuilt,
                        observed_locator,
                    )
                    .await
                }
                Err(error) => Err(error),
            };
            if result.is_err() {
                self.refuse();
            }
            result
        }
    }

    async fn prepare_checkpoint(
        &mut self,
        store: &dyn DurableStore,
        positions: CheckpointPositions,
        body: Vec<u8>,
        prefix_digest: [u8; 32],
        rebuilt: bool,
        observed_locator: Option<&[u8]>,
    ) -> Result<(), CheckpointStageError> {
        if self.route_key != checkpoint_conversation_route(positions.conversation_id) {
            return Err(CheckpointStageError::RouteMismatch);
        }
        let body_key = CheckpointBodyKey::for_body(&body);
        let locator = CheckpointLocator::new(
            positions.conversation_id,
            positions.covered_log_sequence,
            positions.extension_sequence,
            prefix_digest,
            body_key,
        )?;
        let prior = if rebuilt {
            observed_locator.map(<[u8]>::to_vec)
        } else {
            store
                .get_routed(&self.route_key, CheckpointLocator::routed_key())
                .await?
        };
        if !rebuilt {
            if let Some(bytes) = prior.as_ref() {
                let previous =
                    CheckpointLocator::decode_canonical(bytes, positions.conversation_id)?;
                if previous.covered_sequence() > positions.covered_log_sequence
                    || previous.extension_boundary() > positions.extension_sequence
                {
                    return Err(CheckpointStageError::PositionRegression);
                }
            }
        }
        let key = body_key.routed_key();
        let existing_body = store.get_routed(&self.route_key, &key).await?;
        if !rebuilt && existing_body.as_ref().is_some_and(|bytes| bytes != &body) {
            return Err(CheckpointStageError::BodyMismatch);
        }
        self.stage(RoutedBatchItem {
            key,
            expected: existing_body.as_deref().map(Hash::of),
            value: body,
            ttl: None,
        })?;
        self.stage(RoutedBatchItem {
            key: CheckpointLocator::routed_key().to_vec(),
            expected: prior.as_deref().map(Hash::of),
            value: locator.encode_canonical().to_vec(),
            ttl: None,
        })?;
        Ok(())
    }

    /// Determines whether the operation crosses a checkpoint boundary.
    ///
    /// # Errors
    /// Refuses a regression in the covered operation count.
    pub(super) fn checkpoint_due(
        before: u64,
        after: u64,
        sealed: bool,
    ) -> Result<bool, DurabilityError> {
        if after < before {
            return Err(DurabilityError::ConfigError(
                "covered operation count regressed".into(),
            ));
        }
        Ok(sealed || before / CHECKPOINT_EVERY != after / CHECKPOINT_EVERY)
    }

    /// Publishes all staged rows together, without a separate flush or retry.
    /// The submission callback runs after local preflight, when the future is polled.
    ///
    /// # Errors
    /// Propagates the atomic store's refusal or unresolved outcome unchanged.
    pub(super) async fn commit(
        mut self,
        store: &dyn DurableStore,
        on_submit: impl FnOnce(),
    ) -> Result<(), DurabilityError> {
        if self.staging_refused {
            return Err(DurabilityError::ConfigError(
                "operation batch staging was refused".into(),
            ));
        }
        if self.items.is_empty() && self.streams.is_empty() {
            return Err(DurabilityError::EmptyOperationBatch);
        }
        for stream in self.streams.into_values() {
            self.items.push(RoutedBatchItem {
                key: stream.head_key,
                expected: stream.expected,
                value: routed_stream_codec::encode_head(stream.next_sequence).to_vec(),
                ttl: None,
            });
        }
        on_submit();
        // One durable operation, settled like every other write: an unknown
        // outcome is read back at every key before it is answered.
        apply_routed_batch_settled(store, &self.route_key, self.items).await
    }
}

#[cfg(test)]
mod tests;

#[cfg(test)]
mod stream_tests;

#[cfg(test)]
mod overlay_tests;

#[cfg(test)]
mod log_tests;

#[cfg(test)]
mod checkpoint_tests;

#[cfg(test)]
mod bootstrap_tests;

#[cfg(test)]
mod serving_atomic_tests;