liminal-rs 0.16.0

//! Authoritative idempotency-key lookup reads a channel's partitions page by page,
//! without changing admission state and without holding the channel's append lock.
use std::sync::Arc;

use super::{AppendResolution, DurableChannel, MessageEnvelope};
use crate::durability::head_checkpoint::READ_BATCH_SIZE;
use crate::durability::{DurabilityError, DurableStore};

/// The streams a key lookup reads, taken from a channel so the lookup can run
/// after the channel's append lock is released.
#[derive(Debug)]
pub struct KeyLookup {
    stream_keys: Vec<String>,
    store: Arc<dyn DurableStore>,
}

impl DurableChannel {
    /// The streams that answer for this channel's idempotency keys.
    ///
    /// A key is looked up only once its pending append is settled, and a held
    /// key is never appended again, so rows appended while the lookup reads
    /// cannot carry it: the lookup needs no lock.
    #[must_use]
    pub fn key_lookup(&self) -> KeyLookup {
        KeyLookup {
            stream_keys: (0..self.partition_count())
                .map(|partition| self.stream_key_for(partition))
                .collect(),
            store: Arc::clone(&self.store),
        }
    }

    /// Find the original durable position of a producer's key across
    /// partitions, reading each partition from its position in `from`.
    ///
    /// # Errors
    /// Propagates authoritative read and envelope decode errors. A failed read
    /// never establishes absence.
    pub async fn find_idempotency_key(
        &self,
        key: &str,
        from: &[u64],
    ) -> Result<AppendResolution, DurabilityError> {
        self.key_lookup().find(key, from).await
    }
}

impl KeyLookup {
    /// Find the original durable position of a producer's key, reading each
    /// partition from its position in `from` in pages of the recovery read
    /// size.
    ///
    /// `from` is where each partition stood when the key was claimed, which
    /// is before the key's row could be appended, so no row before it can
    /// carry the key. A partition `from` does not name is read from its first
    /// row.
    ///
    /// # Errors
    /// Propagates authoritative read and envelope decode errors. A failed read
    /// never establishes absence.
    pub async fn find(&self, key: &str, from: &[u64]) -> Result<AppendResolution, DurabilityError> {
        for (partition, stream_key) in self.stream_keys.iter().enumerate() {
            let mut sequence = from.get(partition).copied().unwrap_or(0);
            loop {
                let page = self
                    .store
                    .read_from(stream_key, sequence, READ_BATCH_SIZE)
                    .await?;
                for row in &page {
                    if row.sequence != sequence {
                        return Err(DurabilityError::ConfigError(format!(
                            "stream {stream_key} read at {sequence} returned sequence {}",
                            row.sequence
                        )));
                    }
                    let envelope = MessageEnvelope::deserialize(&row.payload)?;
                    if envelope.idempotency_key.as_deref() == Some(key) {
                        return Ok(AppendResolution::Committed {
                            position: row.sequence,
                        });
                    }
                    sequence = sequence.checked_add(1).ok_or_else(|| {
                        DurabilityError::ConfigError(format!(
                            "stream {stream_key} sequence overflow during key lookup"
                        ))
                    })?;
                }
                if page.len() < READ_BATCH_SIZE {
                    break;
                }
            }
        }
        Ok(AppendResolution::Absent)
    }
}