1use 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
31const CHANGE_OVERHEAD: usize = 192;
33const PRESSURE_DIVISOR: usize = 2;
35const RESIDENT_DIVISOR: usize = 16;
37
38struct Change {
39 write: PreparedRecordWrite,
40 identity: PrivateRecordRevision,
41}
42
43type Records = BudgetedSharedMap<RecordKey, Change>;
44
45#[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 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
76pub 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
101struct State {
103 records: Records,
104 resident: usize,
106 runs: RunSet,
107 sources: Sources,
108 revision: Option<PrivateRecordRevision>,
109 savepoints: BudgetedVec<Savepoint>,
110}
111
112impl State {
113 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 fn least_spill(&self) -> usize {
124 self.records.budget().limit() / RESIDENT_DIVISOR
125 }
126
127 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
179pub 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 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 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 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 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 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 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
426pub struct PrivateRecordSnapshot {
428 records: Records,
429 runs: RunSet,
430 sources: RetainedSources,
431 revision: Option<PrivateRecordRevision>,
432 _memory: MemoryReservation,
433}
434
435impl PrivateRecordSnapshot {
436 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 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 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 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 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
649fn 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;