linera_views/backends/
memory.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Implements [`crate::store::KeyValueDatabase`] in memory.
5
6use std::{
7    collections::BTreeMap,
8    sync::{Arc, LazyLock, Mutex, RwLock},
9};
10
11use serde::{Deserialize, Serialize};
12use thiserror::Error;
13
14#[cfg(with_testing)]
15use crate::store::TestKeyValueDatabase;
16use crate::{
17    batch::{Batch, WriteOperation},
18    common::get_key_range_for_prefix,
19    store::{
20        KeyValueDatabase, KeyValueStoreError, ReadableKeyValueStore, WithError,
21        WritableKeyValueStore,
22    },
23};
24
25/// The initial configuration of the system
26#[derive(Debug, Clone, Deserialize, Serialize)]
27pub struct MemoryStoreConfig {
28    /// Whether a namespace should be immediately cleaned up from memory when the
29    /// connection object is dropped.
30    pub kill_on_drop: bool,
31}
32
33/// The values in a partition.
34type MemoryStoreMap = BTreeMap<Vec<u8>, Vec<u8>>;
35
36/// The container for the `MemoryStoreMap`s by namespace and then root key
37#[derive(Default)]
38struct MemoryDatabases {
39    databases: BTreeMap<String, BTreeMap<Vec<u8>, Arc<RwLock<MemoryStoreMap>>>>,
40}
41
42/// A connection to a namespace of key-values in memory.
43#[derive(Clone, Debug)]
44pub struct MemoryDatabase {
45    /// The current namespace.
46    namespace: String,
47    /// Whether to remove the namespace on drop.
48    kill_on_drop: bool,
49}
50
51impl MemoryDatabases {
52    fn sync_open(
53        &mut self,
54        namespace: &str,
55        root_key: &[u8],
56    ) -> Result<MemoryStore, MemoryStoreError> {
57        let Some(stores) = self.databases.get_mut(namespace) else {
58            return Err(MemoryStoreError::NamespaceNotFound);
59        };
60        let store = stores.entry(root_key.to_vec()).or_insert_with(|| {
61            let map = MemoryStoreMap::new();
62            Arc::new(RwLock::new(map))
63        });
64        let map = store.clone();
65        Ok(MemoryStore {
66            map,
67            root_key: root_key.to_vec(),
68        })
69    }
70
71    fn sync_list_all(&self) -> Vec<String> {
72        self.databases.keys().cloned().collect::<Vec<_>>()
73    }
74
75    fn sync_list_root_keys(&self, namespace: &str) -> Vec<Vec<u8>> {
76        match self.databases.get(namespace) {
77            None => Vec::new(),
78            Some(map) => map.keys().cloned().collect::<Vec<_>>(),
79        }
80    }
81
82    fn sync_exists(&self, namespace: &str) -> bool {
83        self.databases.contains_key(namespace)
84    }
85
86    fn sync_create(&mut self, namespace: &str) {
87        self.databases
88            .insert(namespace.to_string(), BTreeMap::new());
89    }
90
91    fn sync_delete(&mut self, namespace: &str) {
92        self.databases.remove(namespace);
93    }
94}
95
96/// The global table of namespaces.
97static MEMORY_DATABASES: LazyLock<Mutex<MemoryDatabases>> =
98    LazyLock::new(|| Mutex::new(MemoryDatabases::default()));
99
100/// A virtual DB client where data are persisted in memory.
101#[derive(Clone)]
102pub struct MemoryStore {
103    /// The map used for storing the data.
104    map: Arc<RwLock<MemoryStoreMap>>,
105    /// The root key.
106    root_key: Vec<u8>,
107}
108
109impl WithError for MemoryDatabase {
110    type Error = MemoryStoreError;
111}
112
113impl WithError for MemoryStore {
114    type Error = MemoryStoreError;
115}
116
117impl ReadableKeyValueStore for MemoryStore {
118    const MAX_KEY_SIZE: usize = usize::MAX;
119
120    fn root_key(&self) -> Result<Vec<u8>, MemoryStoreError> {
121        Ok(self.root_key.clone())
122    }
123
124    async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, MemoryStoreError> {
125        let map = self
126            .map
127            .read()
128            .expect("MemoryStore lock should not be poisoned");
129        Ok(map.get(key).cloned())
130    }
131
132    async fn contains_key(&self, key: &[u8]) -> Result<bool, MemoryStoreError> {
133        let map = self
134            .map
135            .read()
136            .expect("MemoryStore lock should not be poisoned");
137        Ok(map.contains_key(key))
138    }
139
140    async fn contains_keys(&self, keys: &[Vec<u8>]) -> Result<Vec<bool>, MemoryStoreError> {
141        let map = self
142            .map
143            .read()
144            .expect("MemoryStore lock should not be poisoned");
145        Ok(keys
146            .iter()
147            .map(|key| map.contains_key(key))
148            .collect::<Vec<_>>())
149    }
150
151    async fn read_multi_values_bytes(
152        &self,
153        keys: &[Vec<u8>],
154    ) -> Result<Vec<Option<Vec<u8>>>, MemoryStoreError> {
155        let map = self
156            .map
157            .read()
158            .expect("MemoryStore lock should not be poisoned");
159        let mut result = Vec::new();
160        for key in keys {
161            result.push(map.get(key).cloned());
162        }
163        Ok(result)
164    }
165
166    async fn find_keys_by_prefix(
167        &self,
168        key_prefix: &[u8],
169    ) -> Result<Vec<Vec<u8>>, MemoryStoreError> {
170        let map = self
171            .map
172            .read()
173            .expect("MemoryStore lock should not be poisoned");
174        let mut values = Vec::new();
175        let len = key_prefix.len();
176        for (key, _value) in map.range(get_key_range_for_prefix(key_prefix.to_vec())) {
177            values.push(key[len..].to_vec())
178        }
179        Ok(values)
180    }
181
182    async fn find_key_values_by_prefix(
183        &self,
184        key_prefix: &[u8],
185    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, MemoryStoreError> {
186        let map = self
187            .map
188            .read()
189            .expect("MemoryStore lock should not be poisoned");
190        let mut key_values = Vec::new();
191        let len = key_prefix.len();
192        for (key, value) in map.range(get_key_range_for_prefix(key_prefix.to_vec())) {
193            let key_value = (key[len..].to_vec(), value.to_vec());
194            key_values.push(key_value);
195        }
196        Ok(key_values)
197    }
198}
199
200impl WritableKeyValueStore for MemoryStore {
201    const MAX_VALUE_SIZE: usize = usize::MAX;
202
203    async fn write_batch(&self, batch: Batch) -> Result<(), MemoryStoreError> {
204        let mut map = self
205            .map
206            .write()
207            .expect("MemoryStore lock should not be poisoned");
208        for ent in batch.operations {
209            match ent {
210                WriteOperation::Put { key, value } => {
211                    map.insert(key, value);
212                }
213                WriteOperation::Delete { key } => {
214                    map.remove(&key);
215                }
216                WriteOperation::DeletePrefix { key_prefix } => {
217                    let key_list = map
218                        .range(get_key_range_for_prefix(key_prefix))
219                        .map(|x| x.0.to_vec())
220                        .collect::<Vec<_>>();
221                    for key in key_list {
222                        map.remove(&key);
223                    }
224                }
225            }
226        }
227        Ok(())
228    }
229
230    async fn clear_journal(&self) -> Result<(), MemoryStoreError> {
231        Ok(())
232    }
233}
234
235impl MemoryStore {
236    /// Creates a `MemoryStore` that doesn't belong to any registered namespace.
237    #[cfg(with_testing)]
238    pub fn new_for_testing() -> Self {
239        Self {
240            map: Arc::default(),
241            root_key: Vec::new(),
242        }
243    }
244}
245
246impl Drop for MemoryDatabase {
247    fn drop(&mut self) {
248        if self.kill_on_drop {
249            let mut databases = MEMORY_DATABASES
250                .lock()
251                .expect("MEMORY_DATABASES lock should not be poisoned");
252            databases.databases.remove(&self.namespace);
253        }
254    }
255}
256
257impl KeyValueDatabase for MemoryDatabase {
258    type Config = MemoryStoreConfig;
259
260    type Store = MemoryStore;
261
262    fn get_name() -> String {
263        "memory".to_string()
264    }
265
266    async fn connect(config: &Self::Config, namespace: &str) -> Result<Self, MemoryStoreError> {
267        let databases = MEMORY_DATABASES
268            .lock()
269            .expect("MEMORY_DATABASES lock should not be poisoned");
270        if !databases.sync_exists(namespace) {
271            return Err(MemoryStoreError::NamespaceNotFound);
272        };
273        Ok(MemoryDatabase {
274            namespace: namespace.to_string(),
275            kill_on_drop: config.kill_on_drop,
276        })
277    }
278
279    fn open_shared(&self, root_key: &[u8]) -> Result<Self::Store, MemoryStoreError> {
280        let mut databases = MEMORY_DATABASES
281            .lock()
282            .expect("MEMORY_DATABASES lock should not be poisoned");
283        databases.sync_open(&self.namespace, root_key)
284    }
285
286    fn open_exclusive(&self, root_key: &[u8]) -> Result<Self::Store, MemoryStoreError> {
287        self.open_shared(root_key)
288    }
289
290    async fn list_all(_config: &Self::Config) -> Result<Vec<String>, MemoryStoreError> {
291        let databases = MEMORY_DATABASES
292            .lock()
293            .expect("MEMORY_DATABASES lock should not be poisoned");
294        Ok(databases.sync_list_all())
295    }
296
297    async fn list_root_keys(&self) -> Result<Vec<Vec<u8>>, MemoryStoreError> {
298        let databases = MEMORY_DATABASES
299            .lock()
300            .expect("MEMORY_DATABASES lock should not be poisoned");
301        Ok(databases.sync_list_root_keys(&self.namespace))
302    }
303
304    async fn exists(_config: &Self::Config, namespace: &str) -> Result<bool, MemoryStoreError> {
305        let databases = MEMORY_DATABASES
306            .lock()
307            .expect("MEMORY_DATABASES lock should not be poisoned");
308        Ok(databases.sync_exists(namespace))
309    }
310
311    async fn create(_config: &Self::Config, namespace: &str) -> Result<(), MemoryStoreError> {
312        let mut databases = MEMORY_DATABASES
313            .lock()
314            .expect("MEMORY_DATABASES lock should not be poisoned");
315        if databases.sync_exists(namespace) {
316            return Err(MemoryStoreError::StoreAlreadyExists);
317        }
318        databases.sync_create(namespace);
319        Ok(())
320    }
321
322    async fn delete(_config: &Self::Config, namespace: &str) -> Result<(), MemoryStoreError> {
323        let mut databases = MEMORY_DATABASES
324            .lock()
325            .expect("MEMORY_DATABASES lock should not be poisoned");
326        databases.sync_delete(namespace);
327        Ok(())
328    }
329}
330
331#[cfg(with_testing)]
332impl TestKeyValueDatabase for MemoryDatabase {
333    async fn new_test_config() -> Result<MemoryStoreConfig, MemoryStoreError> {
334        Ok(MemoryStoreConfig {
335            kill_on_drop: false,
336        })
337    }
338}
339
340/// The error type for [`MemoryStore`].
341#[derive(Error, Debug)]
342pub enum MemoryStoreError {
343    /// Store already exists during a create operation
344    #[error("Store already exists during a create operation")]
345    StoreAlreadyExists,
346
347    /// Serialization error with BCS.
348    #[error(transparent)]
349    BcsError(#[from] bcs::Error),
350
351    /// The namespace does not exist
352    #[error("The namespace does not exist")]
353    NamespaceNotFound,
354}
355
356impl KeyValueStoreError for MemoryStoreError {
357    const BACKEND: &'static str = "memory";
358}
linera_views/backends/memory.rs

linera_views/backends/
memory.rs
