Skip to main content

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