uqa_storage/mvcc/
overlay.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Transaction-private evaluated replacements, retained command views and undo.
8
9use std::num::NonZeroU64;
10use std::sync::atomic::{AtomicU64, Ordering};
11
12use parking_lot::Mutex;
13use triomphe::Arc as StrongArc;
14use uqa_core::memory::{BudgetedSharedMap, BudgetedVec, MemoryBudget, MemoryReservation};
15
16use crate::read_control::StorageReadControl;
17use crate::StorageSavepointId;
18
19use super::key::RecordKey;
20use super::{PreparedRecordCommit, PreparedRecordWrite, RecordWrite, VersionError, VersionResult};
21
22mod retained;
23use retained::{RetainedSources, Sources};
24
25struct Change {
26    write: PreparedRecordWrite,
27    identity: PrivateRecordRevision,
28}
29
30type Records = BudgetedSharedMap<RecordKey, Change>;
31
32/// Process-local identity of an evaluated private batch. Identities are never reused across transactions or undo branches; they are not durable commit sequences.
33#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
34pub struct PrivateRecordRevision(NonZeroU64);
35
36impl PrivateRecordRevision {
37    fn allocate() -> VersionResult<Self> {
38        static NEXT: AtomicU64 = AtomicU64::new(1);
39        NEXT.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
40            .map(|id| Self(NonZeroU64::new(id).expect("private revisions start at one")))
41            .map_err(|_| VersionError::PrivateRevisionExhausted)
42    }
43
44    #[cfg(test)]
45    pub(crate) fn for_tests() -> Self {
46        Self::allocate().expect("test private revision")
47    }
48
49    pub fn as_u64(self) -> u64 {
50        self.0.get()
51    }
52}
53
54/// A changed key and its private batch identity, without retaining or copying its payload.
55pub struct PrivateRecordKey {
56    key: RecordKey,
57    revision: PrivateRecordRevision,
58}
59
60impl PrivateRecordKey {
61    pub fn key(&self) -> &[u8] {
62        self.key.bytes()
63    }
64
65    pub fn revision(&self) -> PrivateRecordRevision {
66        self.revision
67    }
68}
69
70struct Savepoint {
71    id: StorageSavepointId,
72    records: Records,
73    sources: RetainedSources,
74    revision: Option<PrivateRecordRevision>,
75}
76
77struct State {
78    records: Records,
79    sources: Sources,
80    revision: Option<PrivateRecordRevision>,
81    savepoints: BudgetedVec<Savepoint>,
82}
83
84impl State {
85    fn savepoint_position(&self, id: StorageSavepointId) -> VersionResult<usize> {
86        self.savepoints
87            .iter()
88            .rposition(|savepoint| savepoint.id == id)
89            .ok_or(VersionError::SavepointMissing(id))
90    }
91
92    fn truncate_savepoints(&mut self, len: usize) {
93        if len == 0 {
94            self.savepoints = BudgetedVec::new(self.records.budget());
95        } else {
96            self.savepoints.truncate(len);
97        }
98    }
99}
100
101struct Owner {
102    state: Mutex<State>,
103}
104
105/// One transaction's record changes. No provider lock or native transaction is retained between operations.
106///
107/// Command views and savepoints share immutable ordered roots. Replacements copy only their search paths; obsolete payloads are released as soon as the last referencing root or returned record is dropped. Rollback restores a root without allocating. This primitive does not spill retained values.
108pub struct PrivateRecordChanges {
109    owner: StrongArc<Owner>,
110}
111
112impl PrivateRecordChanges {
113    pub(super) fn share_owner(&self) -> Self {
114        Self {
115            owner: StrongArc::clone(&self.owner),
116        }
117    }
118
119    pub(super) fn write_kind(
120        &self,
121        key: &[u8],
122        control: &StorageReadControl,
123    ) -> VersionResult<Option<super::commit::RecordWriteKind>> {
124        control.check()?;
125        Ok(self
126            .owner
127            .state
128            .lock()
129            .records
130            .get(key)
131            .map(|change| change.write.kind()))
132    }
133
134    pub fn new(memory: &MemoryBudget) -> Self {
135        Self {
136            owner: StrongArc::new(Owner {
137                state: Mutex::new(State {
138                    records: Records::new(memory),
139                    sources: Sources::new(memory),
140                    revision: None,
141                    savepoints: BudgetedVec::new(memory),
142                }),
143            }),
144        }
145    }
146
147    /// Stage one atomic batch of already evaluated replacements. Repeated changes to a private record must preserve its original committed precondition.
148    pub fn apply(
149        &self,
150        writes: &[RecordWrite<'_>],
151        control: &StorageReadControl,
152    ) -> VersionResult<()> {
153        let memory = self.owner.state.lock().records.budget().clone();
154        let owned_control = StorageReadControl::new(&memory, control.cancellation());
155        let prepared = PreparedRecordCommit::new(writes, &owned_control)?;
156        self.apply_owned(prepared.records(), control)
157    }
158
159    pub(crate) fn apply_owned(
160        &self,
161        writes: &[PreparedRecordWrite],
162        control: &StorageReadControl,
163    ) -> VersionResult<()> {
164        if writes.is_empty() {
165            return Ok(());
166        }
167        let mut state = self.owner.state.lock();
168        for (index, write) in writes.iter().enumerate() {
169            control.cancellation().check()?;
170            if let Some(previous) = state.records.get(write.key()) {
171                if previous.write.expected() != write.expected() {
172                    return Err(VersionError::WriteConflict {
173                        mutation: index,
174                        expected: write.expected(),
175                        actual: previous.write.expected(),
176                    });
177                }
178            }
179        }
180        let identity = PrivateRecordRevision::allocate()?;
181        let invalidates_source = writes
182            .iter()
183            .any(|write| state.sources.get(write.key()).is_some_and(Option::is_some));
184        if writes.len() == 1 && !invalidates_source {
185            let write = &writes[0];
186            control.cancellation().check()?;
187            state.records.try_insert(
188                write.shared_key(),
189                Change {
190                    write: write.clone(),
191                    identity,
192                },
193            )?;
194            state.revision = Some(identity);
195            return Ok(());
196        }
197        let mut records = state.records.clone();
198        let mut sources = state.sources.clone();
199        for write in writes {
200            control.cancellation().check()?;
201            records.try_insert(
202                write.shared_key(),
203                Change {
204                    write: write.clone(),
205                    identity,
206                },
207            )?;
208            if sources.get(write.key()).is_some_and(Option::is_some) {
209                sources.try_insert(write.shared_key(), None)?;
210            }
211        }
212        control.cancellation().check()?;
213        // Candidate roots own every reservation before this single atomic publication.
214        state.records = records;
215        state.sources = sources;
216        state.revision = Some(identity);
217        Ok(())
218    }
219
220    pub fn has_written(&self) -> bool {
221        !self.owner.state.lock().records.is_empty()
222    }
223
224    pub fn snapshot(&self) -> VersionResult<PrivateRecordSnapshot> {
225        let state = self.owner.state.lock();
226        let memory = state
227            .records
228            .budget()
229            .reserve(std::mem::size_of::<PrivateRecordSnapshot>())?;
230        Ok(PrivateRecordSnapshot {
231            records: state.records.clone(),
232            sources: state.sources.snapshot(),
233            revision: state.revision,
234            _memory: memory,
235        })
236    }
237
238    pub fn savepoint(&self, id: StorageSavepointId) -> VersionResult<()> {
239        let mut state = self.owner.state.lock();
240        let records = state.records.clone();
241        let sources = state.sources.snapshot();
242        let revision = state.revision;
243        state.savepoints.push(Savepoint {
244            id,
245            records,
246            sources,
247            revision,
248        })?;
249        Ok(())
250    }
251
252    /// Release the nearest matching identity and its nested savepoints, retaining every write.
253    pub fn release_savepoint(&self, id: StorageSavepointId) -> VersionResult<()> {
254        let mut state = self.owner.state.lock();
255        let position = state.savepoint_position(id)?;
256        state.truncate_savepoints(position);
257        Ok(())
258    }
259
260    /// Undo only this owner's changes, retaining the target savepoint for another rollback.
261    pub fn rollback_to_savepoint(&self, id: StorageSavepointId) -> VersionResult<()> {
262        let mut state = self.owner.state.lock();
263        let position = state.savepoint_position(id)?;
264        state.records = state.savepoints[position].records.clone();
265        let sources = state.savepoints[position].sources.clone();
266        state.sources.restore(&sources);
267        state.revision = state.savepoints[position].revision;
268        state.truncate_savepoints(position + 1);
269        Ok(())
270    }
271
272    pub fn rollback(&self) -> VersionResult<()> {
273        let mut state = self.owner.state.lock();
274        let memory = state.records.budget().clone();
275        state.records = Records::new(&memory);
276        state.sources = Sources::new(&memory);
277        state.revision = None;
278        state.truncate_savepoints(0);
279        Ok(())
280    }
281
282    /// Capture immutable final replacements without reevaluating or copying their payloads.
283    pub fn prepare(&self, control: &StorageReadControl) -> VersionResult<PreparedRecordCommit> {
284        control.cancellation().check()?;
285        let state = self.owner.state.lock();
286        let mut writes = BudgetedVec::new(control.memory());
287        for (_, change) in &state.records {
288            control.cancellation().check()?;
289            writes.push(change.write.clone())?;
290        }
291        PreparedRecordCommit::from_unique_owned(writes, control)
292    }
293}
294
295/// Fixed private visibility for a command or retained source cursor; tombstones remain distinguishable from an unchanged key.
296pub struct PrivateRecordSnapshot {
297    records: Records,
298    sources: RetainedSources,
299    revision: Option<PrivateRecordRevision>,
300    _memory: MemoryReservation,
301}
302
303impl PrivateRecordSnapshot {
304    /// Identity of this complete private root. Undo restores the saved identity; successful replacement batches receive identities never used by another root or transaction.
305    pub fn revision(&self) -> Option<PrivateRecordRevision> {
306        self.revision
307    }
308
309    pub(super) fn last_before(
310        &self,
311        prefix: &[u8],
312        before: Option<&[u8]>,
313        control: &StorageReadControl,
314    ) -> VersionResult<Option<PreparedRecordWrite>> {
315        control.check()?;
316        let mut upper = BudgetedVec::new(control.memory());
317        upper.extend_from_slice(prefix)?;
318        let upper = match upper.iter().rposition(|byte| *byte != u8::MAX) {
319            Some(last) => {
320                upper[last] += 1;
321                upper.truncate(last + 1);
322                Some(upper)
323            }
324            None => None,
325        };
326        let end = match (before, upper.as_deref()) {
327            (Some(before), Some(upper)) => Some(before.min(upper)),
328            (before, upper) => before.or(upper),
329        };
330        let bound = end.map_or(std::ops::Bound::Unbounded, std::ops::Bound::Excluded);
331        Ok(self
332            .records
333            .last_before::<[u8]>(bound)
334            .filter(|(key, _)| key.bytes().starts_with(prefix))
335            .map(|(_, change)| change.write.clone()))
336    }
337
338    /// Retain this exact private revision without following later writes or copying its values.
339    pub(super) fn try_clone(&self) -> VersionResult<Self> {
340        let memory = self.records.budget().reserve(std::mem::size_of::<Self>())?;
341        Ok(Self {
342            records: self.records.clone(),
343            sources: self.sources.clone(),
344            revision: self.revision,
345            _memory: memory,
346        })
347    }
348
349    /// Select changed keys on this retained command boundary, including deletions. Undo restores their earlier identities, while later branches and other transactions receive distinct identities.
350    pub fn scan_keys(
351        &self,
352        prefix: &[u8],
353        after: Option<&[u8]>,
354        limit: usize,
355        control: &StorageReadControl,
356    ) -> VersionResult<BudgetedVec<PrivateRecordKey>> {
357        control.check()?;
358        let mut result = BudgetedVec::new(control.memory());
359        if limit == 0 {
360            return Ok(result);
361        }
362        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
363        for (key, change) in self.records.range_from(std::ops::Bound::Included(start)) {
364            control.check()?;
365            let key = key.bytes();
366            if !key.starts_with(prefix) {
367                break;
368            }
369            if after.is_some_and(|after| key <= after) {
370                continue;
371            }
372            result.push(PrivateRecordKey {
373                key: change.write.shared_key(),
374                revision: change.identity,
375            })?;
376            if result.len() == limit {
377                break;
378            }
379        }
380        Ok(result)
381    }
382
383    pub(super) fn visit_merged<P: super::projection::Projection>(
384        &self,
385        committed: &dyn super::CommittedRecordSnapshot,
386        prefix: &[u8],
387        after: Option<&[u8]>,
388        limit: usize,
389        control: &StorageReadControl,
390        visit: &mut super::projection::Visitor<'_, P>,
391    ) -> VersionResult<()> {
392        control.cancellation().check()?;
393        if limit == 0 {
394            return Ok(());
395        }
396        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
397        let mut entries = self.records.range_from(std::ops::Bound::Included(start));
398        let mut next_private = || -> VersionResult<Option<&PreparedRecordWrite>> {
399            for (key, change) in entries.by_ref() {
400                control.cancellation().check()?;
401                if !key.bytes().starts_with(prefix) {
402                    return Ok(None);
403                }
404                if after.is_some_and(|after| key.bytes() <= after) {
405                    continue;
406                }
407                return Ok(Some(&change.write));
408            }
409            Ok(None)
410        };
411        let mut pending = next_private()?;
412        let mut count = 0;
413        let mut running = true;
414        let mut emit = |key: &[u8], record: P::Record<'_>| -> VersionResult<bool> {
415            control.cancellation().check()?;
416            count += 1;
417            let more = visit(key, record)?;
418            control.cancellation().check()?;
419            Ok(more && count < limit)
420        };
421        P::visit(committed, prefix, after, control, &mut |key, record| {
422            while let Some(write) = pending.filter(|write| write.key() < key) {
423                running = emit(write.key(), P::private(write))?;
424                if !running {
425                    return Ok(false);
426                }
427                pending = next_private()?;
428            }
429            if let Some(write) = pending.filter(|write| write.key() == key) {
430                running = emit(key, P::private(write))?;
431                pending = next_private()?;
432            } else {
433                running = emit(key, record)?;
434            }
435            Ok(running)
436        })?;
437        while running {
438            let Some(write) = pending else {
439                break;
440            };
441            running = emit(write.key(), P::private(write))?;
442            pending = next_private()?;
443        }
444        Ok(())
445    }
446
447    pub fn get(
448        &self,
449        key: &[u8],
450        control: &StorageReadControl,
451    ) -> VersionResult<Option<PreparedRecordWrite>> {
452        control.cancellation().check()?;
453        Ok(self.records.get(key).map(|change| change.write.clone()))
454    }
455
456    /// Return a bounded ordered page of private replacements, including deletions.
457    pub fn scan(
458        &self,
459        prefix: &[u8],
460        after: Option<&[u8]>,
461        limit: usize,
462        control: &StorageReadControl,
463    ) -> VersionResult<BudgetedVec<PreparedRecordWrite>> {
464        control.cancellation().check()?;
465        let mut result = BudgetedVec::new(control.memory());
466        if limit == 0 {
467            return Ok(result);
468        }
469        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
470        for (key, change) in self.records.range_from(std::ops::Bound::Included(start)) {
471            control.cancellation().check()?;
472            let key = key.bytes();
473            if !key.starts_with(prefix) {
474                break;
475            }
476            if after.is_some_and(|after| key <= after) {
477                continue;
478            }
479            result.push(change.write.clone())?;
480            if result.len() == limit {
481                break;
482            }
483        }
484        Ok(result)
485    }
486}
487
488#[cfg(test)]
489mod tests {
490    use super::*;
491
492    #[test]
493    fn shared_changes_keep_one_strong_counter_until_the_final_owner_drops() {
494        let memory = MemoryBudget::new(1024);
495        let bytes = size_of::<Owner>() + size_of::<usize>();
496        let mut changes = None;
497        let allocated = allocation_counter::measure(|| {
498            changes = Some(PrivateRecordChanges::new(&memory));
499        });
500        assert_eq!(allocated.count_total, 1);
501        assert_eq!(allocated.bytes_total, bytes as u64);
502        let retained = changes.as_ref().unwrap().share_owner();
503        let shared = allocation_counter::measure(|| drop(changes));
504        assert_eq!(shared.count_total, 0);
505        assert_eq!(shared.bytes_current, 0);
506        let released = allocation_counter::measure(|| drop(retained));
507        assert_eq!(released.count_current, -1);
508        assert_eq!(released.bytes_current, -(bytes as i64));
509    }
510}
uqa_storage/mvcc/overlay.rs

uqa_storage/mvcc/
overlay.rs
