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

uqa_storage/key_value/
memory_store.rs
