liminal-rs 0.14.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 one batch per stream in the fixture builder.

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

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

type RoutedValues =
    std::collections::BTreeMap<(String, Vec<u8>), (Vec<u8>, Option<std::time::Instant>)>;

/// 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>,
    routed: Mutex<RoutedValues>,
}

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 apply_routed_batch(
        &self,
        route_key: &str,
        items: Vec<RoutedBatchItem>,
    ) -> Result<(), DurabilityError> {
        if items.is_empty() {
            return Ok(());
        }
        let now = std::time::Instant::now();
        let mut keys = std::collections::BTreeSet::new();
        let mut prepared = Vec::with_capacity(items.len());
        for item in items {
            if item.key.is_empty() || !keys.insert(item.key.clone()) {
                return Err(DurabilityError::ConfigError(
                    "routed batch has an empty or repeated key".into(),
                ));
            }
            let expires = item
                .ttl
                .map(|ttl| {
                    now.checked_add(ttl).ok_or_else(|| {
                        DurabilityError::ConfigError("routed expiry exceeds the clock range".into())
                    })
                })
                .transpose()?;
            prepared.push((item, expires));
        }
        let mut values = self
            .routed
            .lock()
            .map_err(|_| DurabilityError::ConfigError("routed memory lock poisoned".into()))?;
        for (item, _) in &prepared {
            let actual = values
                .get(&(route_key.to_owned(), item.key.clone()))
                .filter(|(_, expires)| expires.is_none_or(|expiry| expiry > now))
                .map(|(value, _)| haematite::tree::Hash::of(value));
            if actual != item.expected {
                return Err(DurabilityError::ConfigError(format!(
                    "routed batch content-hash mismatch: expected {:?}, actual {actual:?}",
                    item.expected
                )));
            }
        }
        for (item, expires) in prepared {
            values.insert((route_key.to_owned(), item.key), (item.value, expires));
        }
        drop(values);
        Ok(())
    }

    async fn get_routed(
        &self,
        route_key: &str,
        key: &[u8],
    ) -> Result<Option<Vec<u8>>, DurabilityError> {
        let values = self
            .routed
            .lock()
            .map_err(|_| DurabilityError::ConfigError("routed memory lock poisoned".into()))?;
        let now = std::time::Instant::now();
        Ok(values
            .get(&(route_key.to_owned(), key.to_vec()))
            .filter(|(_, expires)| expires.is_none_or(|expiry| expiry > now))
            .map(|(value, _)| value.clone()))
    }

    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_stream_next_seq(&self, stream_key: &str) -> Result<Option<u64>, DurabilityError> {
        let streams = self.streams.lock().map_err(|_| lock_error())?;
        streams
            .get(stream_key)
            .filter(|stream| !stream.is_empty())
            .map(|stream| to_u64(stream.len()))
            .transpose()
    }

    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())?;
        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(())
    }
}

#[cfg(not(test))]
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
    }
}

fn page_in_memory_streams(
    streams: &std::collections::HashMap<String, Vec<StoredEntry>>,
    prefix: &str,
    cursor: Option<&ScanCursor>,
    limit: usize,
) -> Result<ScanPage, DurabilityError> {
    if limit == 0 {
        return Err(DurabilityError::ConfigError(
            "scan_page limit must be nonzero".to_owned(),
        ));
    }
    let mut stream_keys: Vec<&String> = streams
        .keys()
        .filter(|stream_key| stream_key.starts_with(prefix))
        .collect();
    stream_keys.sort_unstable();

    let mut entries: Vec<StoredEntry> = Vec::new();
    let mut resume_at: Option<ScanCursor> = None;
    for stream_key in stream_keys {
        let Some(remaining) = limit.checked_sub(entries.len()).filter(|rest| *rest > 0) else {
            break;
        };
        let offset = match cursor {
            Some(cursor) if *stream_key < cursor.stream_key => continue,
            Some(cursor) if *stream_key == cursor.stream_key => cursor.next_sequence,
            _ => 0,
        };
        let Some(stream) = streams.get(stream_key) else {
            continue;
        };
        let page: Vec<StoredEntry> = stream
            .iter()
            .filter(|entry| entry.sequence >= offset)
            .take(remaining)
            .cloned()
            .collect();
        if let Some(last) = page.last() {
            let next_sequence = last.sequence.checked_add(1).ok_or_else(|| {
                DurabilityError::ConfigError(format!(
                    "fake scan page sequence overflow for stream {stream_key}"
                ))
            })?;
            resume_at = Some(ScanCursor {
                stream_key: stream_key.clone(),
                next_sequence,
            });
        }
        entries.extend(page);
    }

    let next = if entries.len() == limit {
        resume_at
    } else {
        None
    };
    Ok(ScanPage { entries, next })
}