liminal/durability/dedup/

in_flight.rs

//! The in-flight index: where a start finds the claims still in flight without
//! walking the dedup namespace.
//!
//! A claim made with evidence first appends a row naming its key and the
//! sequence its claim takes in the key's own stream to the namespace's index
//! stream, then writes the claim there. A start reads the index from its
//! recorded start and, for each named key, reads one bounded page of two rows
//! from the key's stream at that claim sequence: the claim and the row after
//! it. It never reads a key's head or its history, so a start costs the same
//! however many claims a key has had (DUR-008 A12). A key whose claim was
//! completed or released (a row after the claim, or a receipt on it) or never
//! written (no row at the claim) is dropped there; the others are in flight.
//! A completion or release writes nothing to the index: the process forgets the
//! key, and the recorded start moves past rows that name no claim it holds.
//!
//! The start is recorded by count, as a channel's head is: once the rows behind
//! the oldest claim held in flight reach [`IN_FLIGHT_CHECKPOINT_EVERY`]. A
//! claim held so long that the index grows twice that far past the recorded
//! start is carried: its key is appended again at the tail, with every other
//! such key in one append, and the start moves past its old row. A start
//! therefore reads at most about twice [`IN_FLIGHT_CHECKPOINT_EVERY`] rows plus
//! one per claim in flight, however much the namespace holds. Nothing here runs
//! on a clock.
//!
//! An index row written before the claim sequence was carried (`LDX1`, by a
//! binary older than A12's) names its key alone; a start looks that key up in
//! full, as it did before. Every row this binary writes carries the sequence,
//! so such a row is read until the start moves past it.
//!
//! Crash windows, each repaired by the next start's look-up of the named keys:
//! after the index row and before the claim, the key's stream holds no claim
//! and the row is dropped; after a completion or release and before the start
//! moves, the claim is found settled and the row is dropped; after a carry and
//! before the start moves, the key is read twice and kept once.
//!
//! A store written before the index existed (liminal-server 0.16.0) has no
//! recorded start. The first start over it walks the namespace once, appends a
//! row for every claim it finds in flight in one append, and records the start
//! at those rows. A crash before that record repeats the walk; nothing after it
//! does.
//!
//! A namespace the caller knows is empty (a store this process created, or a
//! namespace no earlier process used) is neither read nor walked: its index
//! starts empty, and its start is recorded by its first claim, before that
//! claim is written, so a later start never mistakes it for a store that
//! predates the index.
use std::collections::{BTreeMap, HashMap};
use std::sync::{Mutex, MutexGuard, PoisonError};

use super::codec::{self, DedupRecord};
use super::{ClaimEvidence, DedupCache, DedupEntry};
use crate::durability::DurabilityError;
use crate::durability::head_checkpoint::READ_BATCH_SIZE;

/// How many index rows may lie behind the oldest claim held in flight before
/// the start is recorded again. This bounds the rows a start reads.
pub(super) const IN_FLIGHT_CHECKPOINT_EVERY: u64 = 256;

/// The rows a start reads from a named key's stream: the claim the index row
/// names and the one after it, which settles it.
const CLAIM_PAGE: usize = 2;

/// How many times one index append retries after another writer took its row.
const INDEX_APPEND_ATTEMPTS: usize = 64;

/// A claim this process holds in flight: the index row naming it and the
/// sequence of the claim in its key's stream.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
struct OpenClaim {
    row: u64,
    claim: u64,
}

/// What the process knows of its namespace's index.
#[derive(Debug)]
struct IndexState {
    /// The first row a start reads, as last recorded.
    start: u64,
    /// The index stream's next row, as this process last saw it.
    tail: u64,
    /// Claims held in flight, by key.
    open: HashMap<String, OpenClaim>,
    /// Appends under way, by the row they were reserved at: none of them lands
    /// below it, so the start never moves past it.
    reserved: BTreeMap<u64, usize>,
    /// Whether a carry is under way.
    carrying: bool,
    /// Whether the start is in the store. Only an index begun empty by this
    /// process lacks it, until its first join records it.
    recorded: bool,
}

impl IndexState {
    /// The index of a namespace nothing has written to.
    fn empty() -> Self {
        Self {
            start: 0,
            tail: 0,
            open: HashMap::new(),
            reserved: BTreeMap::new(),
            carrying: false,
            recorded: false,
        }
    }

    /// The lowest row a start must still read.
    fn floor(&self) -> u64 {
        let open = self.open.values().map(|claim| claim.row).min();
        let reserved = self.reserved.keys().next().copied();
        [open, reserved]
            .into_iter()
            .flatten()
            .fold(self.tail, u64::min)
    }

    fn reserve(&mut self) -> u64 {
        let floor = self.tail;
        *self.reserved.entry(floor).or_insert(0) += 1;
        floor
    }

    fn unreserve(&mut self, floor: u64) {
        if let Some(count) = self.reserved.get_mut(&floor) {
            *count = count.saturating_sub(1);
            if *count == 0 {
                self.reserved.remove(&floor);
            }
        }
    }

    fn appended(&mut self, after: u64) {
        self.tail = self.tail.max(after);
    }

    fn leave(&mut self, key: &str, settled: u64) -> bool {
        if self.open.get(key).is_some_and(|open| open.claim <= settled) {
            self.open.remove(key);
            return true;
        }
        false
    }

    fn plan(&mut self) -> Option<Advance> {
        if self.carrying {
            return None;
        }
        let floor = self.floor();
        if floor >= self.start.saturating_add(IN_FLIGHT_CHECKPOINT_EVERY) {
            return Some(Advance::Record(floor));
        }
        if self.tail < self.start.saturating_add(2 * IN_FLIGHT_CHECKPOINT_EVERY) {
            return None;
        }
        let horizon = self.tail.saturating_sub(IN_FLIGHT_CHECKPOINT_EVERY);
        let mut keys: Vec<(String, OpenClaim)> = self
            .open
            .iter()
            .filter(|(_, open)| open.row < horizon)
            .map(|(key, open)| (key.clone(), *open))
            .collect();
        if keys.is_empty() {
            return None;
        }
        keys.sort_unstable();
        self.carrying = true;
        let at = self.reserve();
        Some(Advance::Carry { keys, at })
    }

    fn carried(
        &mut self,
        keys: &[(String, OpenClaim)],
        first: Option<u64>,
        at: u64,
    ) -> Option<u64> {
        if let Some(first) = first {
            let mut row = first;
            for (key, carried) in keys {
                if let Some(open) = self.open.get_mut(key) {
                    if open.row == carried.row {
                        open.row = row;
                    }
                }
                row = row.saturating_add(1);
            }
            self.appended(row);
        }
        self.unreserve(at);
        self.carrying = false;
        let floor = self.floor();
        (floor > self.start).then_some(floor)
    }
}

/// What moving the start needs next.
enum Advance {
    /// Record the start at this row.
    Record(u64),
    /// Append these keys again, each with the claim held for it, at the
    /// reserved row, then record.
    Carry {
        keys: Vec<(String, OpenClaim)>,
        at: u64,
    },
}

/// The process's view of one namespace's index, shared by the cache's clones.
#[derive(Debug, Default)]
pub(super) struct InFlightIndex(Mutex<Option<IndexState>>);

impl InFlightIndex {
    /// The index of a namespace the caller knows is empty, loaded without
    /// reading the store.
    pub(super) fn empty() -> Self {
        Self(Mutex::new(Some(IndexState::empty())))
    }

    fn state(&self) -> MutexGuard<'_, Option<IndexState>> {
        // The state is rebuilt from the durable index on the next start; a
        // panic elsewhere cannot make it claim less than the store holds.
        self.0.lock().unwrap_or_else(PoisonError::into_inner)
    }

    fn is_loaded(&self) -> bool {
        self.state().is_some()
    }

    /// Installs a loaded state unless another load got there first.
    fn install(&self, loaded: IndexState) -> bool {
        let mut state = self.state();
        if state.is_some() {
            return false;
        }
        *state = Some(loaded);
        true
    }

    fn reserve(&self) -> Option<u64> {
        self.state().as_mut().map(IndexState::reserve)
    }

    fn unreserve(&self, floor: u64) {
        if let Some(state) = self.state().as_mut() {
            state.unreserve(floor);
        }
    }

    fn joined(&self, key: &str, row: u64, claim: u64) {
        if let Some(state) = self.state().as_mut() {
            state.appended(row.saturating_add(1));
            state.open.insert(key.to_owned(), OpenClaim { row, claim });
        }
    }

    /// Forgets a join whose claim was not made, when it is still the one held.
    fn unjoin(&self, key: &str, row: u64) {
        if let Some(state) = self.state().as_mut() {
            if state.open.get(key).is_some_and(|open| open.row == row) {
                state.open.remove(key);
            }
        }
    }

    /// Forgets `key` when the claim held for it is the one settled at
    /// `settled` or an earlier one; a claim made after the settlement stays.
    fn leave(&self, key: &str, settled: u64) -> bool {
        self.with_state(|state| state.leave(key, settled))
            .unwrap_or(false)
    }

    /// The claims held, each key with its claim sequence.
    fn open_claims(&self) -> Vec<(String, u64)> {
        self.state()
            .as_ref()
            .map(|state| {
                state
                    .open
                    .iter()
                    .map(|(key, open)| (key.clone(), open.claim))
                    .collect()
            })
            .unwrap_or_default()
    }

    /// Whether this is an index begun empty whose start is not yet recorded.
    fn awaits_record(&self) -> bool {
        self.state().as_ref().is_some_and(|state| !state.recorded)
    }

    fn start_recorded(&self) {
        if let Some(state) = self.state().as_mut() {
            state.recorded = true;
        }
    }

    fn start(&self) -> Option<u64> {
        self.state().as_ref().map(|state| state.start)
    }

    fn recorded(&self, start: u64) {
        if let Some(state) = self.state().as_mut() {
            state.start = state.start.max(start);
        }
    }

    fn plan(&self) -> Option<Advance> {
        self.with_state(IndexState::plan).flatten()
    }

    /// Moves each carried key that is still held to its new row, answering
    /// the start the carry lets the index record.
    fn carried(&self, keys: &[(String, OpenClaim)], first: Option<u64>, at: u64) -> Option<u64> {
        self.with_state(|state| state.carried(keys, first, at))
            .flatten()
    }

    fn with_state<T>(&self, apply: impl FnOnce(&mut IndexState) -> T) -> Option<T> {
        self.state().as_mut().map(apply)
    }
}

impl DedupCache {
    /// The stream holding this namespace's in-flight index rows. It sits
    /// outside the `namespace:` prefix the key streams share, so no walk of
    /// the key streams reads it.
    fn index_stream(&self) -> String {
        format!("{}/in-flight", self.namespace)
    }

    /// The value key recording the index start, kept as the start plus one
    /// because the value store keeps no physical zero.
    fn index_start_key(&self) -> String {
        format!("{}/in-flight-start", self.namespace)
    }

    /// The claims in flight in this namespace with the evidence each carries,
    /// in key order, read from the in-flight index and never by walking the
    /// namespace (except once, on the first start over a store written before
    /// the index existed).
    ///
    /// # Errors
    /// Refuses an unreadable index or key stream, or a malformed record; never
    /// answers a partial list.
    pub async fn in_flight_claims(
        &self,
    ) -> Result<Vec<(String, Option<ClaimEvidence>)>, DurabilityError> {
        if !self.in_flight.is_loaded() {
            let (state, claims) = self.load_in_flight_index().await?;
            if self.in_flight.install(state) {
                self.advance_in_flight().await;
                return Ok(claims);
            }
        }
        let mut held = self.in_flight.open_claims();
        held.sort_unstable();
        let mut claims = Vec::new();
        for (key, claim) in held {
            if let Some(entry) = self.claim_at(&key, claim).await? {
                claims.push((key, entry.evidence().cloned()));
            }
        }
        Ok(claims)
    }

    /// Loads the index state when this process has not yet.
    pub(super) async fn ensure_in_flight_index(&self) -> Result<(), DurabilityError> {
        if self.in_flight.is_loaded() {
            return Ok(());
        }
        let (state, _) = self.load_in_flight_index().await?;
        if self.in_flight.install(state) {
            self.advance_in_flight().await;
        }
        Ok(())
    }

    /// Appends the index row naming `key` before its claim is written,
    /// answering the row. `claim` is the sequence the claim will take in the
    /// key's stream.
    pub(super) async fn join_in_flight(
        &self,
        key: &str,
        claim: u64,
    ) -> Result<u64, DurabilityError> {
        self.ensure_in_flight_index().await?;
        self.record_empty_start().await?;
        let floor = self.in_flight.reserve().ok_or_else(|| {
            DurabilityError::ConfigError("in-flight index is not loaded".to_owned())
        })?;
        let appended = self
            .append_index_rows(vec![codec::encode_join(key, claim)?], floor)
            .await;
        if let Ok(row) = appended {
            self.in_flight.joined(key, row, claim);
        }
        self.in_flight.unreserve(floor);
        appended
    }

    /// Records the start of an index begun empty, before the namespace's
    /// first claim is written: a store holding claims and no recorded start
    /// reads as one written before the index existed, and its next start
    /// would walk the namespace. Does nothing for any other index.
    pub(super) async fn record_empty_start(&self) -> Result<(), DurabilityError> {
        if !self.in_flight.awaits_record() {
            return Ok(());
        }
        match self.store.cas(&self.index_start_key(), 0, 1).await {
            // Another join recorded it first, at the same start.
            Ok(()) | Err(DurabilityError::CursorRegression { .. }) => {
                self.in_flight.start_recorded();
                Ok(())
            }
            Err(error) => Err(error),
        }
    }

    /// Forgets the join at `row` whose claim was not made.
    pub(super) fn unjoin_in_flight(&self, key: &str, row: u64) {
        self.in_flight.unjoin(key, row);
    }

    /// Forgets `key`'s claim settled at `settled` and moves the index start
    /// when enough rows lie behind it.
    pub(super) async fn leave_in_flight(&self, key: &str, settled: u64) {
        if self.in_flight.leave(key, settled) {
            self.advance_in_flight().await;
        }
    }

    /// Moves the recorded start as far as the claims held allow. The rows are
    /// already durable, so a failure here fails nothing: the old start stays
    /// valid and the next start reads more rows.
    async fn advance_in_flight(&self) {
        if let Err(error) = self.try_advance_in_flight().await {
            tracing::warn!(
                namespace = %self.namespace,
                %error,
                "the in-flight index start could not be recorded; the next start reads the \
                 rows since the previous record"
            );
        }
    }

    async fn try_advance_in_flight(&self) -> Result<(), DurabilityError> {
        let target = match self.in_flight.plan() {
            None => return Ok(()),
            Some(Advance::Record(target)) => target,
            Some(Advance::Carry { keys, at }) => {
                let rows = keys
                    .iter()
                    .map(|(key, open)| codec::encode_join(key, open.claim))
                    .collect::<Result<Vec<_>, _>>();
                let appended = match rows {
                    Ok(rows) => self.append_index_rows(rows, at).await,
                    Err(error) => Err(error),
                };
                let target = self
                    .in_flight
                    .carried(&keys, appended.as_ref().ok().copied(), at);
                appended?;
                let Some(target) = target else {
                    return Ok(());
                };
                target
            }
        };
        let Some(start) = self.in_flight.start() else {
            return Ok(());
        };
        match self
            .store
            .cas(
                &self.index_start_key(),
                start.saturating_add(1),
                target.saturating_add(1),
            )
            .await
        {
            Ok(()) => self.in_flight.recorded(target),
            // Another advance recorded first; take its start.
            Err(DurabilityError::CursorRegression { stored, .. }) => {
                self.in_flight.recorded(stored.saturating_sub(1));
            }
            Err(error) => return Err(error),
        }
        Ok(())
    }

    /// Appends `rows` to the index at the first row the stream has free at or
    /// after `expected`, answering where the first one landed.
    async fn append_index_rows(
        &self,
        rows: Vec<Vec<u8>>,
        expected: u64,
    ) -> Result<u64, DurabilityError> {
        let stream = self.index_stream();
        let mut expected = expected;
        let mut last_conflict = None;
        for _ in 0..INDEX_APPEND_ATTEMPTS {
            match self.store.append_all(&stream, rows.clone(), expected).await {
                Ok(_) => return Ok(expected),
                // Another writer took the row; the stream's head is later, so
                // the append moves there and never below its reservation.
                Err(DurabilityError::SequenceConflict { actual, .. }) if actual > expected => {
                    last_conflict = Some(actual);
                    expected = actual;
                }
                Err(error) => return Err(error),
            }
        }
        Err(DurabilityError::SequenceConflict {
            expected,
            actual: last_conflict.unwrap_or(expected),
        })
    }

    /// Reads the index from its recorded start, migrating a store that has
    /// none, and reads each named key's claim from its own stream: one page
    /// of two rows at the claim sequence the row carries, or, for a row
    /// written before the sequence was carried, the key's full look-up.
    async fn load_in_flight_index(
        &self,
    ) -> Result<(IndexState, Vec<(String, Option<ClaimEvidence>)>), DurabilityError> {
        let start = match self.store.read_value(&self.index_start_key()).await? {
            Some(recorded) if recorded > 0 => recorded - 1,
            _ => self.migrate_in_flight_index().await?,
        };
        let (named, tail) = self.read_index_rows(start).await?;
        let mut open = HashMap::new();
        let mut claims = Vec::new();
        for (key, (row, claim)) in named {
            let held = match claim {
                Some(claim) => self
                    .claim_at(&key, claim)
                    .await?
                    .map(|entry| (claim, entry)),
                None => self.claim_without_sequence(&key).await?,
            };
            let Some((claim, entry)) = held else {
                continue;
            };
            open.insert(key.clone(), OpenClaim { row, claim });
            claims.push((key, entry.evidence().cloned()));
        }
        let state = IndexState {
            start,
            tail,
            open,
            reserved: BTreeMap::new(),
            carrying: false,
            recorded: true,
        };
        Ok((state, claims))
    }

    /// The index rows from `start` to the stream's end: each named key with
    /// the last row naming it and that row's claim sequence, and the row after
    /// the last one read.
    async fn read_index_rows(
        &self,
        start: u64,
    ) -> Result<(BTreeMap<String, (u64, Option<u64>)>, u64), DurabilityError> {
        let stream = self.index_stream();
        let mut named = BTreeMap::new();
        let mut tail = start;
        loop {
            let page = self.store.read_from(&stream, tail, READ_BATCH_SIZE).await?;
            for row in &page {
                if row.sequence < tail {
                    return Err(DurabilityError::ConfigError(format!(
                        "stream {stream} answered sequence {} to a read from {tail}",
                        row.sequence
                    )));
                }
                let (key, claim) = codec::decode_join(&row.payload)?;
                named.insert(key, (row.sequence, claim));
                tail = row.sequence.checked_add(1).ok_or_else(|| {
                    DurabilityError::ConfigError(format!("stream {stream} sequence overflow"))
                })?;
            }
            if page.len() < READ_BATCH_SIZE {
                break;
            }
        }
        Ok((named, tail))
    }

    /// `key`'s claim at `claim` when it is still in flight, from one bounded
    /// read of two rows: the claim's and the one after it. The key's head and
    /// its history before `claim` are never read.
    ///
    /// No row at `claim` (the crash window after the index row and before
    /// the claim), a receipt on it, a tombstone there, or a row after it (a
    /// completion or release) answers `None`.
    ///
    /// # Errors
    /// Refuses a row at `claim` naming another key
    /// ([`DurabilityError::DedupCollision`]), a row answered at a sequence it
    /// was not read at, or a malformed record.
    async fn claim_at(&self, key: &str, claim: u64) -> Result<Option<DedupEntry>, DurabilityError> {
        let stream = self.stream_key_for(key);
        let rows = self.store.read_from(&stream, claim, CLAIM_PAGE).await?;
        let mut expected = claim;
        for row in &rows {
            if row.sequence != expected {
                return Err(DurabilityError::ConfigError(format!(
                    "stream {stream} answered sequence {} where {expected} was read",
                    row.sequence
                )));
            }
            expected = expected.saturating_add(1);
        }
        let Some(claimed) = rows.first() else {
            return Ok(None);
        };
        let record = DedupRecord::deserialize(&claimed.payload)?;
        if record.idempotency_key() != key {
            return Err(DurabilityError::DedupCollision {
                key: key.to_owned(),
            });
        }
        if rows.len() > 1 {
            return Ok(None);
        }
        Ok(record
            .into_active()
            .filter(|entry| entry.receipt().is_none()))
    }

    /// `key`'s claim and its sequence when it is in flight, looked up in
    /// full: for an index row written before the claim sequence was carried.
    async fn claim_without_sequence(
        &self,
        key: &str,
    ) -> Result<Option<(u64, DedupEntry)>, DurabilityError> {
        let snapshot = self.load_snapshot(&self.stream_key_for(key), key).await?;
        let claim = match snapshot.current.as_ref() {
            Some(entry) if entry.receipt().is_none() => snapshot.claim_sequence()?,
            _ => return Ok(None),
        };
        Ok(snapshot.current.map(|entry| (claim, entry)))
    }

    /// Builds the index of a store that has none: one walk of the namespace,
    /// one append of a row per claim in flight, and the start recorded at
    /// those rows. Answers the recorded start.
    async fn migrate_in_flight_index(&self) -> Result<u64, DurabilityError> {
        let mut keys = Vec::new();
        let mut walk = self.latest_records();
        while let Some(batch) = walk.next_batch().await? {
            for candidate in batch {
                if let DedupRecord::Active(entry) = candidate.record {
                    if entry.receipt().is_none() {
                        keys.push((entry.idempotency_key, candidate.sequence));
                    }
                }
            }
        }
        let start = if keys.is_empty() {
            // Nothing to name. A row left by an earlier walk that crashed
            // before its record is read once and dropped.
            0
        } else {
            let rows = keys
                .iter()
                .map(|(key, claim)| codec::encode_join(key, *claim))
                .collect::<Result<Vec<_>, _>>()?;
            self.append_index_rows(rows, 0).await?
        };
        let recorded = match self
            .store
            .cas(&self.index_start_key(), 0, start.saturating_add(1))
            .await
        {
            Ok(()) => start,
            // A concurrent first start recorded the index first.
            Err(DurabilityError::CursorRegression { stored, .. }) if stored > 0 => stored - 1,
            Err(error) => return Err(error),
        };
        tracing::info!(
            namespace = %self.namespace,
            claims = keys.len(),
            scanned = walk.scanned(),
            "the dedup namespace was walked once to build its in-flight index"
        );
        Ok(recorded)
    }
}