liminal-server 0.23.1

//! DUR-020 R3: the server-scope identity count, carried across a clean stop.
//!
//! Enrollment is the one request a server-scope count can refuse (identity
//! Server), and the ledger that count lives in is exact only once every
//! registered conversation has folded its contribution. A clean stop writes
//! each conversation's identity share here, one routed value on the registry
//! route, so the next process seeds its ledger from it and enrollment decides
//! at once instead of waiting for the warm. A conversation's own restore then
//! replaces its seeded share with the exact one (`fold_conversation`
//! replaces, never adds).
//!
//! The snapshot is true only for the process that starts right after the stop
//! that wrote it. So before this process commits its first enrollment, it
//! overwrites the snapshot with the consumed marker: a crash after that point
//! leaves no snapshot to trust, and the next start waits for its warm. After
//! an unclean stop there is never a valid snapshot, because the stop that
//! would have written one did not run.
//!
//! Layout, big-endian, in the style of the registry snapshot: magic `LSCS`,
//! version 1.0 or 1.1 (u16 each), state (u8: 1 valid, 2 consumed), seven zero bytes,
//! share count (u64), incarnation (u64 in version 1.1; zero padding in 1.0),
//! then each share as conversation id and
//! identity count (u64 each, ids strictly ascending), closed by a blake3
//! digest over a domain and every byte before it.

use std::collections::BTreeMap;

use haematite::tree::Hash;
use liminal::durability::bridge::block_on;
use liminal::durability::{DurabilityError, DurableStore, RoutedBatchItem};

use crate::server::participant::ParticipantSemanticError;

use super::handler::bridge_error;
use super::operation_batch::OperationBatch;
use super::registry::REGISTRY_STREAM_KEY;

/// The snapshot's local key under the registry route.
pub(super) const CAPACITY_KEY: &[u8] = b"registry/capacity/v1";
const MAGIC: [u8; 4] = *b"LSCS";
const VERSION: (u16, u16) = (1, 0);
const VALID: u8 = 1;
const CONSUMED: u8 = 2;
const DIGEST_DOMAIN: &[u8] = b"liminal-capacity-snapshot-v1";
const HEADER_BYTES: usize = 32;
const DIGEST_BYTES: usize = 32;
const SHARE_BYTES: usize = 16;
const STATE_AT: usize = 8;
const COUNT_AT: usize = 16;

/// One decoded capacity snapshot.
#[derive(Clone, Debug, PartialEq, Eq)]
pub(super) enum CapacitySnapshot {
    /// A clean stop whose codec has no incarnation field.
    LegacyCleanStop(BTreeMap<u64, u64>),
    /// A clean stop with the incarnation that published these shares.
    CleanStop {
        server_incarnation: u64,
        shares: BTreeMap<u64, u64>,
    },
    /// Overwritten by the process that seeded from it, before its first
    /// enrollment committed.
    Consumed,
}

/// Named refusal of a capacity snapshot. A refused snapshot is never
/// trusted; enrollment then waits for the warm.
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub(super) enum CapacitySnapshotError {
    #[error("capacity snapshot: length {actual} differs from {expected}")]
    Length { expected: u64, actual: u64 },
    #[error("capacity snapshot: magic differs")]
    Magic,
    #[error("capacity snapshot: codec {major}.{minor} is unsupported")]
    Version { major: u16, minor: u16 },
    #[error("capacity snapshot: state {state} is unknown")]
    State { state: u8 },
    #[error("capacity snapshot: padding is not zero")]
    Padding,
    #[error("capacity snapshot: conversation {id} at index {index} is not above {previous}")]
    Order {
        index: usize,
        previous: u64,
        id: u64,
    },
    #[error("capacity snapshot: a consumed snapshot carries {count} shares")]
    ConsumedShares { count: u64 },
    #[error("capacity snapshot: a clean stop names incarnation zero")]
    IncarnationZero,
    #[error("capacity snapshot: codec 1.1 carries state {state} instead of a clean stop")]
    IncarnationState { state: u8 },
    #[error("capacity snapshot: digest differs")]
    Digest,
    #[error("capacity snapshot: {count} shares cannot be framed")]
    TooLarge { count: usize },
}

impl CapacitySnapshot {
    /// The canonical bytes: `64 + 16 x shares` long.
    ///
    /// # Errors
    ///
    /// Refuses a share map too large to frame.
    pub(super) fn encode(&self) -> Result<Vec<u8>, CapacitySnapshotError> {
        let (state, shares, incarnation) = match self {
            Self::LegacyCleanStop(shares) => (VALID, Some(shares), None),
            Self::CleanStop {
                server_incarnation: 0,
                ..
            } => return Err(CapacitySnapshotError::IncarnationZero),
            Self::CleanStop {
                server_incarnation,
                shares,
            } => (VALID, Some(shares), Some(*server_incarnation)),
            Self::Consumed => (CONSUMED, None, None),
        };
        let version = if incarnation.is_some() {
            (1_u16, 1_u16)
        } else {
            VERSION
        };
        let len = shares.map_or(0, BTreeMap::len);
        let count =
            u64::try_from(len).map_err(|_| CapacitySnapshotError::TooLarge { count: len })?;
        let mut bytes = Vec::with_capacity(
            HEADER_BYTES
                .saturating_add(len.saturating_mul(SHARE_BYTES))
                .saturating_add(DIGEST_BYTES),
        );
        bytes.extend_from_slice(&MAGIC);
        bytes.extend_from_slice(&version.0.to_be_bytes());
        bytes.extend_from_slice(&version.1.to_be_bytes());
        bytes.push(state);
        bytes.extend_from_slice(&[0; 7]);
        bytes.extend_from_slice(&count.to_be_bytes());
        bytes.extend_from_slice(&incarnation.map_or([0; 8], u64::to_be_bytes));
        for (conversation_id, identities) in shares.into_iter().flatten() {
            bytes.extend_from_slice(&conversation_id.to_be_bytes());
            bytes.extend_from_slice(&identities.to_be_bytes());
        }
        let sealed = digest(&bytes);
        bytes.extend_from_slice(&sealed);
        Ok(bytes)
    }

    /// Reads one canonical snapshot, refusing every malformed byte by name,
    /// in a fixed order: length, magic, version, state, padding, count
    /// against length, share order, then the digest.
    ///
    /// # Errors
    ///
    /// The [`CapacitySnapshotError`] naming the first check the bytes fail.
    pub(super) fn decode(input: &[u8]) -> Result<Self, CapacitySnapshotError> {
        let minimum = HEADER_BYTES + DIGEST_BYTES;
        if input.len() < minimum {
            return Err(length(width(minimum), input));
        }
        if input.get(..4) != Some(MAGIC.as_slice()) {
            return Err(CapacitySnapshotError::Magic);
        }
        let major = u16::from_be_bytes(take(input, 4)?);
        let minor = u16::from_be_bytes(take(input, 6)?);
        if major != 1 || minor > 1 {
            return Err(CapacitySnapshotError::Version { major, minor });
        }
        let [state] = take::<1>(input, STATE_AT)?;
        if state != VALID && state != CONSUMED {
            return Err(CapacitySnapshotError::State { state });
        }
        let incarnation_bytes = take::<8>(input, COUNT_AT + 8)?;
        if take::<7>(input, STATE_AT + 1)? != [0; 7] || (minor == 0 && incarnation_bytes != [0; 8])
        {
            return Err(CapacitySnapshotError::Padding);
        }
        let incarnation = if minor == 0 {
            None
        } else {
            let value = u64::from_be_bytes(incarnation_bytes);
            if value == 0 {
                return Err(CapacitySnapshotError::IncarnationZero);
            }
            if state != VALID {
                return Err(CapacitySnapshotError::IncarnationState { state });
            }
            Some(value)
        };
        let count = u64::from_be_bytes(take(input, COUNT_AT)?);
        let expected = count
            .checked_mul(width(SHARE_BYTES))
            .and_then(|bytes| bytes.checked_add(width(minimum)))
            .unwrap_or(u64::MAX);
        if width(input.len()) != expected {
            return Err(length(expected, input));
        }
        if state == CONSUMED && count != 0 {
            return Err(CapacitySnapshotError::ConsumedShares { count });
        }
        let (body, stored) = input
            .split_at_checked(input.len().saturating_sub(DIGEST_BYTES))
            .ok_or_else(|| length(expected, input))?;
        let shares = ascending_shares(body.get(HEADER_BYTES..).unwrap_or_default())?;
        if digest(body).as_slice() != stored {
            return Err(CapacitySnapshotError::Digest);
        }
        Ok(if state == VALID {
            match incarnation {
                None => Self::LegacyCleanStop(shares),
                Some(server_incarnation) => Self::CleanStop {
                    server_incarnation,
                    shares,
                },
            }
        } else {
            Self::Consumed
        })
    }
}

/// A stored snapshot's bytes hash, the compare-and-swap witness for the next
/// write, with its decoding.
pub(super) type StoredCapacity = (Hash, Result<CapacitySnapshot, CapacitySnapshotError>);

/// The stored snapshot's bytes hash, the compare-and-swap witness for the
/// next write, and its decoding; `None` when no snapshot was ever written.
///
/// # Errors
///
/// The store's refusal or unknown outcome.
pub(super) fn read_capacity_snapshot(
    store: &dyn DurableStore,
) -> Result<Option<StoredCapacity>, ParticipantSemanticError> {
    let stored = block_on(store.get_routed(REGISTRY_STREAM_KEY, CAPACITY_KEY))
        .map_err(|error| bridge_error(&error))?
        .map_err(|error| store_error(&error))?;
    Ok(stored.map(|bytes| (Hash::of(&bytes), CapacitySnapshot::decode(&bytes))))
}

/// Writes `snapshot` over the bytes whose hash is `expected` (absent when no
/// snapshot was ever written), as one batch on the registry route, and
/// answers the hash of the bytes it wrote.
///
/// # Errors
///
/// An unframeable share map, a changed witness, or the store's refusal or
/// unknown outcome.
pub(super) fn write_capacity_snapshot(
    store: &dyn DurableStore,
    expected: Option<Hash>,
    snapshot: &CapacitySnapshot,
) -> Result<Hash, ParticipantSemanticError> {
    let value = snapshot
        .encode()
        .map_err(|error| ParticipantSemanticError::Internal {
            message: error.to_string(),
        })?;
    let hash = Hash::of(&value);
    // The held witness already proves these bytes were committed. A later
    // shutdown phase owes another write only when the ledger changed.
    if expected == Some(hash) {
        return Ok(hash);
    }
    let mut batch = OperationBatch::new(REGISTRY_STREAM_KEY);
    batch
        .stage(RoutedBatchItem {
            key: CAPACITY_KEY.to_vec(),
            expected,
            value,
            ttl: None,
        })
        .map_err(|error| store_error(&error))?;
    block_on(batch.commit(store, || {}))
        .map_err(|error| bridge_error(&error))?
        .map_err(|error| store_error(&error))?;
    Ok(hash)
}

fn store_error(error: &DurabilityError) -> ParticipantSemanticError {
    ParticipantSemanticError::Internal {
        message: format!("capacity snapshot store operation failed: {error}"),
    }
}

fn digest(body: &[u8]) -> [u8; DIGEST_BYTES] {
    let mut hash = blake3::Hasher::new();
    hash.update(DIGEST_DOMAIN);
    hash.update(body);
    *hash.finalize().as_bytes()
}

fn ascending_shares(run: &[u8]) -> Result<BTreeMap<u64, u64>, CapacitySnapshotError> {
    let mut shares = BTreeMap::new();
    let mut previous: Option<u64> = None;
    for (index, chunk) in run.chunks_exact(SHARE_BYTES).enumerate() {
        let id = u64::from_be_bytes(take(chunk, 0)?);
        let identities = u64::from_be_bytes(take(chunk, 8)?);
        if let Some(previous) = previous.filter(|previous| id <= *previous) {
            return Err(CapacitySnapshotError::Order {
                index,
                previous,
                id,
            });
        }
        previous = Some(id);
        shares.insert(id, identities);
    }
    Ok(shares)
}

fn length(expected: u64, input: &[u8]) -> CapacitySnapshotError {
    CapacitySnapshotError::Length {
        expected,
        actual: width(input.len()),
    }
}

/// A byte count as a frame width; no supported target has a wider `usize`.
fn width(bytes: usize) -> u64 {
    u64::try_from(bytes).unwrap_or(u64::MAX)
}

fn take<const N: usize>(input: &[u8], start: usize) -> Result<[u8; N], CapacitySnapshotError> {
    start
        .checked_add(N)
        .and_then(|end| input.get(start..end))
        .and_then(|bytes| <[u8; N]>::try_from(bytes).ok())
        .ok_or_else(|| length(width(HEADER_BYTES + DIGEST_BYTES), input))
}

/// What this process knows of the clean-stop snapshot: the witness its next
/// write must name, whether it seeded the ledger, and whether it has consumed
/// the snapshot it seeded from.
#[derive(Debug, Default)]
pub(super) struct CapacitySeed {
    state: std::sync::Mutex<SeedState>,
}

#[derive(Debug, Default)]
struct SeedState {
    witness: Option<Hash>,
    seeded: bool,
    consumed: bool,
}

impl CapacitySeed {
    fn state(&self) -> std::sync::MutexGuard<'_, SeedState> {
        // Every critical section below assigns plain fields and cannot leave
        // them torn, so a poisoned guard is recovered rather than refused.
        self.state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }
}

impl super::handler::ProductionParticipantHandler {
    /// Construction: seeds the server-scope ledger from the clean-stop
    /// snapshot when one is valid. Any other answer leaves the ledger empty,
    /// and enrollment then waits for the warm; nothing here is fatal.
    pub(super) fn seed_capacity_from_snapshot(&self) {
        let read = read_capacity_snapshot(self.store.as_ref());
        let mut state = self.capacity_seed.state();
        match read {
            Ok(None) => {}
            Ok(Some((
                witness,
                Ok(
                    CapacitySnapshot::LegacyCleanStop(shares)
                    | CapacitySnapshot::CleanStop { shares, .. },
                ),
            ))) => {
                state.witness = Some(witness);
                match self.capacity.seed_identities(&shares) {
                    Ok(()) => {
                        state.seeded = true;
                        tracing::info!(
                            step = "participant restore",
                            conversations = shares.len(),
                            "capacity seeded from the clean stop's snapshot; enrollment waits for nothing"
                        );
                    }
                    Err(error) => tracing::warn!(
                        error = %error,
                        "capacity snapshot refused at seeding; enrollment waits for the warm"
                    ),
                }
            }
            Ok(Some((witness, Ok(CapacitySnapshot::Consumed)))) => {
                state.witness = Some(witness);
                tracing::info!(
                    step = "participant restore",
                    "no clean stop since the capacity snapshot was consumed; enrollment waits for the warm"
                );
            }
            Ok(Some((witness, Err(refusal)))) => {
                state.witness = Some(witness);
                tracing::warn!(
                    refusal = %refusal,
                    "capacity snapshot refused by name; enrollment waits for the warm"
                );
            }
            Err(error) => tracing::warn!(
                error = %error,
                "capacity snapshot unreadable; enrollment waits for the warm"
            ),
        }
        drop(state);
    }

    /// Whether the ledger is exact without the warm: seeded from a valid
    /// snapshot that this process has consumed. The first call consumes it,
    /// before any enrollment commits; a refused consumption falls back to
    /// waiting for the warm for the rest of this process.
    pub(super) fn capacity_ready_without_warm(&self) -> bool {
        let mut state = self.capacity_seed.state();
        if !state.seeded {
            return false;
        }
        if state.consumed {
            return true;
        }
        // The consumed marker must be durable before any enrollment of this
        // process can be: under a buffered commit policy the marker's shard
        // could otherwise lose it while the enrollment's shard keeps its row,
        // and the next start would trust a stale snapshot.
        let consumed = write_capacity_snapshot(
            self.store.as_ref(),
            state.witness,
            &CapacitySnapshot::Consumed,
        )
        .and_then(|witness| {
            block_on(self.store.flush())
                .map_err(|error| bridge_error(&error))?
                .map_err(|error| store_error(&error))?;
            Ok(witness)
        });
        let ready = match consumed {
            Ok(witness) => {
                state.witness = Some(witness);
                state.consumed = true;
                true
            }
            Err(error) => {
                state.seeded = false;
                tracing::warn!(
                    error = %error,
                    "capacity snapshot could not be consumed; enrollment waits for the warm"
                );
                false
            }
        };
        // Held to here on purpose: a second enrollment must not consume the
        // snapshot again while this one writes the marker.
        drop(state);
        ready
    }

    /// The clean stop: records every conversation's identity share, when the
    /// ledger holds every one of them (the warm finished, or the ledger was
    /// seeded). Otherwise it writes nothing and the next start waits for its
    /// warm, said by name.
    pub(super) fn record_clean_stop_capacity(&self) {
        let mut state = self.capacity_seed.state();
        // A refused conversation never folded its identities, so a ledger
        // that names one is not complete either.
        let complete = state.seeded
            || (self.pending.remaining().is_ok_and(|ids| ids.is_empty())
                && self.unloadable_record().snapshot().is_empty());
        if !complete {
            tracing::info!(
                "clean stop writes no capacity snapshot: some conversations were never restored or refused, so the next start waits for its warm"
            );
            return;
        }
        let shares = self.capacity.identity_shares();
        let snapshot = match self.pending.owed_repair() {
            Ok(Some(server_incarnation)) => CapacitySnapshot::CleanStop {
                server_incarnation,
                shares,
            },
            Ok(None) => CapacitySnapshot::LegacyCleanStop(shares),
            Err(error) => {
                tracing::warn!(%error, "clean stop could not read its incarnation; no capacity snapshot was written");
                return;
            }
        };
        match write_capacity_snapshot(self.store.as_ref(), state.witness, &snapshot) {
            Ok(witness) => state.witness = Some(witness),
            Err(error) => tracing::warn!(
                error = %error,
                "clean stop could not write its capacity snapshot; the next start waits for its warm"
            ),
        }
        drop(state);
    }
}