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::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;
23pub(in crate::mvcc) mod run;
24mod run_set;
25mod tiers;
26use retained::{RetainedSources, Sources};
27use run::SpilledRunWriter;
28use run_set::RunSet;
29use tiers::{TieredChange, TieredCursor};
30
31/// The resident bytes of a memory-tier change beyond its key and value, as an estimate: its tree node and entry, the shared key and value handles and their allocation headers.
32const CHANGE_OVERHEAD: usize = 192;
33/// The memory tier spills once the transaction's allowance is more than this fraction used.
34const PRESSURE_DIVISOR: usize = 2;
35/// The memory tier spills only once it holds at least this fraction of the allowance, so that memory held elsewhere does not spill a run for every write.
36const RESIDENT_DIVISOR: usize = 16;
37
38struct Change {
39    write: PreparedRecordWrite,
40    identity: PrivateRecordRevision,
41}
42
43type Records = BudgetedSharedMap<RecordKey, Change>;
44
45/// Process-local identity of an evaluated private batch. Identities are never reused across transactions or undo branches; they are not durable commit sequences.
46#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
47pub struct PrivateRecordRevision(NonZeroU64);
48
49impl PrivateRecordRevision {
50    pub(in crate::mvcc) fn allocate() -> VersionResult<Self> {
51        static NEXT: AtomicU64 = AtomicU64::new(1);
52        NEXT.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
53            .map(|id| Self(NonZeroU64::new(id).expect("private revisions start at one")))
54            .map_err(|_| VersionError::PrivateRevisionExhausted)
55    }
56
57    #[cfg(test)]
58    pub(crate) fn for_tests() -> Self {
59        Self::allocate().expect("test private revision")
60    }
61
62    pub fn as_u64(self) -> u64 {
63        self.0.get()
64    }
65
66    /// Decode an identity written to a spill file of this process.
67    fn from_u64(value: u64) -> VersionResult<Self> {
68        NonZeroU64::new(value)
69            .map(Self)
70            .ok_or(VersionError::InvalidEncoding(
71                "zero private record revision",
72            ))
73    }
74}
75
76/// A changed key and its private batch identity, without retaining or copying its payload.
77pub struct PrivateRecordKey {
78    key: RecordKey,
79    revision: PrivateRecordRevision,
80}
81
82impl PrivateRecordKey {
83    pub fn key(&self) -> &[u8] {
84        self.key.bytes()
85    }
86
87    pub fn revision(&self) -> PrivateRecordRevision {
88        self.revision
89    }
90}
91
92struct Savepoint {
93    id: StorageSavepointId,
94    records: Records,
95    resident: usize,
96    runs: RunSet,
97    sources: RetainedSources,
98    revision: Option<PrivateRecordRevision>,
99}
100
101/// The changes of one transaction in two tiers: the newest in a persistent ordered root in memory, older ones in immutable spilled runs. A key's change in memory shadows its changes in runs, and a newer run shadows an older one.
102struct State {
103    records: Records,
104    /// The estimated bytes the memory tier holds, which decide when it spills.
105    resident: usize,
106    runs: RunSet,
107    sources: Sources,
108    revision: Option<PrivateRecordRevision>,
109    savepoints: BudgetedVec<Savepoint>,
110}
111
112impl State {
113    /// The newest change of `key`, from either tier.
114    fn change(
115        &self,
116        key: &[u8],
117        control: &StorageReadControl,
118    ) -> VersionResult<Option<TieredChange<'_>>> {
119        tiers::lookup(&self.records, &self.runs, key, control)
120    }
121
122    /// The smallest memory tier that spills, which is also the size of the smallest runs.
123    fn least_spill(&self) -> usize {
124        self.records.budget().limit() / RESIDENT_DIVISOR
125    }
126
127    /// Move the memory tier into a new run when the transaction's allowance is under pressure: more than half of it is used and the memory tier holds a share of it worth a run. The allowance measures what this tier actually holds, so changes that share their values with an earlier owner do not spill on their own. Failure leaves both tiers unchanged.
128    fn make_room(&mut self, control: &StorageReadControl) -> VersionResult<()> {
129        let budget = self.records.budget();
130        if self.records.is_empty()
131            || self.resident < self.least_spill()
132            || budget.used() <= budget.limit() / PRESSURE_DIVISOR
133        {
134            return Ok(());
135        }
136        let memory = self.records.budget().clone();
137        let mut writer =
138            SpilledRunWriter::new(self.records.len() as u64, self.resident as u64, &memory)?;
139        for (_, change) in &self.records {
140            writer.push(
141                change.write.key(),
142                change.write.expected(),
143                change.write.kind(),
144                change.identity,
145                change.write.value(),
146                control,
147            )?;
148        }
149        let run = writer
150            .finish()?
151            .expect("a nonempty memory tier spills a nonempty run");
152        let base = self.least_spill() as u64;
153        self.runs = self.runs.with_run(run, base, &memory, control)?;
154        self.records = Records::new(&memory);
155        self.resident = 0;
156        Ok(())
157    }
158
159    fn savepoint_position(&self, id: StorageSavepointId) -> VersionResult<usize> {
160        self.savepoints
161            .iter()
162            .rposition(|savepoint| savepoint.id == id)
163            .ok_or(VersionError::SavepointMissing(id))
164    }
165
166    fn truncate_savepoints(&mut self, len: usize) {
167        if len == 0 {
168            self.savepoints = BudgetedVec::new(self.records.budget());
169        } else {
170            self.savepoints.truncate(len);
171        }
172    }
173}
174
175struct Owner {
176    state: Mutex<State>,
177}
178
179/// One transaction's record changes. No provider lock or native transaction is retained between operations.
180///
181/// Command views and savepoints share immutable ordered roots and the spilled runs beneath them. 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. Once more than half of the transaction's allowance is used, the memory tier moves into a run in an encrypted temporary file, so the size of a transaction is bounded by disk rather than by its allowance.
182pub struct PrivateRecordChanges {
183    owner: StrongArc<Owner>,
184}
185
186impl PrivateRecordChanges {
187    pub(super) fn share_owner(&self) -> Self {
188        Self {
189            owner: StrongArc::clone(&self.owner),
190        }
191    }
192
193    pub(super) fn write_kind(
194        &self,
195        key: &[u8],
196        control: &StorageReadControl,
197    ) -> VersionResult<Option<super::commit::RecordWriteKind>> {
198        control.check()?;
199        let state = self.owner.state.lock();
200        Ok(state.change(key, control)?.map(|change| change.kind()))
201    }
202
203    pub fn new(memory: &MemoryBudget) -> Self {
204        Self {
205            owner: StrongArc::new(Owner {
206                state: Mutex::new(State {
207                    records: Records::new(memory),
208                    resident: 0,
209                    runs: RunSet::empty(),
210                    sources: Sources::new(memory),
211                    revision: None,
212                    savepoints: BudgetedVec::new(memory),
213                }),
214            }),
215        }
216    }
217
218    /// Stage one atomic batch of already evaluated replacements. Repeated changes to a private record must preserve its original committed precondition.
219    pub fn apply(
220        &self,
221        writes: &[RecordWrite<'_>],
222        control: &StorageReadControl,
223    ) -> VersionResult<()> {
224        let memory = self.owner.state.lock().records.budget().clone();
225        let owned_control = StorageReadControl::new(&memory, control.cancellation());
226        let prepared = PreparedRecordCommit::new(writes, &owned_control)?;
227        let writes = prepared.resident().ok_or(VersionError::InvalidEncoding(
228            "a supplied batch is in memory",
229        ))?;
230        self.apply_owned(writes, control)
231    }
232
233    pub(crate) fn apply_owned(
234        &self,
235        writes: &[PreparedRecordWrite],
236        control: &StorageReadControl,
237    ) -> VersionResult<()> {
238        if writes.is_empty() {
239            return Ok(());
240        }
241        let mut state = self.owner.state.lock();
242        let mut incoming = 0_usize;
243        for (index, write) in writes.iter().enumerate() {
244            control.cancellation().check()?;
245            if let Some(previous) = state.change(write.key(), control)? {
246                if previous.expected() != write.expected() {
247                    return Err(VersionError::WriteConflict {
248                        mutation: index,
249                        expected: write.expected(),
250                        actual: previous.expected(),
251                    });
252                }
253            }
254            incoming = incoming.saturating_add(resident_bytes(write));
255        }
256        state.make_room(control)?;
257        let identity = PrivateRecordRevision::allocate()?;
258        let invalidates_source = writes
259            .iter()
260            .any(|write| state.sources.get(write.key()).is_some_and(Option::is_some));
261        if writes.len() == 1 && !invalidates_source {
262            let write = &writes[0];
263            control.cancellation().check()?;
264            state.records.try_insert(
265                write.shared_key(),
266                Change {
267                    write: write.clone(),
268                    identity,
269                },
270            )?;
271            state.resident = state.resident.saturating_add(incoming);
272            state.revision = Some(identity);
273            return Ok(());
274        }
275        let mut records = state.records.clone();
276        let mut sources = state.sources.clone();
277        for write in writes {
278            control.cancellation().check()?;
279            records.try_insert(
280                write.shared_key(),
281                Change {
282                    write: write.clone(),
283                    identity,
284                },
285            )?;
286            if sources.get(write.key()).is_some_and(Option::is_some) {
287                sources.try_insert(write.shared_key(), None)?;
288            }
289        }
290        control.cancellation().check()?;
291        // Candidate roots own every reservation before this single atomic publication.
292        state.records = records;
293        state.sources = sources;
294        state.resident = state.resident.saturating_add(incoming);
295        state.revision = Some(identity);
296        Ok(())
297    }
298
299    /// Stage every write of `prepared` one at a time, so that a batch larger than the allowance spills as it is staged. Unlike `apply`, a failure can leave a prefix of the writes staged; it serves changes that are discarded when it fails.
300    pub(in crate::mvcc) fn apply_prepared(
301        &self,
302        prepared: &PreparedRecordCommit,
303        control: &StorageReadControl,
304    ) -> VersionResult<()> {
305        let mut writes = prepared.writes();
306        while let Some(write) = writes.next(control)? {
307            self.apply_owned(std::slice::from_ref(&write), control)?;
308        }
309        Ok(())
310    }
311
312    pub fn has_written(&self) -> bool {
313        let state = self.owner.state.lock();
314        !state.records.is_empty() || !state.runs.is_empty()
315    }
316
317    pub fn snapshot(&self) -> VersionResult<PrivateRecordSnapshot> {
318        let state = self.owner.state.lock();
319        let memory = state
320            .records
321            .budget()
322            .reserve(std::mem::size_of::<PrivateRecordSnapshot>())?;
323        Ok(PrivateRecordSnapshot {
324            records: state.records.clone(),
325            runs: state.runs.clone(),
326            sources: state.sources.snapshot(),
327            revision: state.revision,
328            _memory: memory,
329        })
330    }
331
332    pub fn savepoint(&self, id: StorageSavepointId) -> VersionResult<()> {
333        let mut state = self.owner.state.lock();
334        let records = state.records.clone();
335        let resident = state.resident;
336        let runs = state.runs.clone();
337        let sources = state.sources.snapshot();
338        let revision = state.revision;
339        state.savepoints.push(Savepoint {
340            id,
341            records,
342            resident,
343            runs,
344            sources,
345            revision,
346        })?;
347        Ok(())
348    }
349
350    /// Release the nearest matching identity and its nested savepoints, retaining every write.
351    pub fn release_savepoint(&self, id: StorageSavepointId) -> VersionResult<()> {
352        let mut state = self.owner.state.lock();
353        let position = state.savepoint_position(id)?;
354        state.truncate_savepoints(position);
355        Ok(())
356    }
357
358    /// Undo only this owner's changes, retaining the target savepoint for another rollback.
359    pub fn rollback_to_savepoint(&self, id: StorageSavepointId) -> VersionResult<()> {
360        let mut state = self.owner.state.lock();
361        let position = state.savepoint_position(id)?;
362        state.records = state.savepoints[position].records.clone();
363        state.resident = state.savepoints[position].resident;
364        state.runs = state.savepoints[position].runs.clone();
365        let sources = state.savepoints[position].sources.clone();
366        state.sources.restore(&sources);
367        state.revision = state.savepoints[position].revision;
368        state.truncate_savepoints(position + 1);
369        Ok(())
370    }
371
372    pub fn rollback(&self) -> VersionResult<()> {
373        let mut state = self.owner.state.lock();
374        let memory = state.records.budget().clone();
375        state.records = Records::new(&memory);
376        state.resident = 0;
377        state.runs = RunSet::empty();
378        state.sources = Sources::new(&memory);
379        state.revision = None;
380        state.truncate_savepoints(0);
381        Ok(())
382    }
383
384    /// Capture immutable final replacements without reevaluating or copying their payloads. Changes that spilled are merged into one spilled run of final replacements, which commits read one block at a time.
385    pub fn prepare(&self, control: &StorageReadControl) -> VersionResult<PreparedRecordCommit> {
386        control.cancellation().check()?;
387        let state = self.owner.state.lock();
388        let mut changes = TieredCursor::new(
389            Some(&state.records),
390            &state.runs,
391            std::ops::Bound::Unbounded,
392            control,
393        )?;
394        if state.runs.is_empty() {
395            let mut writes = BudgetedVec::new(control.memory());
396            writes.reserve(state.records.len())?;
397            while let Some(change) = changes.next(control)? {
398                writes.push(change.write(control)?)?;
399            }
400            return PreparedRecordCommit::from_unique_owned(writes, control);
401        }
402        let (entries, entry_bytes) = state.runs.size();
403        let mut writer = SpilledRunWriter::new(
404            entries.saturating_add(state.records.len() as u64),
405            entry_bytes.saturating_add(state.resident as u64),
406            control.memory(),
407        )?;
408        while let Some(change) = changes.next(control)? {
409            match &change {
410                TieredChange::Resident(resident) => writer.push(
411                    resident.write.key(),
412                    resident.write.expected(),
413                    resident.write.kind(),
414                    resident.identity,
415                    resident.write.value(),
416                    control,
417                )?,
418                TieredChange::Spilled { entry, run } => writer.push_spilled(entry, run, control)?,
419            }
420        }
421        let run = writer.finish()?.expect("a spilled transaction has changes");
422        PreparedRecordCommit::from_spilled_run(run, control)
423    }
424}
425
426/// Fixed private visibility for a command or retained source cursor; tombstones remain distinguishable from an unchanged key.
427pub struct PrivateRecordSnapshot {
428    records: Records,
429    runs: RunSet,
430    sources: RetainedSources,
431    revision: Option<PrivateRecordRevision>,
432    _memory: MemoryReservation,
433}
434
435impl PrivateRecordSnapshot {
436    /// Identity of this complete private root. Undo restores the saved identity; successful replacement batches receive identities never used by another root or transaction.
437    pub fn revision(&self) -> Option<PrivateRecordRevision> {
438        self.revision
439    }
440
441    pub(super) fn last_before(
442        &self,
443        prefix: &[u8],
444        before: Option<&[u8]>,
445        control: &StorageReadControl,
446    ) -> VersionResult<Option<PreparedRecordWrite>> {
447        control.check()?;
448        let mut upper = BudgetedVec::new(control.memory());
449        upper.extend_from_slice(prefix)?;
450        let upper = match upper.iter().rposition(|byte| *byte != u8::MAX) {
451            Some(last) => {
452                upper[last] += 1;
453                upper.truncate(last + 1);
454                Some(upper)
455            }
456            None => None,
457        };
458        let end = match (before, upper.as_deref()) {
459            (Some(before), Some(upper)) => Some(before.min(upper)),
460            (before, upper) => before.or(upper),
461        };
462        let bound = end.map_or(std::ops::Bound::Unbounded, std::ops::Bound::Excluded);
463        let resident = self
464            .records
465            .last_before::<[u8]>(bound)
466            .map(|(_, change)| TieredChange::Resident(change));
467        let spilled = self.runs.last_before(bound, control)?;
468        // The memory tier is newer, so it wins a tie.
469        let last = match (resident, spilled) {
470            (Some(resident), Some(spilled)) if spilled.key() > resident.key() => Some(spilled),
471            (Some(resident), _) => Some(resident),
472            (None, spilled) => spilled,
473        };
474        last.filter(|change| change.key().starts_with(prefix))
475            .map(|change| change.write(control))
476            .transpose()
477    }
478
479    /// Retain this exact private revision without following later writes or copying its values.
480    pub(super) fn try_clone(&self) -> VersionResult<Self> {
481        let memory = self.records.budget().reserve(std::mem::size_of::<Self>())?;
482        Ok(Self {
483            records: self.records.clone(),
484            runs: self.runs.clone(),
485            sources: self.sources.clone(),
486            revision: self.revision,
487            _memory: memory,
488        })
489    }
490
491    /// Select changed keys on this retained command boundary, including deletions. Undo restores their earlier identities, while later branches and other transactions receive distinct identities.
492    pub fn scan_keys(
493        &self,
494        prefix: &[u8],
495        after: Option<&[u8]>,
496        limit: usize,
497        control: &StorageReadControl,
498    ) -> VersionResult<BudgetedVec<PrivateRecordKey>> {
499        control.check()?;
500        let mut result = BudgetedVec::new(control.memory());
501        if limit == 0 {
502            return Ok(result);
503        }
504        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
505        let mut changes = TieredCursor::new(
506            Some(&self.records),
507            &self.runs,
508            std::ops::Bound::Included(start),
509            control,
510        )?;
511        while let Some(change) = changes.next(control)? {
512            control.check()?;
513            let key = change.key();
514            if !key.starts_with(prefix) {
515                break;
516            }
517            if after.is_some_and(|after| key <= after) {
518                continue;
519            }
520            result.push(PrivateRecordKey {
521                key: change.shared_key(),
522                revision: change.identity(),
523            })?;
524            if result.len() == limit {
525                break;
526            }
527        }
528        Ok(result)
529    }
530
531    pub(super) fn visit_merged<P: super::projection::Projection>(
532        &self,
533        committed: &dyn super::CommittedRecordSnapshot,
534        prefix: &[u8],
535        after: Option<&[u8]>,
536        limit: usize,
537        control: &StorageReadControl,
538        visit: &mut super::projection::Visitor<'_, P>,
539    ) -> VersionResult<()> {
540        control.cancellation().check()?;
541        if limit == 0 {
542            return Ok(());
543        }
544        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
545        let mut changes = TieredCursor::new(
546            Some(&self.records),
547            &self.runs,
548            std::ops::Bound::Included(start),
549            control,
550        )?;
551        let mut next_private = || -> VersionResult<Option<PreparedRecordWrite>> {
552            while let Some(change) = changes.next(control)? {
553                control.cancellation().check()?;
554                if !change.key().starts_with(prefix) {
555                    return Ok(None);
556                }
557                if after.is_some_and(|after| change.key() <= after) {
558                    continue;
559                }
560                return change.write(control).map(Some);
561            }
562            Ok(None)
563        };
564        let mut pending = next_private()?;
565        let mut count = 0;
566        let mut running = true;
567        let mut emit = |key: &[u8], record: P::Record<'_>| -> VersionResult<bool> {
568            control.cancellation().check()?;
569            count += 1;
570            let more = visit(key, record)?;
571            control.cancellation().check()?;
572            Ok(more && count < limit)
573        };
574        P::visit(committed, prefix, after, control, &mut |key, record| {
575            while let Some(write) = pending.as_ref().filter(|write| write.key() < key) {
576                running = emit(write.key(), P::private(write))?;
577                if !running {
578                    return Ok(false);
579                }
580                pending = next_private()?;
581            }
582            if let Some(write) = pending.as_ref().filter(|write| write.key() == key) {
583                running = emit(key, P::private(write))?;
584                pending = next_private()?;
585            } else {
586                running = emit(key, record)?;
587            }
588            Ok(running)
589        })?;
590        while running {
591            let Some(write) = pending.as_ref() else {
592                break;
593            };
594            running = emit(write.key(), P::private(write))?;
595            pending = next_private()?;
596        }
597        Ok(())
598    }
599
600    pub fn get(
601        &self,
602        key: &[u8],
603        control: &StorageReadControl,
604    ) -> VersionResult<Option<PreparedRecordWrite>> {
605        control.cancellation().check()?;
606        tiers::lookup(&self.records, &self.runs, key, control)?
607            .map(|change| change.write(control))
608            .transpose()
609    }
610
611    /// Return a bounded ordered page of private replacements, including deletions.
612    pub fn scan(
613        &self,
614        prefix: &[u8],
615        after: Option<&[u8]>,
616        limit: usize,
617        control: &StorageReadControl,
618    ) -> VersionResult<BudgetedVec<PreparedRecordWrite>> {
619        control.cancellation().check()?;
620        let mut result = BudgetedVec::new(control.memory());
621        if limit == 0 {
622            return Ok(result);
623        }
624        let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
625        let mut changes = TieredCursor::new(
626            Some(&self.records),
627            &self.runs,
628            std::ops::Bound::Included(start),
629            control,
630        )?;
631        while let Some(change) = changes.next(control)? {
632            control.cancellation().check()?;
633            let key = change.key();
634            if !key.starts_with(prefix) {
635                break;
636            }
637            if after.is_some_and(|after| key <= after) {
638                continue;
639            }
640            result.push(change.write(control)?)?;
641            if result.len() == limit {
642                break;
643            }
644        }
645        Ok(result)
646    }
647}
648
649/// The estimated memory a change takes in the memory tier.
650fn resident_bytes(write: &PreparedRecordWrite) -> usize {
651    write
652        .key()
653        .len()
654        .saturating_add(write.value().map_or(0, <[u8]>::len))
655        .saturating_add(CHANGE_OVERHEAD)
656}
657
658#[cfg(test)]
659mod tests;