liminal-server 0.23.1

//! The store-backed fixtures the registry snapshot pins share (DUR-019 R1,
//! R2, R3): a recording store, a log capture, and the flat, routed and
//! snapshot readers.

use std::collections::BTreeSet;
use std::error::Error;
use std::io;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

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

use super::registry::REGISTRY_STREAM_KEY;
use super::registry_snapshot::{RegistrySnapshot, SNAPSHOT_KEY};
use super::routed_stream_codec::head_key;
use super::tests_memory_store::MemoryStore;

/// The registry's row schema version (registry.rs `SCHEMA_VERSION`).
pub(super) const ROW_SCHEMA: u8 = 1;

pub(super) type TestResult<T = ()> = Result<T, Box<dyn Error>>;

pub(super) const KEY: &str = REGISTRY_STREAM_KEY;

/// What a [`Recording`] store saw read.
#[derive(Debug, Default)]
pub(super) struct ReadCounts {
    /// `read_from` and `read_at` calls on the registry stream.
    pub(super) row_reads: u64,
    /// Rows those calls returned.
    pub(super) rows: u64,
    /// `read_stream_next_seq` calls on the registry stream.
    pub(super) head_reads: u64,
    /// Every routed key read on the registry route, in order.
    pub(super) routed_keys: Vec<Vec<u8>>,
}

/// Passes every call through, recording each routed batch, each flush and
/// each registry read; armed, it refuses the next batch carrying the
/// snapshot key before that batch applies.
#[derive(Debug)]
pub(super) struct Recording {
    inner: Arc<dyn DurableStore>,
    batches: Mutex<Vec<(String, Vec<RoutedBatchItem>)>>,
    pub(super) flushes: AtomicU64,
    reads: Mutex<ReadCounts>,
    pub(super) refuse_snapshot: AtomicBool,
}

impl Recording {
    pub(super) fn new(inner: Arc<dyn DurableStore>) -> Arc<Self> {
        Arc::new(Self {
            inner,
            batches: Mutex::new(Vec::new()),
            flushes: AtomicU64::new(0),
            reads: Mutex::new(ReadCounts::default()),
            refuse_snapshot: AtomicBool::new(false),
        })
    }

    pub(super) fn shared(self: &Arc<Self>) -> Arc<dyn DurableStore> {
        Arc::clone(self) as Arc<dyn DurableStore>
    }

    pub(super) fn take_batches(&self) -> TestResult<Vec<(String, Vec<RoutedBatchItem>)>> {
        Ok(std::mem::take(
            &mut *self.batches.lock().map_err(|_| "batch log poisoned")?,
        ))
    }

    pub(super) fn take_reads(&self) -> TestResult<ReadCounts> {
        Ok(std::mem::take(
            &mut *self.reads.lock().map_err(|_| "read log poisoned")?,
        ))
    }

    fn note(&self, work: impl FnOnce(&mut ReadCounts)) -> Result<(), DurabilityError> {
        let mut reads = self
            .reads
            .lock()
            .map_err(|_| DurabilityError::ConfigError("read log poisoned".to_owned()))?;
        work(&mut reads);
        drop(reads);
        Ok(())
    }
}

#[async_trait::async_trait]
impl DurableStore for Recording {
    async fn apply_routed_batch(
        &self,
        route_key: &str,
        items: Vec<RoutedBatchItem>,
    ) -> Result<(), DurabilityError> {
        self.batches
            .lock()
            .map_err(|_| DurabilityError::ConfigError("batch log poisoned".to_owned()))?
            .push((route_key.to_owned(), items.clone()));
        if items.iter().any(|item| item.key == SNAPSHOT_KEY)
            && self.refuse_snapshot.swap(false, Ordering::SeqCst)
        {
            return Err(DurabilityError::ConfigError(
                "the registry snapshot batch was refused before it applied".to_owned(),
            ));
        }
        self.inner.apply_routed_batch(route_key, items).await
    }

    async fn get_routed(
        &self,
        route_key: &str,
        key: &[u8],
    ) -> Result<Option<Vec<u8>>, DurabilityError> {
        if route_key == KEY {
            self.note(|reads| reads.routed_keys.push(key.to_vec()))?;
        }
        self.inner.get_routed(route_key, key).await
    }

    async fn append(
        &self,
        stream_key: &str,
        payload: Vec<u8>,
        expected_seq: u64,
    ) -> Result<u64, DurabilityError> {
        self.inner.append(stream_key, payload, expected_seq).await
    }

    async fn append_all(
        &self,
        stream_key: &str,
        payloads: Vec<Vec<u8>>,
        expected_seq: u64,
    ) -> Result<u64, DurabilityError> {
        self.inner
            .append_all(stream_key, payloads, expected_seq)
            .await
    }

    async fn read_from(
        &self,
        stream_key: &str,
        offset: u64,
        limit: usize,
    ) -> Result<Vec<StoredEntry>, DurabilityError> {
        let rows = self.inner.read_from(stream_key, offset, limit).await?;
        if stream_key == KEY {
            let count = u64::try_from(rows.len()).unwrap_or(u64::MAX);
            self.note(|reads| {
                reads.row_reads += 1;
                reads.rows += count;
            })?;
        }
        Ok(rows)
    }

    async fn read_at(
        &self,
        stream_key: &str,
        sequence: u64,
    ) -> Result<Option<StoredEntry>, DurabilityError> {
        let row = self.inner.read_at(stream_key, sequence).await?;
        if stream_key == KEY {
            let count = u64::from(row.is_some());
            self.note(|reads| {
                reads.row_reads += 1;
                reads.rows += count;
            })?;
        }
        Ok(row)
    }

    async fn cas(&self, key: &str, old_value: u64, new_value: u64) -> Result<(), DurabilityError> {
        self.inner.cas(key, old_value, new_value).await
    }

    async fn read_stream_next_seq(&self, stream_key: &str) -> Result<Option<u64>, DurabilityError> {
        if stream_key == KEY {
            self.note(|reads| reads.head_reads += 1)?;
        }
        self.inner.read_stream_next_seq(stream_key).await
    }

    async fn read_value(&self, key: &str) -> Result<Option<u64>, DurabilityError> {
        self.inner.read_value(key).await
    }

    async fn scan_page(
        &self,
        prefix: &str,
        cursor: Option<ScanCursor>,
        limit: usize,
    ) -> Result<ScanPage, DurabilityError> {
        self.inner.scan_page(prefix, cursor, limit).await
    }

    async fn flush(&self) -> Result<(), DurabilityError> {
        self.flushes.fetch_add(1, Ordering::SeqCst);
        self.inner.flush().await
    }

    async fn truncate_before(
        &self,
        stream_key: &str,
        before_seq: u64,
    ) -> Result<u64, DurabilityError> {
        self.inner.truncate_before(stream_key, before_seq).await
    }

    async fn scan_routed(
        &self,
        route_key: &str,
        from: &[u8],
        to: &[u8],
    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, DurabilityError> {
        self.inner.scan_routed(route_key, from, to).await
    }

    async fn remove_routed_range(
        &self,
        route_key: &str,
        from: &[u8],
        to: &[u8],
    ) -> Result<u64, DurabilityError> {
        self.inner.remove_routed_range(route_key, from, to).await
    }

    fn store_counts(&self) -> Option<liminal::durability::StoreCountsHandle> {
        self.inner.store_counts()
    }
}

/// Collects what the `tracing` subscriber writes during one call.
#[derive(Clone, Default)]
struct Diagnostics(Arc<Mutex<Vec<u8>>>);

impl io::Write for Diagnostics {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        self.0
            .lock()
            .map_err(|_| io::Error::other("diagnostic buffer poisoned"))?
            .extend_from_slice(bytes);
        Ok(bytes.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for Diagnostics {
    type Writer = Self;

    fn make_writer(&'a self) -> Self::Writer {
        self.clone()
    }
}

/// Runs `work` and answers what it returned with every line it logged at
/// warn or above.
pub(super) fn logged<T>(work: impl FnOnce() -> T) -> TestResult<(T, String)> {
    let diagnostics = Diagnostics::default();
    let subscriber = tracing_subscriber::fmt()
        .without_time()
        .with_ansi(false)
        .with_max_level(tracing::Level::WARN)
        .with_writer(diagnostics.clone())
        .finish();
    let answer = tracing::subscriber::with_default(subscriber, work);
    let bytes = diagnostics
        .0
        .lock()
        .map_err(|_| "diagnostic buffer poisoned")?
        .clone();
    Ok((answer, String::from_utf8(bytes)?))
}

pub(super) fn memory() -> Arc<dyn DurableStore> {
    Arc::new(MemoryStore::default())
}

/// The registry row payload an old install wrote for `id`.
fn payload(id: u64) -> TestResult<Vec<u8>> {
    Ok(serde_json::to_vec(&serde_json::json!({
        "schema_version": 1,
        "conversation_id": id,
    }))?)
}

/// Appends flat rows the way an install from before routing wrote them, and
/// answers the native head.
pub(super) fn append_flat(store: &dyn DurableStore, payloads: Vec<Vec<u8>>) -> TestResult<u64> {
    let start = block_on(store.read_stream_next_seq(KEY))??.unwrap_or(0);
    Ok(block_on(store.append_all(KEY, payloads, start))??)
}

pub(super) fn flat_ids(
    store: &dyn DurableStore,
    ids: impl Iterator<Item = u64>,
) -> TestResult<u64> {
    append_flat(store, ids.map(payload).collect::<TestResult<Vec<_>>>()?)
}

/// The stored snapshot, decoded, with the hash of its bytes.
pub(super) fn stored_snapshot(
    store: &dyn DurableStore,
) -> TestResult<Option<(RegistrySnapshot, Hash)>> {
    block_on(store.get_routed(KEY, SNAPSHOT_KEY))??
        .map(|bytes| {
            Ok((
                RegistrySnapshot::decode(&bytes, ROW_SCHEMA)?,
                Hash::of(&bytes),
            ))
        })
        .transpose()
}

/// The snapshot item of a batch, decoded, with its staged witness.
pub(super) fn snapshot_in(
    items: &[RoutedBatchItem],
) -> TestResult<(RegistrySnapshot, Option<Hash>)> {
    let item = items
        .iter()
        .find(|item| item.key == SNAPSHOT_KEY)
        .ok_or("the batch carries no snapshot")?;
    Ok((
        RegistrySnapshot::decode(&item.value, ROW_SCHEMA)?,
        item.expected,
    ))
}

/// The routed row keys among `keys`: everything but the snapshot and head.
pub(super) fn row_keys(keys: &[Vec<u8>]) -> TestResult<Vec<Vec<u8>>> {
    let head = head_key(KEY)?;
    Ok(keys
        .iter()
        .filter(|key| key.as_slice() != SNAPSHOT_KEY && **key != head)
        .cloned()
        .collect())
}

pub(super) fn ascending(range: std::ops::Range<u64>) -> BTreeSet<u64> {
    range.collect()
}