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

uqa_storage/mvcc/
overlay.rs
