uqa_storage/key_value/
memory_store.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! In-memory implementation of the backend-neutral key/value traits.
8
9mod view;
10
11use super::{
12    BTreeMap, KeyValueBatch, KeyValueBatchOperation, KeyValueStore, Mutex, StorageBackendError,
13    StorageBackendResult,
14};
15
16/// In-memory Key/Value store used by trait-level tests and future non-SQL
17/// fixtures.
18#[derive(Debug)]
19pub struct MemoryKeyValueStore {
20    inner: Mutex<MemoryKeyValueState>,
21    control: crate::read_control::StorageReadControl,
22}
23
24#[derive(Debug, Default, Clone)]
25struct MemoryKeyValueState {
26    map: BTreeMap<Vec<u8>, Vec<u8>>,
27    transactions: Vec<BTreeMap<Vec<u8>, Vec<u8>>>,
28    savepoints: Vec<MemorySavepoint>,
29    transaction_read_only: bool,
30    transaction_written: bool,
31    change_version: u64,
32    read_revision: std::sync::Arc<()>,
33}
34
35#[derive(Debug, Clone)]
36struct MemorySavepoint {
37    name: String,
38    snapshot: BTreeMap<Vec<u8>, Vec<u8>>,
39}
40
41impl MemoryKeyValueStore {
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    fn evaluate_mutation(
47        &self,
48        mutate: &mut super::KeyValueMutation<'_>,
49        certify: bool,
50    ) -> StorageBackendResult<Option<super::KeyValueReadRevision>> {
51        self.control.check()?;
52        let mut state = self.inner.lock();
53        self.control.check()?;
54        check_write(&state)?;
55        let mut batch = MemoryKeyValueBatch {
56            store: self,
57            operations: Vec::new(),
58        };
59        mutate(
60            &view::MemoryRead {
61                state: &state,
62                control: &self.control,
63            },
64            &mut batch,
65        )?;
66        self.control.check()?;
67        apply_operations(&mut state, batch.operations)?;
68        Ok(certify.then(|| super::KeyValueReadRevision::memory(&state.read_revision)))
69    }
70}
71
72impl Default for MemoryKeyValueStore {
73    fn default() -> Self {
74        Self {
75            inner: Mutex::new(MemoryKeyValueState::default()),
76            control: crate::read_control::StorageReadControl::with_limit(64 << 20),
77        }
78    }
79}
80
81impl KeyValueStore for MemoryKeyValueStore {
82    fn with_read_view(&self, read: &mut super::KeyValueReadScope<'_>) -> StorageBackendResult<()> {
83        self.control.check()?;
84        let state = self.inner.lock();
85        self.control.check()?;
86        read(&view::MemoryRead {
87            state: &state,
88            control: &self.control,
89        })?;
90        self.control.check()
91    }
92
93    fn with_mutation(&self, mutate: &mut super::KeyValueMutation<'_>) -> StorageBackendResult<()> {
94        self.evaluate_mutation(mutate, false).map(|_| ())
95    }
96
97    fn with_mutation_revision(
98        &self,
99        _prefixes: &[&[u8]],
100        mutate: &mut super::KeyValueMutation<'_>,
101    ) -> StorageBackendResult<Option<super::KeyValueReadRevision>> {
102        self.evaluate_mutation(mutate, true)
103    }
104
105    fn visit_value(
106        &self,
107        key: &[u8],
108        control: &crate::read_control::StorageReadControl,
109        visit: &mut crate::read_control::ValueReadVisitor<'_>,
110    ) -> StorageBackendResult<()> {
111        control.check()?;
112        let state = self.inner.lock();
113        control.check()?;
114        visit(state.map.get(key).map(Vec::as_slice))?;
115        control.check()
116    }
117
118    fn visit_prefix_after(
119        &self,
120        prefix: &[u8],
121        after: Option<&[u8]>,
122        limit: usize,
123        control: &crate::read_control::StorageReadControl,
124        visit: &mut crate::read_control::KeyValueReadVisitor<'_>,
125    ) -> StorageBackendResult<()> {
126        use std::ops::Bound::{Excluded, Included, Unbounded};
127        control.check()?;
128        if limit == 0 {
129            return Ok(());
130        }
131        let state = self.inner.lock();
132        let lower = match after {
133            Some(after) if after >= prefix => Excluded(after),
134            _ => Included(prefix),
135        };
136        for (key, value) in state.map.range::<[u8], _>((lower, Unbounded)).take(limit) {
137            control.check()?;
138            if !key.starts_with(prefix) {
139                break;
140            }
141            visit(key, value)?;
142        }
143        control.check()
144    }
145
146    fn get(&self, key: &[u8]) -> StorageBackendResult<Option<Vec<u8>>> {
147        Ok(self.inner.lock().map.get(key).cloned())
148    }
149
150    fn contains_key(&self, key: &[u8]) -> StorageBackendResult<bool> {
151        Ok(self.inner.lock().map.contains_key(key))
152    }
153
154    fn put(&self, key: &[u8], value: &[u8]) -> StorageBackendResult<()> {
155        let mut inner = self.inner.lock();
156        prepare_write(&mut inner)?;
157        inner.map.insert(key.to_vec(), value.to_vec());
158        finish_autocommit_write(&mut inner);
159        Ok(())
160    }
161
162    fn delete(&self, key: &[u8]) -> StorageBackendResult<()> {
163        let mut inner = self.inner.lock();
164        prepare_write(&mut inner)?;
165        inner.map.remove(key);
166        finish_autocommit_write(&mut inner);
167        Ok(())
168    }
169
170    fn scan_prefix(&self, prefix: &[u8]) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>> {
171        Ok(self
172            .inner
173            .lock()
174            .map
175            .range(prefix.to_vec()..)
176            .take_while(|(key, _)| key.starts_with(prefix))
177            .map(|(key, value)| (key.clone(), value.clone()))
178            .collect())
179    }
180
181    fn scan_prefix_after(
182        &self,
183        prefix: &[u8],
184        after: Option<&[u8]>,
185        limit: usize,
186    ) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>> {
187        use std::ops::Bound::{Excluded, Included, Unbounded};
188
189        if limit == 0 {
190            return Ok(Vec::new());
191        }
192        let inner = self.inner.lock();
193        let lower = match after {
194            Some(after) if after >= prefix => Excluded(after.to_vec()),
195            Some(_) | None => Included(prefix.to_vec()),
196        };
197        Ok(inner
198            .map
199            .range((lower, Unbounded))
200            .take_while(|(key, _)| key.starts_with(prefix))
201            .take(limit)
202            .map(|(key, value)| (key.clone(), value.clone()))
203            .collect())
204    }
205
206    fn scan_prefix_keys_after(
207        &self,
208        prefix: &[u8],
209        after: Option<&[u8]>,
210        limit: usize,
211    ) -> StorageBackendResult<Vec<Vec<u8>>> {
212        use std::ops::Bound::{Excluded, Included, Unbounded};
213
214        if limit == 0 {
215            return Ok(Vec::new());
216        }
217        let inner = self.inner.lock();
218        let lower = match after {
219            Some(after) if after >= prefix => Excluded(after.to_vec()),
220            Some(_) | None => Included(prefix.to_vec()),
221        };
222        Ok(inner
223            .map
224            .range((lower, Unbounded))
225            .take_while(|(key, _)| key.starts_with(prefix))
226            .take(limit)
227            .map(|(key, _)| key.clone())
228            .collect())
229    }
230
231    fn first_prefix_after(
232        &self,
233        prefix: &[u8],
234        after: Option<&[u8]>,
235    ) -> StorageBackendResult<Option<(Vec<u8>, Vec<u8>)>> {
236        use std::ops::Bound::{Excluded, Included, Unbounded};
237
238        let inner = self.inner.lock();
239        let lower = match after {
240            Some(after) if after >= prefix => Excluded(after.to_vec()),
241            Some(_) | None => Included(prefix.to_vec()),
242        };
243        Ok(inner
244            .map
245            .range((lower, Unbounded))
246            .next()
247            .filter(|(key, _)| key.starts_with(prefix))
248            .map(|(key, value)| (key.clone(), value.clone())))
249    }
250
251    fn delete_prefix(&self, prefix: &[u8]) -> StorageBackendResult<usize> {
252        let mut inner = self.inner.lock();
253        prepare_write(&mut inner)?;
254        let keys = inner
255            .map
256            .range(prefix.to_vec()..)
257            .take_while(|(key, _)| key.starts_with(prefix))
258            .map(|(key, _)| key.clone())
259            .collect::<Vec<_>>();
260        for key in &keys {
261            inner.map.remove(key);
262        }
263        finish_autocommit_write(&mut inner);
264        Ok(keys.len())
265    }
266
267    fn batch(&self) -> Box<dyn KeyValueBatch + '_> {
268        Box::new(MemoryKeyValueBatch {
269            store: self,
270            operations: Vec::new(),
271        })
272    }
273
274    fn begin_transaction(&self) -> StorageBackendResult<()> {
275        let mut inner = self.inner.lock();
276        if !inner.transactions.is_empty() {
277            return Err(StorageBackendError::Other(
278                "a KeyValue transaction is already open".into(),
279            ));
280        }
281        let snapshot = inner.map.clone();
282        inner.transactions.push(snapshot);
283        inner.transaction_read_only = false;
284        inner.transaction_written = false;
285        Ok(())
286    }
287
288    fn begin_read_transaction(&self) -> StorageBackendResult<()> {
289        let mut inner = self.inner.lock();
290        if !inner.transactions.is_empty() {
291            return Err(StorageBackendError::Other(
292                "a KeyValue transaction is already open".into(),
293            ));
294        }
295        let snapshot = inner.map.clone();
296        inner.transactions.push(snapshot);
297        inner.transaction_read_only = true;
298        inner.transaction_written = false;
299        Ok(())
300    }
301
302    fn in_transaction(&self) -> bool {
303        !self.inner.lock().transactions.is_empty()
304    }
305
306    fn transaction_has_written(&self) -> StorageBackendResult<bool> {
307        Ok(self.inner.lock().transaction_written)
308    }
309
310    fn change_version(&self) -> StorageBackendResult<Option<u64>> {
311        Ok(Some(self.inner.lock().change_version))
312    }
313
314    fn commit_transaction(&self) -> StorageBackendResult<()> {
315        let mut inner = self.inner.lock();
316        inner.transactions.pop().ok_or_else(|| {
317            StorageBackendError::Other("no open KeyValue transaction to commit".into())
318        })?;
319        if inner.transaction_written {
320            inner.change_version = inner.change_version.wrapping_add(1);
321        }
322        inner.transaction_read_only = false;
323        inner.transaction_written = false;
324        inner.savepoints.clear();
325        Ok(())
326    }
327
328    fn rollback_transaction(&self) -> StorageBackendResult<()> {
329        let mut inner = self.inner.lock();
330        let snapshot = inner.transactions.pop().ok_or_else(|| {
331            StorageBackendError::Other("no open KeyValue transaction to roll back".into())
332        })?;
333        inner.map = snapshot;
334        inner.read_revision = std::sync::Arc::new(());
335        inner.transaction_read_only = false;
336        inner.transaction_written = false;
337        inner.savepoints.clear();
338        Ok(())
339    }
340
341    fn savepoint(&self, name: &str) -> StorageBackendResult<()> {
342        let mut inner = self.inner.lock();
343        if inner.transactions.is_empty() {
344            return Err(StorageBackendError::Other(
345                "cannot create a savepoint outside a KeyValue transaction".into(),
346            ));
347        }
348        let snapshot = inner.map.clone();
349        inner.savepoints.push(MemorySavepoint {
350            name: name.to_string(),
351            snapshot,
352        });
353        Ok(())
354    }
355
356    fn release_savepoint(&self, name: &str) -> StorageBackendResult<()> {
357        let mut inner = self.inner.lock();
358        let position = inner
359            .savepoints
360            .iter()
361            .rposition(|savepoint| savepoint.name == name)
362            .ok_or_else(|| StorageBackendError::Other(format!("unknown savepoint `{name}`")))?;
363        inner.savepoints.truncate(position);
364        Ok(())
365    }
366
367    fn rollback_to_savepoint(&self, name: &str) -> StorageBackendResult<()> {
368        let mut inner = self.inner.lock();
369        let position = inner
370            .savepoints
371            .iter()
372            .rposition(|savepoint| savepoint.name == name)
373            .ok_or_else(|| StorageBackendError::Other(format!("unknown savepoint `{name}`")))?;
374        inner.map = inner.savepoints[position].snapshot.clone();
375        inner.read_revision = std::sync::Arc::new(());
376        inner.savepoints.truncate(position + 1);
377        Ok(())
378    }
379}
380
381struct MemoryKeyValueBatch<'a> {
382    store: &'a MemoryKeyValueStore,
383    operations: Vec<KeyValueBatchOperation>,
384}
385
386impl KeyValueBatch for MemoryKeyValueBatch<'_> {
387    fn put(&mut self, key: &[u8], value: &[u8]) -> StorageBackendResult<()> {
388        self.operations
389            .push(KeyValueBatchOperation::Put(key.to_vec(), value.to_vec()));
390        Ok(())
391    }
392
393    fn delete(&mut self, key: &[u8]) -> StorageBackendResult<()> {
394        self.operations
395            .push(KeyValueBatchOperation::Delete(key.to_vec()));
396        Ok(())
397    }
398
399    fn delete_prefix(&mut self, prefix: &[u8]) -> StorageBackendResult<()> {
400        self.operations
401            .push(KeyValueBatchOperation::DeletePrefix(prefix.to_vec()));
402        Ok(())
403    }
404
405    fn commit(self: Box<Self>) -> StorageBackendResult<()> {
406        let mut inner = self.store.inner.lock();
407        apply_operations(&mut inner, self.operations)
408    }
409}
410
411fn apply_operations(
412    inner: &mut MemoryKeyValueState,
413    operations: Vec<KeyValueBatchOperation>,
414) -> StorageBackendResult<()> {
415    prepare_write(inner)?;
416    for operation in operations {
417        match operation {
418            KeyValueBatchOperation::Put(key, value) => {
419                inner.map.insert(key, value);
420            }
421            KeyValueBatchOperation::Delete(key) => {
422                inner.map.remove(&key);
423            }
424            KeyValueBatchOperation::DeletePrefix(prefix) => {
425                let keys = inner
426                    .map
427                    .range(prefix.clone()..)
428                    .take_while(|(key, _)| key.starts_with(&prefix))
429                    .map(|(key, _)| key.clone())
430                    .collect::<Vec<_>>();
431                for key in keys {
432                    inner.map.remove(&key);
433                }
434            }
435        }
436    }
437    finish_autocommit_write(inner);
438    Ok(())
439}
440
441fn prepare_write(inner: &mut MemoryKeyValueState) -> StorageBackendResult<()> {
442    check_write(inner)?;
443    if !inner.transactions.is_empty() {
444        inner.transaction_written = true;
445    }
446    Ok(())
447}
448
449fn check_write(inner: &MemoryKeyValueState) -> StorageBackendResult<()> {
450    if !inner.transactions.is_empty() && inner.transaction_read_only {
451        return Err(StorageBackendError::Other(
452            "cannot write in a read-only KeyValue transaction".into(),
453        ));
454    }
455    Ok(())
456}
457
458fn finish_autocommit_write(inner: &mut MemoryKeyValueState) {
459    inner.read_revision = std::sync::Arc::new(());
460    if inner.transactions.is_empty() {
461        inner.change_version = inner.change_version.wrapping_add(1);
462    }
463}
uqa_storage/key_value/memory_store.rs

uqa_storage/key_value/
memory_store.rs
