millipede_storage_memory/
kvs.rs

1use bytes::Bytes;
2use millipede_core::storage::{
3    KeyInfo, KeyList, KeyValueStore, KvEntry, ListKeysOptions, StorageResult,
4};
5use std::{collections::HashMap, sync::Mutex};
6
7/// An in-process byte-oriented key-value store.
8pub struct MemoryKeyValueStore {
9    name: String,
10    inner: Mutex<HashMap<String, KvEntry>>,
11}
12
13impl MemoryKeyValueStore {
14    /// Creates an empty key-value store with the supplied name.
15    #[must_use]
16    pub fn new(name: impl Into<String>) -> Self {
17        Self {
18            name: name.into(),
19            inner: Mutex::new(HashMap::new()),
20        }
21    }
22
23    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<String, KvEntry>> {
24        let _store_name = &self.name;
25        // A panic while holding this lock is a programming bug, so poisoning is unrecoverable.
26        self.inner
27            .lock()
28            .expect("MemoryKeyValueStore mutex poisoned")
29    }
30
31    pub(crate) fn clear(&self) {
32        self.lock().clear();
33    }
34}
35
36#[async_trait::async_trait]
37impl KeyValueStore for MemoryKeyValueStore {
38    async fn get_bytes(&self, key: &str) -> StorageResult<Option<KvEntry>> {
39        Ok(self.lock().get(key).cloned())
40    }
41
42    async fn set_bytes(&self, key: &str, value: Bytes, content_type: &str) -> StorageResult<()> {
43        self.lock().insert(
44            key.to_owned(),
45            KvEntry {
46                key: key.to_owned(),
47                value,
48                content_type: content_type.to_owned(),
49            },
50        );
51        Ok(())
52    }
53
54    async fn delete(&self, key: &str) -> StorageResult<()> {
55        self.lock().remove(key);
56        Ok(())
57    }
58
59    /// Lists keys in lexical order after the exclusive cursor.
60    ///
61    /// A zero limit returns an empty, non-truncated page. This avoids claiming
62    /// that a caller can continue when the page contains no key to use as its
63    /// next exclusive cursor.
64    async fn list_keys(&self, opts: ListKeysOptions) -> StorageResult<KeyList> {
65        if opts.limit == Some(0) {
66            return Ok(KeyList {
67                keys: Vec::new(),
68                is_truncated: false,
69                next_exclusive_start_key: None,
70            });
71        }
72
73        let entries = self.lock();
74        let mut keys: Vec<_> = entries.keys().cloned().collect();
75        keys.sort_unstable();
76        let start = opts.exclusive_start_key.as_deref();
77        let mut filtered = keys
78            .into_iter()
79            .filter(|key| start.is_none_or(|start| key.as_str() > start));
80        let limit = opts.limit.unwrap_or(usize::MAX);
81        let selected: Vec<_> = filtered.by_ref().take(limit).collect();
82        let is_truncated = filtered.next().is_some();
83        let next_exclusive_start_key = is_truncated.then(|| selected.last().cloned()).flatten();
84        let keys = selected
85            .into_iter()
86            .map(|key| KeyInfo {
87                size: entries.get(&key).expect("selected key exists").value.len() as u64,
88                key,
89            })
90            .collect();
91        Ok(KeyList {
92            keys,
93            is_truncated,
94            next_exclusive_start_key,
95        })
96    }
97}
millipede_storage_memory/kvs.rs

millipede_storage_memory/
kvs.rs
