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.cancellation(),
112            )?;
113        }
114        let writes = commit.records();
115        let validate = || {
116            commit.validate(control.cancellation(), |key| {
117                Ok(state
118                    .records
119                    .get(key)
120                    .and_then(|entry| entry.history.head())
121                    .map(RecordVersion::sequence))
122            })
123        };
124        if writes.is_empty() {
125            validate()?;
126            return Ok(state.sequence);
127        }
128        let sequence = state.sequence.successor()?;
129        validate()?;
130        let mut prepared = BudgetedVec::new(&self.database.memory);
131        prepared.reserve(writes.len())?;
132        for (write, value) in writes.iter().zip(values.iter()) {
133            control.cancellation().check()?;
134            let key = RecordKey::new(write.key(), &self.database.memory)?;
135            let value = value.clone();
136            let history = if let Some(entry) = state.records.get(write.key()) {
137                entry.history.fork_appending(sequence, value)?
138            } else {
139                let mut history = History::new(&self.database.memory);
140                history.append(sequence, value)?;
141                history
142            };
143            let memory = self
144                .database
145                .memory
146                .reserve(std::mem::size_of::<(RecordKey, RecordEntry)>())?;
147            prepared.push((
148                key,
149                RecordEntry {
150                    history,
151                    _memory: memory,
152                },
153            ))?;
154        }
155        control.cancellation().check()?;
156        // Publication is infallible after every payload and history reservation
157        // succeeds. Tree-node bookkeeping follows the collection allocator.
158        while let Some((key, entry)) = prepared.pop() {
159            state.records.insert(key, entry);
160        }
161        state.sequence = sequence;
162        Ok(sequence)
163    }
164
165    fn retain_values(
166        &self,
167        commit: &PreparedRecordCommit,
168        control: &StorageReadControl,
169    ) -> VersionResult<BudgetedVec<Option<SharedRecordValue>>> {
170        let mut values = BudgetedVec::new(&self.database.memory);
171        values.reserve(commit.records().len())?;
172        for write in commit.records() {
173            control.cancellation().check()?;
174            let mut value = write.shared_value();
175            if let Some(source) = value
176                .as_ref()
177                .filter(|value| !value.budget().shares_allowance(&self.database.memory))
178            {
179                let mut owned = BudgetedVec::new(&self.database.memory);
180                owned.reserve(source.len())?;
181                for chunk in source.chunks(65536) {
182                    control.cancellation().check()?;
183                    owned.extend_from_slice(chunk)?;
184                }
185                value = Some(Arc::new(owned));
186            }
187            values.push(value)?;
188        }
189        Ok(values)
190    }
191
192    /// Reclaim histories under the same gate that admits new snapshots.
193    pub fn reclaim(&self) -> VersionResult<usize> {
194        let mut state = self.database.state.lock();
195        let horizon = state
196            .snapshots
197            .first_key_value()
198            .map_or(state.sequence, |(sequence, _)| *sequence);
199        state.records.values_mut().try_fold(0, |removed, entry| {
200            Ok(removed + entry.history.reclaim_before(horizon)?)
201        })
202    }
203}
204
205/// One pinned committed sequence. Dropping the last owner releases its lease.
206pub struct MemoryRecordSnapshot {
207    database: Arc<Database>,
208    sequence: CommitSequence,
209    reclamation_epoch: u64,
210    _memory: MemoryReservation,
211}
212
213impl MemoryRecordSnapshot {
214    pub fn sequence(&self) -> CommitSequence {
215        self.sequence
216    }
217
218    /// Return a revision even if its payload is a tombstone.
219    pub fn get(&self, key: &[u8]) -> Option<RecordVersion<SharedRecordValue>> {
220        self.database
221            .state
222            .lock()
223            .records
224            .get(key)
225            .and_then(|entry| entry.history.visible_at(self.sequence))
226            .cloned()
227    }
228
229    /// Read a bounded ordered page, including tombstones for revision checks.
230    pub fn scan(
231        &self,
232        prefix: &[u8],
233        after: Option<&[u8]>,
234        limit: usize,
235        control: &StorageReadControl,
236    ) -> VersionResult<BudgetedVec<ScannedRecord>> {
237        control.cancellation().check()?;
238        let mut result = BudgetedVec::new(control.memory());
239        if limit == 0 {
240            return Ok(result);
241        }
242        let state = self.database.state.lock();
243        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
244        for (key, entry) in state
245            .records
246            .range::<[u8], _>((std::ops::Bound::Included(start), std::ops::Bound::Unbounded))
247        {
248            control.cancellation().check()?;
249            let bytes: &[u8] = key.borrow();
250            if !bytes.starts_with(prefix) {
251                break;
252            }
253            if after.is_some_and(|after| bytes <= after) {
254                continue;
255            }
256            let Some(version) = entry.history.visible_at(self.sequence) else {
257                continue;
258            };
259            result.push(ScannedRecord {
260                key: copy_bytes(bytes, control.memory())?,
261                version: version.clone(),
262            })?;
263            if result.len() == limit {
264                break;
265            }
266        }
267        Ok(result)
268    }
269}
270
271impl CommittedRecordSnapshot for MemoryRecordSnapshot {
272    fn reclamation_epoch(&self) -> Option<u64> {
273        Some(self.reclamation_epoch)
274    }
275    fn sequence(&self) -> CommitSequence {
276        Self::sequence(self)
277    }
278
279    fn get(
280        &self,
281        key: &[u8],
282        control: &StorageReadControl,
283    ) -> VersionResult<Option<RecordVersion<SharedRecordValue>>> {
284        control.cancellation().check()?;
285        Ok(Self::get(self, key))
286    }
287
288    fn visit_value_bounded(
289        &self,
290        key: &[u8],
291        max_bytes: usize,
292        control: &StorageReadControl,
293        visit: &mut super::RecordValueVisitor<'_>,
294    ) -> VersionResult<()> {
295        self.visit_value(key, control, &mut |record| {
296            if let Some(value) = record.and_then(|record| record.value) {
297                control.check_value_size(value.len(), max_bytes)?;
298            }
299            visit(record)
300        })
301    }
302
303    fn scan(
304        &self,
305        prefix: &[u8],
306        after: Option<&[u8]>,
307        limit: usize,
308        control: &StorageReadControl,
309    ) -> VersionResult<BudgetedVec<ScannedRecord>> {
310        Self::scan(self, prefix, after, limit, control)
311    }
312
313    fn visit_prefix(
314        &self,
315        prefix: &[u8],
316        after: Option<&[u8]>,
317        limit: usize,
318        control: &StorageReadControl,
319        visit: &mut super::RecordScanVisitor<'_>,
320    ) -> VersionResult<()> {
321        control.cancellation().check()?;
322        if limit == 0 {
323            return Ok(());
324        }
325        let state = self.database.state.lock();
326        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
327        let mut count = 0;
328        for (key, history) in state
329            .records
330            .range::<[u8], _>((std::ops::Bound::Included(start), std::ops::Bound::Unbounded))
331        {
332            control.cancellation().check()?;
333            let key = key.bytes();
334            if !key.starts_with(prefix) {
335                break;
336            }
337            if after.is_some_and(|after| key <= after) {
338                continue;
339            }
340            if let Some(record) = history.history.visible_at(self.sequence) {
341                let more = visit(
342                    key,
343                    super::BorrowedRecord {
344                        revision: Some(record.sequence()),
345                        value: record.value().map(|value| &***value),
346                    },
347                )?;
348                control.cancellation().check()?;
349                count += 1;
350                if !more || count == limit {
351                    break;
352                }
353            }
354        }
355        Ok(())
356    }
357}
358
359impl Drop for MemoryRecordSnapshot {
360    fn drop(&mut self) {
361        let mut state = self.database.state.lock();
362        let count = state
363            .snapshots
364            .get_mut(&self.sequence)
365            .expect("snapshot owns its lease");
366        *count -= 1;
367        if *count == 0 {
368            state.snapshots.remove(&self.sequence);
369        }
370    }
371}
372
373fn copy_bytes(value: &[u8], memory: &MemoryBudget) -> VersionResult<BudgetedVec<u8>> {
374    let mut owned = BudgetedVec::new(memory);
375    owned.extend_from_slice(value)?;
376    Ok(owned)
377}
uqa_storage/mvcc/memory.rs

uqa_storage/mvcc/
memory.rs
