liminal-rs 0.11.0

//! A store that keeps streams and values in memory, for tests: the same
//! answers as the durable stores and no disk. A test that needs a large
//! seed builds it here through the real code, then lands it on the durable
//! store under test with [`MemoryStore::land_on`].

use std::collections::HashMap;
use std::sync::Mutex;

use super::{DurabilityError, DurableStore, ScanCursor, ScanPage, StoredEntry};

/// A store that keeps streams and values in memory.
#[derive(Debug, Default)]
pub struct MemoryStore {
    streams: Mutex<HashMap<String, Vec<StoredEntry>>>,
    values: Mutex<HashMap<String, u64>>,
    flushes: Mutex<u64>,
}

pub fn lock_error() -> DurabilityError {
    DurabilityError::ConfigError("memory store lock poisoned".to_owned())
}

fn to_u64(value: usize) -> Result<u64, DurabilityError> {
    u64::try_from(value).map_err(|_| DurabilityError::ConfigError("length overflow".to_owned()))
}

#[async_trait::async_trait]
impl DurableStore for MemoryStore {
    async fn append(
        &self,
        stream_key: &str,
        payload: Vec<u8>,
        expected_seq: u64,
    ) -> Result<u64, DurabilityError> {
        let next = self
            .append_all(stream_key, vec![payload], expected_seq)
            .await?;
        Ok(next - 1)
    }

    async fn read_from(
        &self,
        stream_key: &str,
        offset: u64,
        limit: usize,
    ) -> Result<Vec<StoredEntry>, DurabilityError> {
        let streams = self.streams.lock().map_err(|_| lock_error())?;
        Ok(streams
            .get(stream_key)
            .map(|stream| {
                stream
                    .iter()
                    .filter(|entry| entry.sequence >= offset)
                    .take(limit)
                    .cloned()
                    .collect()
            })
            .unwrap_or_default())
    }

    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 stream = streams.entry(stream_key.to_owned()).or_default();
        let mut next = to_u64(stream.len())?;
        if next != expected_seq {
            return Err(DurabilityError::SequenceConflict {
                expected: expected_seq,
                actual: next,
            });
        }
        for payload in payloads {
            stream.push(StoredEntry {
                payload,
                sequence: next,
                timestamp: 0,
            });
            next += 1;
        }
        drop(streams);
        Ok(next)
    }

    async fn read_at(
        &self,
        stream_key: &str,
        sequence: u64,
    ) -> Result<Option<StoredEntry>, DurabilityError> {
        let streams = self.streams.lock().map_err(|_| lock_error())?;
        Ok(streams
            .get(stream_key)
            .and_then(|stream| stream.iter().find(|entry| entry.sequence == sequence))
            .cloned())
    }

    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_value(&self, key: &str) -> Result<Option<u64>, DurabilityError> {
        let values = self.values.lock().map_err(|_| lock_error())?;
        Ok(values.get(key).copied())
    }

    async fn scan_page(
        &self,
        prefix: &str,
        cursor: Option<ScanCursor>,
        limit: usize,
    ) -> Result<ScanPage, DurabilityError> {
        let streams = self.streams.lock().map_err(|_| lock_error())?;
        crate::durability::store::page_in_memory_streams(&streams, prefix, cursor.as_ref(), limit)
    }

    async fn flush(&self) -> Result<(), DurabilityError> {
        let mut flushes = self.flushes.lock().map_err(|_| lock_error())?;
        *flushes += 1;
        drop(flushes);
        Ok(())
    }
}

impl MemoryStore {
    /// Writes every stream onto `target` with one `append_all` each and every
    /// value with one `cas`, then flushes once: the fewest durable operations
    /// the store's trait allows, since no operation spans streams.
    ///
    /// # Errors
    ///
    /// Returns the target's error, or [`DurabilityError::ConfigError`] when a
    /// stream lands a different number of rows than it holds.
    pub async fn land_on(&self, target: &dyn DurableStore) -> Result<(), DurabilityError> {
        let streams: Vec<(String, Vec<Vec<u8>>)> = {
            let streams = self.streams.lock().map_err(|_| lock_error())?;
            streams
                .iter()
                .map(|(key, entries)| {
                    (
                        key.clone(),
                        entries.iter().map(|entry| entry.payload.clone()).collect(),
                    )
                })
                .collect()
        };
        for (key, payloads) in streams {
            let rows = to_u64(payloads.len())?;
            let next = target.append_all(&key, payloads, 0).await?;
            if next != rows {
                return Err(DurabilityError::ConfigError(format!(
                    "stream {key} landed {next} rows of {rows}"
                )));
            }
        }
        let values: Vec<(String, u64)> = {
            let values = self.values.lock().map_err(|_| lock_error())?;
            values
                .iter()
                .map(|(key, value)| (key.clone(), *value))
                .collect()
        };
        for (key, value) in values {
            target.cas(&key, 0, value).await?;
        }
        target.flush().await
    }
}