xz-memory-engine 0.2.0

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

use xz_memory_core::StoreError;
use xz_memory_core::traits::store::EntryStore;
use xz_memory_core::types::entry::*;

/// In-memory [`EntryStore`] implementation backed by a `HashMap`.
///
/// Partitioned entries are stored in-memory using `tokio::sync::Mutex`.
/// Not suitable for production persistence — data is lost on drop.
pub struct InMemoryEntryStore {
    data: Mutex<HashMap<String, Vec<Entry>>>,
}

impl InMemoryEntryStore {
    /// Create a new, empty `InMemoryEntryStore`.
    pub fn new() -> Self {
        InMemoryEntryStore { data: Mutex::new(HashMap::new()) }
    }
}

#[async_trait]
impl EntryStore for InMemoryEntryStore {
    async fn append(&self, entry: Entry) -> Result<(), StoreError> {
        let mut data = self.data.lock().await;
        data.entry(entry.partition.clone()).or_default().push(entry);
        Ok(())
    }

    async fn query(
        &self,
        partition: &str,
        range: &TimeRange,
        opts: &QueryOptions,
    ) -> Result<Vec<Entry>, StoreError> {
        let data = self.data.lock().await;
        let mut entries: Vec<Entry> = data.get(partition).map_or(vec![], |v| {
            v.iter()
                .filter(|e| {
                    let after_start = range.start.is_none_or(|s| e.recorded_at >= s);
                    let before_end = range.end.is_none_or(|e2| e.recorded_at <= e2);
                    after_start && before_end
                })
                .cloned()
                .collect()
        });
        match opts.sort {
            SortOrder::Ascending => entries.sort_by_key(|e| e.recorded_at),
            SortOrder::Descending => entries.sort_by_key(|b| std::cmp::Reverse(b.recorded_at)),
        }
        entries.truncate(opts.limit);
        Ok(entries)
    }

    async fn evict(&self, partition: &str, keep: usize) -> Result<usize, StoreError> {
        let mut data = self.data.lock().await;
        if let Some(entries) = data.get_mut(partition) {
            if entries.len() <= keep {
                return Ok(0);
            }
            let remove_count = entries.len() - keep;
            entries.drain(..remove_count);
            Ok(remove_count)
        } else {
            Ok(0)
        }
    }

    async fn delete(&self, id: &str) -> Result<(), StoreError> {
        let mut data = self.data.lock().await;
        for entries in data.values_mut() {
            entries.retain(|e| e.id != id);
        }
        Ok(())
    }

    async fn clear_partition(&self, partition: &str) -> Result<(), StoreError> {
        let mut data = self.data.lock().await;
        data.remove(partition);
        Ok(())
    }
}

impl Default for InMemoryEntryStore {
    fn default() -> Self {
        Self::new()
    }
}