uqa_storage/mvcc/
memory.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! In-memory reference owner for conditional commits and pinned record reads.
8
9use std::borrow::Borrow;
10use std::collections::BTreeMap;
11use std::sync::Arc;
12
13use parking_lot::Mutex;
14use uqa_core::memory::{BudgetedVec, MemoryBudget, MemoryReservation};
15
16use crate::read_control::StorageReadControl;
17
18use super::key::RecordKey;
19use super::{
20    CommitSequence, CommittedRecordSnapshot, PreparedRecordCommit, RecordHistory, RecordVersion,
21    RecordWrite, ScannedRecord, SharedRecordValue, VersionResult,
22};
23
24mod reclamation;
25
26type History = RecordHistory<SharedRecordValue>;
27
28struct RecordEntry {
29    history: History,
30    _memory: MemoryReservation,
31}
32
33#[derive(Default)]
34struct State {
35    sequence: CommitSequence,
36    records: BTreeMap<RecordKey, RecordEntry>,
37    snapshots: BTreeMap<CommitSequence, usize>,
38    reclamation_epoch: u64,
39    reclamation_domains: BTreeMap<RecordKey, (u64, MemoryReservation)>,
40}
41
42struct Database {
43    state: Mutex<State>,
44    memory: MemoryBudget,
45}
46
47/// Reference storage for the record contract, without physical durability.
48///
49/// Snapshot admission, conditional validation and publication share one mutex; no lock is retained between calls. SQL execution and transaction-local overlays are deliberately outside this committed-record owner.
50#[derive(Clone)]
51pub struct MemoryVersionStore {
52    database: Arc<Database>,
53}
54
55impl MemoryVersionStore {
56    pub fn new(memory: &MemoryBudget) -> Self {
57        Self {
58            database: Arc::new(Database {
59                state: Mutex::new(State::default()),
60                memory: memory.clone(),
61            }),
62        }
63    }
64
65    pub fn snapshot(&self) -> VersionResult<MemoryRecordSnapshot> {
66        let memory = self
67            .database
68            .memory
69            .reserve(std::mem::size_of::<MemoryRecordSnapshot>())?;
70        let mut state = self.database.state.lock();
71        let sequence = state.sequence;
72        let count = state.snapshots.entry(sequence).or_default();
73        *count = count
74            .checked_add(1)
75            .ok_or(uqa_core::memory::MemoryError::SizeOverflow)?;
76        Ok(MemoryRecordSnapshot {
77            database: Arc::clone(&self.database),
78            sequence,
79            reclamation_epoch: state.reclamation_epoch,
80            _memory: memory,
81        })
82    }
83
84    /// Validate every original revision before publishing any record.
85    ///
86    /// Values and replacement histories are prepared privately. Failed validation, allocation or cancellation preserves all published records.
87    pub fn commit(
88        &self,
89        writes: &[RecordWrite<'_>],
90        control: &StorageReadControl,
91    ) -> VersionResult<CommitSequence> {
92        let preparation = StorageReadControl::new(&self.database.memory, control.cancellation());
93        let prepared = PreparedRecordCommit::new(writes, &preparation)?;
94        self.commit_prepared(&prepared, control)
95    }
96
97    pub fn commit_prepared(
98        &self,
99        commit: &PreparedRecordCommit,
100        control: &StorageReadControl,
101    ) -> VersionResult<CommitSequence> {
102        control.cancellation().check()?;
103        let values = self.retain_values(commit, control)?;
104        let mut state = self.database.state.lock();
105        commit.validate_snapshot(state.sequence)?;
106        for (prefix, (minimum, _)) in &state.reclamation_domains {
107            commit.validate_reclamation_epoch(
108                prefix.bytes(),
109                *minimum,
110                state.reclamation_epoch,
111                control,
112            )?;
113        }
114        let validate = || {
115            commit.validate(control, |key| {
116                Ok(state
117                    .records
118                    .get(key)
119                    .and_then(|entry| entry.history.head())
120                    .map(RecordVersion::sequence))
121            })
122        };
123        if values.is_empty() {
124            validate()?;
125            return Ok(state.sequence);
126        }
127        let sequence = state.sequence.successor()?;
128        validate()?;
129        let mut prepared = BudgetedVec::new(&self.database.memory);
130        prepared.reserve(values.len())?;
131        for (key, value) in values.iter() {
132            control.cancellation().check()?;
133            let value = value.clone();
134            let history = if let Some(entry) = state.records.get(key.bytes()) {
135                entry.history.fork_appending(sequence, value)?
136            } else {
137                let mut history = History::new(&self.database.memory);
138                history.append(sequence, value)?;
139                history
140            };
141            let memory = self
142                .database
143                .memory
144                .reserve(std::mem::size_of::<(RecordKey, RecordEntry)>())?;
145            prepared.push((
146                key.clone(),
147                RecordEntry {
148                    history,
149                    _memory: memory,
150                },
151            ))?;
152        }
153        control.cancellation().check()?;
154        // Publication is infallible after every payload and history reservation
155        // succeeds. Tree-node bookkeeping follows the collection allocator.
156        while let Some((key, entry)) = prepared.pop() {
157            state.records.insert(key, entry);
158        }
159        state.sequence = sequence;
160        Ok(sequence)
161    }
162
163    fn retain_values(
164        &self,
165        commit: &PreparedRecordCommit,
166        control: &StorageReadControl,
167    ) -> VersionResult<BudgetedVec<(RecordKey, Option<SharedRecordValue>)>> {
168        let mut values = BudgetedVec::new(&self.database.memory);
169        values.reserve(commit.len())?;
170        let mut writes = commit.writes();
171        while let Some(write) = writes.next(control)? {
172            let key = RecordKey::new(write.key(), &self.database.memory)?;
173            let mut value = write.shared_value();
174            if let Some(source) = value
175                .as_ref()
176                .filter(|value| !value.budget().shares_allowance(&self.database.memory))
177            {
178                let mut owned = BudgetedVec::new(&self.database.memory);
179                owned.reserve(source.len())?;
180                for chunk in source.chunks(65536) {
181                    control.cancellation().check()?;
182                    owned.extend_from_slice(chunk)?;
183                }
184                value = Some(Arc::new(owned));
185            }
186            values.push((key, value))?;
187        }
188        Ok(values)
189    }
190
191    /// Reclaim histories under the same gate that admits new snapshots.
192    pub fn reclaim(&self) -> VersionResult<usize> {
193        let mut state = self.database.state.lock();
194        let horizon = state
195            .snapshots
196            .first_key_value()
197            .map_or(state.sequence, |(sequence, _)| *sequence);
198        state.records.values_mut().try_fold(0, |removed, entry| {
199            Ok(removed + entry.history.reclaim_before(horizon)?)
200        })
201    }
202}
203
204/// One pinned committed sequence. Dropping the last owner releases its lease.
205pub struct MemoryRecordSnapshot {
206    database: Arc<Database>,
207    sequence: CommitSequence,
208    reclamation_epoch: u64,
209    _memory: MemoryReservation,
210}
211
212impl MemoryRecordSnapshot {
213    pub fn sequence(&self) -> CommitSequence {
214        self.sequence
215    }
216
217    /// Return a revision even if its payload is a tombstone.
218    pub fn get(&self, key: &[u8]) -> Option<RecordVersion<SharedRecordValue>> {
219        self.database
220            .state
221            .lock()
222            .records
223            .get(key)
224            .and_then(|entry| entry.history.visible_at(self.sequence))
225            .cloned()
226    }
227
228    /// Read a bounded ordered page, including tombstones for revision checks.
229    pub fn scan(
230        &self,
231        prefix: &[u8],
232        after: Option<&[u8]>,
233        limit: usize,
234        control: &StorageReadControl,
235    ) -> VersionResult<BudgetedVec<ScannedRecord>> {
236        control.cancellation().check()?;
237        let mut result = BudgetedVec::new(control.memory());
238        if limit == 0 {
239            return Ok(result);
240        }
241        let state = self.database.state.lock();
242        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
243        for (key, entry) in state
244            .records
245            .range::<[u8], _>((std::ops::Bound::Included(start), std::ops::Bound::Unbounded))
246        {
247            control.cancellation().check()?;
248            let bytes: &[u8] = key.borrow();
249            if !bytes.starts_with(prefix) {
250                break;
251            }
252            if after.is_some_and(|after| bytes <= after) {
253                continue;
254            }
255            let Some(version) = entry.history.visible_at(self.sequence) else {
256                continue;
257            };
258            result.push(ScannedRecord {
259                key: copy_bytes(bytes, control.memory())?,
260                version: version.clone(),
261            })?;
262            if result.len() == limit {
263                break;
264            }
265        }
266        Ok(result)
267    }
268}
269
270impl CommittedRecordSnapshot for MemoryRecordSnapshot {
271    fn reclamation_epoch(&self) -> Option<u64> {
272        Some(self.reclamation_epoch)
273    }
274    fn sequence(&self) -> CommitSequence {
275        Self::sequence(self)
276    }
277
278    fn get(
279        &self,
280        key: &[u8],
281        control: &StorageReadControl,
282    ) -> VersionResult<Option<RecordVersion<SharedRecordValue>>> {
283        control.cancellation().check()?;
284        Ok(Self::get(self, key))
285    }
286
287    fn visit_value_bounded(
288        &self,
289        key: &[u8],
290        max_bytes: usize,
291        control: &StorageReadControl,
292        visit: &mut super::RecordValueVisitor<'_>,
293    ) -> VersionResult<()> {
294        self.visit_value(key, control, &mut |record| {
295            if let Some(value) = record.and_then(|record| record.value) {
296                control.check_value_size(value.len(), max_bytes)?;
297            }
298            visit(record)
299        })
300    }
301
302    fn scan(
303        &self,
304        prefix: &[u8],
305        after: Option<&[u8]>,
306        limit: usize,
307        control: &StorageReadControl,
308    ) -> VersionResult<BudgetedVec<ScannedRecord>> {
309        Self::scan(self, prefix, after, limit, control)
310    }
311
312    fn visit_prefix(
313        &self,
314        prefix: &[u8],
315        after: Option<&[u8]>,
316        limit: usize,
317        control: &StorageReadControl,
318        visit: &mut super::RecordScanVisitor<'_>,
319    ) -> VersionResult<()> {
320        control.cancellation().check()?;
321        if limit == 0 {
322            return Ok(());
323        }
324        let state = self.database.state.lock();
325        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
326        let mut count = 0;
327        for (key, history) in state
328            .records
329            .range::<[u8], _>((std::ops::Bound::Included(start), std::ops::Bound::Unbounded))
330        {
331            control.cancellation().check()?;
332            let key = key.bytes();
333            if !key.starts_with(prefix) {
334                break;
335            }
336            if after.is_some_and(|after| key <= after) {
337                continue;
338            }
339            if let Some(record) = history.history.visible_at(self.sequence) {
340                let more = visit(
341                    key,
342                    super::BorrowedRecord {
343                        revision: Some(record.sequence()),
344                        value: record.value().map(|value| &***value),
345                    },
346                )?;
347                control.cancellation().check()?;
348                count += 1;
349                if !more || count == limit {
350                    break;
351                }
352            }
353        }
354        Ok(())
355    }
356}
357
358impl Drop for MemoryRecordSnapshot {
359    fn drop(&mut self) {
360        let mut state = self.database.state.lock();
361        let count = state
362            .snapshots
363            .get_mut(&self.sequence)
364            .expect("snapshot owns its lease");
365        *count -= 1;
366        if *count == 0 {
367            state.snapshots.remove(&self.sequence);
368        }
369    }
370}
371
372fn copy_bytes(value: &[u8], memory: &MemoryBudget) -> VersionResult<BudgetedVec<u8>> {
373    let mut owned = BudgetedVec::new(memory);
374    owned.extend_from_slice(value)?;
375    Ok(owned)
376}
uqa_storage/mvcc/memory.rs

uqa_storage/mvcc/
memory.rs
