liminal-rs 0.14.1

//! A restart's reads of a durable partition are bounded by the rows since its recorded
//! head, never by the partition's whole history.

use crate::durability::{
    DurableChannel, MessageEnvelope, ScanCursor, ScanPage, StoredEntry, recover_durable_channel,
};

use std::collections::HashMap;
use std::error::Error;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};

use super::*;

const STREAM: &str = "inbox:0";

/// Rows appended after the last recorded head in the history tests.
const TAIL_ROWS: u64 = 5;

/// A store that keeps streams and values in memory and counts every row it hands out.
#[derive(Debug, Default)]
struct CountingStore {
    streams: Mutex<HashMap<String, Vec<StoredEntry>>>,
    values: Mutex<HashMap<String, u64>>,
    reads: Mutex<u64>,
    rows_read: Mutex<u64>,
}

#[async_trait::async_trait]
impl DurableStore for CountingStore {
    async fn apply_routed_batch(
        &self,
        _: &str,
        _: Vec<crate::durability::RoutedBatchItem>,
    ) -> Result<(), crate::durability::DurabilityError> {
        Err(crate::durability::DurabilityError::BatchUnsupported {
            store: std::any::type_name::<Self>(),
        })
    }

    async fn get_routed(
        &self,
        _: &str,
        _: &[u8],
    ) -> Result<Option<Vec<u8>>, crate::durability::DurabilityError> {
        Err(crate::durability::DurabilityError::BatchUnsupported {
            store: std::any::type_name::<Self>(),
        })
    }

    async fn append(
        &self,
        stream_key: &str,
        payload: Vec<u8>,
        expected_seq: u64,
    ) -> Result<u64, DurabilityError> {
        let mut streams = self.streams.lock().map_err(|_| lock_error())?;
        let stream = streams.entry(stream_key.to_owned()).or_default();
        let actual = to_u64(stream.len())?;
        if actual != expected_seq {
            return Err(DurabilityError::SequenceConflict {
                expected: expected_seq,
                actual,
            });
        }
        stream.push(StoredEntry {
            payload,
            sequence: actual,
            timestamp: 0,
        });
        drop(streams);
        Ok(actual)
    }

    async fn append_all(
        &self,
        stream_key: &str,
        payloads: Vec<Vec<u8>>,
        expected_seq: u64,
    ) -> Result<u64, DurabilityError> {
        let mut streams = self.streams.lock().map_err(|_| lock_error())?;
        let entries = streams.entry(stream_key.to_owned()).or_default();
        let actual = to_u64(entries.len())?;
        if actual != expected_seq {
            return Err(DurabilityError::SequenceConflict {
                expected: expected_seq,
                actual,
            });
        }
        let count = u64::try_from(payloads.len())
            .map_err(|_| DurabilityError::ConfigError("batch length exceeds u64".to_owned()))?;
        let next = actual
            .checked_add(count)
            .ok_or_else(|| DurabilityError::ConfigError("batch sequence overflow".to_owned()))?;
        let mut sequence = actual;
        entries.extend(payloads.into_iter().map(|payload| {
            let entry = StoredEntry {
                payload,
                sequence,
                timestamp: 0,
            };
            sequence += 1;
            entry
        }));
        drop(streams);
        Ok(next)
    }

    async fn read_from(
        &self,
        stream_key: &str,
        offset: u64,
        limit: usize,
    ) -> Result<Vec<StoredEntry>, DurabilityError> {
        let page: Vec<StoredEntry> = self
            .streams
            .lock()
            .map_err(|_| lock_error())?
            .get(stream_key)
            .map_or_else(Vec::new, |stream| {
                stream
                    .iter()
                    .skip_while(|stored| stored.sequence < offset)
                    .take(limit)
                    .cloned()
                    .collect()
            });
        *self.reads.lock().map_err(|_| lock_error())? += 1;
        *self.rows_read.lock().map_err(|_| lock_error())? += to_u64(page.len())?;
        Ok(page)
    }

    async fn cas(&self, key: &str, old_value: u64, new_value: u64) -> Result<(), DurabilityError> {
        let mut values = self.values.lock().map_err(|_| lock_error())?;
        let stored = values.get(key).copied().unwrap_or(0);
        if stored != old_value {
            return Err(DurabilityError::CursorRegression {
                stored,
                attempted: old_value,
            });
        }
        values.insert(key.to_owned(), new_value);
        drop(values);
        Ok(())
    }

    async fn read_stream_next_seq(&self, stream_key: &str) -> Result<Option<u64>, DurabilityError> {
        Err(DurabilityError::ConfigError(format!(
            "stream head read is unsupported by CountingStore for {stream_key}"
        )))
    }

    async fn read_value(&self, key: &str) -> Result<Option<u64>, DurabilityError> {
        Ok(self
            .values
            .lock()
            .map_err(|_| lock_error())?
            .get(key)
            .copied())
    }

    async fn scan_page(
        &self,
        prefix: &str,
        _: Option<ScanCursor>,
        _: usize,
    ) -> Result<ScanPage, DurabilityError> {
        Err(DurabilityError::ConfigError(format!(
            "head recovery never scans, and a scan of {prefix} was asked for"
        )))
    }

    async fn flush(&self) -> Result<(), DurabilityError> {
        Ok(())
    }
}

impl CountingStore {
    fn seed(&self, stream_key: &str, rows: u64) -> Result<(), DurabilityError> {
        let mut entries = Vec::new();
        for sequence in 0..rows {
            entries.push(StoredEntry {
                payload: Vec::new(),
                sequence,
                timestamp: 0,
            });
        }
        self.streams
            .lock()
            .map_err(|_| lock_error())?
            .insert(stream_key.to_owned(), entries);
        Ok(())
    }

    /// Seeds a log pruned below `first`: only rows `first..first + rows` are retained,
    /// and a read from an offset below `first` answers from the first retained row, as
    /// the engine's sequence-keyed read does.
    fn seed_retained(
        &self,
        stream_key: &str,
        first: u64,
        rows: u64,
    ) -> Result<(), DurabilityError> {
        let entries = (first..first + rows)
            .map(|sequence| StoredEntry {
                payload: Vec::new(),
                sequence,
                timestamp: 0,
            })
            .collect();
        self.streams
            .lock()
            .map_err(|_| lock_error())?
            .insert(stream_key.to_owned(), entries);
        Ok(())
    }

    fn record(&self, stream_key: &str, head: u64) -> Result<(), DurabilityError> {
        self.values
            .lock()
            .map_err(|_| lock_error())?
            .insert(head_checkpoint_key(stream_key), head);
        Ok(())
    }

    fn recorded(&self, stream_key: &str) -> Result<Option<u64>, DurabilityError> {
        self.values
            .lock()
            .map(|values| values.get(&head_checkpoint_key(stream_key)).copied())
            .map_err(|_| lock_error())
    }

    /// Clears the read counters and answers the (reads, rows) taken since the last call.
    fn take_counts(&self) -> Result<(u64, u64), DurabilityError> {
        let reads = std::mem::take(&mut *self.reads.lock().map_err(|_| lock_error())?);
        let rows = std::mem::take(&mut *self.rows_read.lock().map_err(|_| lock_error())?);
        Ok((reads, rows))
    }
}

/// The reads and rows a restart takes over a partition of `history` rows, where the
/// last `TAIL_ROWS` were appended after the recorded head.
fn restart_counts(history: u64) -> Result<(u64, u64), Box<dyn Error>> {
    let store = CountingStore::default();
    store.seed(STREAM, history)?;
    store.record(STREAM, history - TAIL_ROWS)?;
    store.take_counts()?;
    let next = block_on(recover_partition_head(&store, STREAM))?;
    assert_eq!(next, history);
    Ok(store.take_counts()?)
}

#[test]
fn a_restart_reads_the_same_whatever_the_history() -> Result<(), Box<dyn Error>> {
    let short = restart_counts(1_000)?;
    let long = restart_counts(100_000)?;

    // One point read verifies the record; one page holds the tail.
    assert_eq!(short, (2, 1 + TAIL_ROWS));
    assert_eq!(long, short);
    Ok(())
}

#[test]
fn a_missing_record_over_rows_is_walked_whole_and_recorded() -> Result<(), Box<dyn Error>> {
    let store = CountingStore::default();
    store.seed(STREAM, 3_000)?;

    let next = block_on(recover_partition_head(&store, STREAM))?;
    let (_, rows) = store.take_counts()?;

    assert_eq!(next, 3_000);
    assert_eq!(rows, 3_000);
    assert_eq!(store.recorded(STREAM)?, Some(3_000));

    let again = block_on(recover_partition_head(&store, STREAM))?;
    assert_eq!(again, 3_000);
    assert_eq!(store.take_counts()?, (2, 1));
    Ok(())
}

/// Rows a pruned log retains in the pruned-history tests: two full pages of
/// `READ_BATCH_SIZE`.
const RETAINED_ROWS: u64 = 2_048;

/// How a pruned log's recorded head stands before a restart.
#[derive(Clone, Copy, Debug)]
enum Head {
    /// Recorded just before the last `TAIL_ROWS` retained rows.
    Correct,
    /// Never recorded.
    Missing,
    /// Recorded past the end of the log.
    Refused,
}

/// The reads and rows a restart takes over a log of `history` logical rows pruned to
/// its last `RETAINED_ROWS`, with its head standing as `head`.
fn pruned_restart_counts(history: u64, head: Head) -> Result<(u64, u64), Box<dyn Error>> {
    let store = CountingStore::default();
    store.seed_retained(STREAM, history - RETAINED_ROWS, RETAINED_ROWS)?;
    match head {
        Head::Correct => store.record(STREAM, history - TAIL_ROWS)?,
        Head::Missing => {}
        Head::Refused => store.record(STREAM, history + TAIL_ROWS)?,
    }
    store.take_counts()?;
    let next = block_on(recover_partition_head(&store, STREAM))?;
    // A short tail past a correct head leaves the head where it was recorded.
    let recorded = match head {
        Head::Correct => history - TAIL_ROWS,
        Head::Missing | Head::Refused => history,
    };
    assert_eq!(
        (next, store.recorded(STREAM)?),
        (history, Some(recorded)),
        "a {head:?} head over a pruned log of {history} rows recovered the wrong head"
    );
    Ok(store.take_counts()?)
}

#[test]
fn a_pruned_log_with_its_head_restarts_from_the_head_whatever_the_history()
-> Result<(), Box<dyn Error>> {
    let short = pruned_restart_counts(10_000, Head::Correct)?;
    let long = pruned_restart_counts(100_000, Head::Correct)?;

    assert_eq!(
        (short, long),
        ((2, 1 + TAIL_ROWS), (2, 1 + TAIL_ROWS)),
        "a pruned log's restart from its recorded head read past the tail"
    );
    Ok(())
}

#[test]
fn a_pruned_log_missing_its_head_walks_each_retained_row_once() -> Result<(), Box<dyn Error>> {
    let short = pruned_restart_counts(10_000, Head::Missing)?;
    let long = pruned_restart_counts(100_000, Head::Missing)?;

    // Two full pages of retained rows, then one empty page that ends the walk.
    assert_eq!(
        (short, long),
        ((3, RETAINED_ROWS), (3, RETAINED_ROWS)),
        "a pruned log missing its head was walked by count from 0, rereading retained pages"
    );
    Ok(())
}

#[test]
fn a_pruned_log_refusing_its_head_walks_each_retained_row_once() -> Result<(), Box<dyn Error>> {
    let short = pruned_restart_counts(10_000, Head::Refused)?;
    let long = pruned_restart_counts(100_000, Head::Refused)?;

    // The point read that refuses the head, two full pages, one empty page.
    assert_eq!(
        (short, long),
        ((4, RETAINED_ROWS), (4, RETAINED_ROWS)),
        "a pruned log refusing its head was walked by count from 0, rereading retained pages"
    );
    Ok(())
}

#[test]
fn an_empty_partition_records_nothing() -> Result<(), Box<dyn Error>> {
    let store = CountingStore::default();

    let next = block_on(recover_partition_head(&store, STREAM))?;

    assert_eq!(next, 0);
    assert_eq!(store.recorded(STREAM)?, None);
    Ok(())
}

#[test]
fn a_record_beyond_the_log_is_refused_and_replaced() -> Result<(), Box<dyn Error>> {
    let store = CountingStore::default();
    store.seed(STREAM, 40)?;
    store.record(STREAM, 900)?;

    let refusal = block_on(verify(&store, STREAM, 900))?;
    assert_eq!(
        refusal,
        Some(HeadCheckpointRefusal::BeyondLog { covered: 900 })
    );

    let next = block_on(recover_partition_head(&store, STREAM))?;
    assert_eq!(next, 40);
    assert_eq!(store.recorded(STREAM)?, Some(40));
    Ok(())
}

#[test]
fn a_record_naming_a_misplaced_row_is_refused() -> Result<(), Box<dyn Error>> {
    let store = CountingStore::default();
    store.seed(STREAM, 10)?;
    let mut streams = store.streams.lock().map_err(|_| lock_error())?;
    let row = streams
        .get_mut(STREAM)
        .and_then(|stream| stream.get_mut(6))
        .ok_or("the seeded stream has no seventh row")?;
    row.sequence = 99;
    drop(streams);

    let refusal = block_on(verify(&store, STREAM, 7))?;

    assert_eq!(
        refusal,
        Some(HeadCheckpointRefusal::WrongSequence {
            covered: 7,
            found: 99
        })
    );
    Ok(())
}

#[test]
fn a_long_tail_is_recorded_again_at_restart() -> Result<(), Box<dyn Error>> {
    let store = CountingStore::default();
    let history = 2 * HEAD_CHECKPOINT_EVERY;
    store.seed(STREAM, history)?;
    store.record(STREAM, 1)?;

    let next = block_on(recover_partition_head(&store, STREAM))?;

    assert_eq!(next, history);
    assert_eq!(store.recorded(STREAM)?, Some(history));
    Ok(())
}

#[test]
fn publishes_record_the_head_by_count() -> Result<(), Box<dyn Error>> {
    let store = Arc::new(CountingStore::default());
    let mut channel = block_on(recover_durable_channel("inbox", 1, store.clone()))?;

    for index in 0..HEAD_CHECKPOINT_EVERY - 1 {
        block_on(channel.publish(&message(index)))?;
    }
    assert_eq!(store.recorded(STREAM)?, None);

    block_on(channel.publish(&message(HEAD_CHECKPOINT_EVERY)))?;
    assert_eq!(store.recorded(STREAM)?, Some(HEAD_CHECKPOINT_EVERY));

    store.take_counts()?;
    let restarted: DurableChannel = block_on(recover_durable_channel("inbox", 1, store.clone()))?;
    assert_eq!(
        restarted.next_expected_sequence(0),
        Some(HEAD_CHECKPOINT_EVERY)
    );
    assert_eq!(store.take_counts()?, (2, 1));
    Ok(())
}

#[test]
fn a_head_already_recorded_further_is_left_alone() -> Result<(), Box<dyn Error>> {
    let store = CountingStore::default();
    store.record(STREAM, 500)?;

    block_on(advance_head(&store, STREAM, 300))?;

    assert_eq!(store.recorded(STREAM)?, Some(500));
    Ok(())
}

fn message(index: u64) -> MessageEnvelope {
    MessageEnvelope {
        payload: format!("message-{index}").into_bytes(),
        causal_context: None,
        timestamp: 1,
        publisher_id: "publisher".to_owned(),
        idempotency_key: None,
    }
}

fn block_on<F: Future>(future: F) -> F::Output {
    let mut context = Context::from_waker(Waker::noop());
    let mut future = Box::pin(future);
    loop {
        match Future::poll(Pin::as_mut(&mut future), &mut context) {
            Poll::Ready(output) => return output,
            Poll::Pending => std::thread::yield_now(),
        }
    }
}

fn to_u64(len: usize) -> Result<u64, DurabilityError> {
    u64::try_from(len).map_err(|error| DurabilityError::ConfigError(format!("length: {error}")))
}

fn lock_error() -> DurabilityError {
    DurabilityError::StoreError(haematite::ApiError::Storage(
        haematite::DatabaseError::IoError(std::io::Error::other("counting store lock poisoned")),
    ))
}