1use std::{
5 cmp::{Ordering, Reverse},
6 collections::{HashMap, HashSet},
7 ops::Bound,
8 sync::Arc,
9};
10
11use reifydb_codec::key::encoded::EncodedKey;
12use reifydb_core::{common::CommitVersion, interface::store::EntryKind};
13use reifydb_value::{Result, reifydb_assertions, util::cowvec::CowVec};
14use tracing::{Span, field, instrument};
15
16use crate::{
17 MultiVersionScope,
18 tier::{
19 HistoricalCursor, RangeBatch, RangeCursor, RawEntry, TierBackend, TierBatch, TierStorage,
20 VersionedGetResult,
21 commit::memory::entry::{CurrentMap, Entries, Entry, HistoricalMap},
22 },
23};
24
25type EvictablePersist = Vec<(EncodedKey, CommitVersion, Option<CowVec<u8>>)>;
26type EvictableDrop = Vec<(EncodedKey, CommitVersion)>;
27
28#[derive(Clone)]
29pub struct MemoryPrimitiveStorage {
30 inner: Arc<MemoryPrimitiveStorageInner>,
31}
32
33struct MemoryPrimitiveStorageInner {
34 entries: Entries,
35}
36
37impl Default for MemoryPrimitiveStorage {
38 fn default() -> Self {
39 Self::new()
40 }
41}
42
43impl MemoryPrimitiveStorage {
44 #[instrument(name = "store::multi::memory::new", level = "debug")]
45 pub fn new() -> Self {
46 Self {
47 inner: Arc::new(MemoryPrimitiveStorageInner {
48 entries: Entries::default(),
49 }),
50 }
51 }
52
53 pub fn count_current(&self, table: EntryKind) -> Result<u64> {
54 Ok(self.inner.entries.data.get(&table).map(|e| e.current.read().len() as u64).unwrap_or(0))
55 }
56
57 pub fn list_all_entry_kinds(&self) -> Result<Vec<EntryKind>> {
58 Ok(self.inner.entries.data.keys())
59 }
60
61 pub fn count_historical(&self, table: EntryKind) -> Result<u64> {
62 Ok(self.inner
63 .entries
64 .data
65 .get(&table)
66 .map(|e| {
67 let hist = e.historical.read();
68 hist.values().map(|m| m.len() as u64).sum()
69 })
70 .unwrap_or(0))
71 }
72
73 #[inline]
74 #[instrument(name = "store::multi::memory::get_or_create_table", level = "trace", skip(self), fields(table = ?table))]
75 fn get_or_create_table(&self, table: EntryKind) -> Entry {
76 self.inner.entries.data.get_or_insert_with(table, Entry::new)
77 }
78
79 #[inline]
80 #[instrument(name = "store::multi::memory::set::table", level = "trace", skip(self, entries), fields(
81 table = ?table,
82 entry_count = entries.len(),
83 ))]
84 fn process_table(
85 &self,
86 table: EntryKind,
87 version: CommitVersion,
88 entries: Vec<(EncodedKey, Option<CowVec<u8>>)>,
89 ) {
90 let table_entry = self.get_or_create_table(table);
91 let mut current = table_entry.current.write();
92 let mut historical = table_entry.historical.write();
93
94 for (key, value) in entries {
95 if let Some((pre_version, pre_value)) = current.get(&key) {
96 if *pre_version < version {
97 let pre_version = *pre_version;
98 let pre_value = pre_value.clone();
99 reifydb_assertions! {
100 assert!(
101 version.0 > pre_version.0,
102 "promoting current entry to historical requires the incoming version to exceed it, otherwise the same version appears in both tiers and point-reads return the wrong entry (version={} pre_version={})",
103 version.0,
104 pre_version.0
105 );
106 }
107 historical
108 .entry(key.clone())
109 .or_default()
110 .insert(Reverse(pre_version), pre_value);
111
112 current.insert(key, (version, value));
113 } else {
114 historical.entry(key).or_default().insert(Reverse(version), value);
115 }
116 } else {
117 current.insert(key, (version, value));
118 }
119 }
120 }
121
122 pub fn collect_evictable_below(
123 &self,
124 table: EntryKind,
125 cutoff: CommitVersion,
126 ) -> (EvictablePersist, EvictableDrop) {
127 let entry = match self.inner.entries.data.get(&table) {
128 Some(e) => e,
129 None => return (Vec::new(), Vec::new()),
130 };
131 let current = entry.current.read();
132 let historical = entry.historical.read();
133
134 let historical_entries: usize = historical.values().map(|versions| versions.len()).sum();
135
136 let mut latest: HashMap<EncodedKey, (CommitVersion, Option<CowVec<u8>>)> =
137 HashMap::with_capacity(current.len() + historical.len());
138 let mut to_drop: Vec<(EncodedKey, CommitVersion)> =
139 Vec::with_capacity(current.len() + historical_entries);
140
141 for (key, (v, val)) in current.iter() {
142 if *v <= cutoff {
143 to_drop.push((key.clone(), *v));
144 latest.insert(key.clone(), (*v, val.clone()));
145 }
146 }
147 for (key, versions) in historical.iter() {
148 for (Reverse(v), val) in versions.iter() {
149 if *v <= cutoff {
150 to_drop.push((key.clone(), *v));
151 match latest.get(key) {
152 Some((best, _)) if *best >= *v => {}
153 _ => {
154 latest.insert(key.clone(), (*v, val.clone()));
155 }
156 }
157 }
158 }
159 }
160
161 let to_persist = latest.into_iter().map(|(key, (v, val))| (key, v, val)).collect();
162 (to_persist, to_drop)
163 }
164}
165
166impl TierStorage for MemoryPrimitiveStorage {
167 #[instrument(name = "store::multi::memory::get", level = "trace", skip(self, key), fields(table = ?table, key_len = key.len(), version = version.0))]
168 fn get(&self, table: EntryKind, key: &[u8], version: CommitVersion) -> Result<VersionedGetResult> {
169 let entry = match self.inner.entries.data.get(&table) {
170 Some(e) => e,
171 None => return Ok(VersionedGetResult::NotFound),
172 };
173
174 let current = entry.current.read();
175 if let Some((cur_version, value)) = current.get(key)
176 && *cur_version <= version
177 {
178 return Ok(match value {
179 Some(v) => VersionedGetResult::Value {
180 value: v.clone(),
181 version: *cur_version,
182 },
183 None => VersionedGetResult::Tombstone,
184 });
185 }
186 drop(current);
187
188 let historical = entry.historical.read();
189 if let Some(versions) = historical.get(key) {
190 for (Reverse(v), value) in versions.range(Reverse(version)..) {
191 if *v <= version {
192 return Ok(match value {
193 Some(val) => VersionedGetResult::Value {
194 value: val.clone(),
195 version: *v,
196 },
197 None => VersionedGetResult::Tombstone,
198 });
199 }
200 }
201 }
202
203 Ok(VersionedGetResult::NotFound)
204 }
205
206 #[instrument(name = "store::multi::memory::contains", level = "trace", skip(self, key), fields(table = ?table, key_len = key.len(), version = version.0), ret)]
207 fn contains(&self, table: EntryKind, key: &[u8], version: CommitVersion) -> Result<bool> {
208 let entry = match self.inner.entries.data.get(&table) {
209 Some(e) => e,
210 None => return Ok(false),
211 };
212
213 let current = entry.current.read();
214 if let Some((cur_version, value)) = current.get(key)
215 && *cur_version <= version
216 {
217 return Ok(value.is_some());
218 }
219 drop(current);
220
221 let historical = entry.historical.read();
222 if let Some(versions) = historical.get(key) {
223 for (Reverse(v), value) in versions.range(Reverse(version)..) {
224 if *v <= version {
225 return Ok(value.is_some());
226 }
227 }
228 }
229
230 Ok(false)
231 }
232
233 #[instrument(name = "store::multi::memory::set", level = "trace", skip(self, batches), fields(
234 table_count = batches.len(),
235 total_entry_count = field::Empty,
236 version = version.0
237 ))]
238 fn set(&self, version: CommitVersion, batches: TierBatch) -> Result<()> {
239 let total_entries: usize = batches.values().map(|v| v.len()).sum();
240
241 batches.into_iter().for_each(|(table, entries)| {
242 self.process_table(table, version, entries);
243 });
244
245 Span::current().record("total_entry_count", total_entries);
246 Ok(())
247 }
248
249 #[instrument(name = "store::multi::memory::range_next", level = "trace", skip(self, cursor, start, end), fields(table = ?table, batch_size = batch_size, scope = ?scope))]
250 fn range_next(
251 &self,
252 table: EntryKind,
253 cursor: &mut RangeCursor,
254 start: Bound<&[u8]>,
255 end: Bound<&[u8]>,
256 scope: MultiVersionScope,
257 batch_size: usize,
258 ) -> Result<RangeBatch> {
259 if cursor.exhausted {
260 return Ok(RangeBatch::empty());
261 }
262
263 let entry = match self.inner.entries.data.get(&table) {
264 Some(e) => e,
265 None => {
266 cursor.exhausted = true;
267 return Ok(RangeBatch::empty());
268 }
269 };
270
271 let cursor_key = cursor.last_key.clone();
272
273 let current = entry.current.read();
274 let historical = entry.historical.read();
275
276 let mut entries: Vec<RawEntry> = Vec::with_capacity(batch_size + 1);
277
278 let iter_start: Bound<&[u8]> = match &cursor_key {
279 Some(last) => Bound::Excluded(last.as_slice()),
280 None => start,
281 };
282
283 let iter_end: Bound<&[u8]> = end;
284
285 let mut cur_iter = current.range::<[u8], _>((iter_start, iter_end)).peekable();
286 let mut hist_iter = historical.range::<[u8], _>((iter_start, iter_end)).peekable();
287
288 while entries.len() <= batch_size {
289 let (take_cur, take_hist) = match (cur_iter.peek(), hist_iter.peek()) {
290 (None, None) => break,
291 (Some(_), None) => (true, false),
292 (None, Some(_)) => (false, true),
293 (Some((kc, _)), Some((kh, _))) => match kc.cmp(kh) {
294 Ordering::Less => (true, false),
295 Ordering::Greater => (false, true),
296 Ordering::Equal => (true, true),
297 },
298 };
299
300 if take_cur && take_hist {
301 let (key, (cur_version, cur_value)) = cur_iter.next().unwrap();
302 let (_, versions) = hist_iter.next().unwrap();
303 if scope.contains(*cur_version) {
304 entries.push(RawEntry {
305 key: key.clone(),
306 version: *cur_version,
307 value: cur_value.clone(),
308 });
309 } else if *cur_version > scope.read() {
310 for (Reverse(v), value) in versions.range(Reverse(scope.read())..) {
311 if scope.contains(*v) {
312 entries.push(RawEntry {
313 key: key.clone(),
314 version: *v,
315 value: value.clone(),
316 });
317 break;
318 }
319 if let MultiVersionScope::Between {
320 after,
321 ..
322 } = scope && *v <= after
323 {
324 break;
325 }
326 }
327 }
328 } else if take_cur {
329 let (key, (cur_version, cur_value)) = cur_iter.next().unwrap();
330 if scope.contains(*cur_version) {
331 entries.push(RawEntry {
332 key: key.clone(),
333 version: *cur_version,
334 value: cur_value.clone(),
335 });
336 }
337 } else {
338 let (key, versions) = hist_iter.next().unwrap();
339 for (Reverse(v), value) in versions.range(Reverse(scope.read())..) {
340 if scope.contains(*v) {
341 entries.push(RawEntry {
342 key: key.clone(),
343 version: *v,
344 value: value.clone(),
345 });
346 break;
347 }
348 if let MultiVersionScope::Between {
349 after,
350 ..
351 } = scope && *v <= after
352 {
353 break;
354 }
355 }
356 }
357 }
358
359 let has_more = entries.len() > batch_size;
360 if has_more {
361 entries.truncate(batch_size);
362 }
363
364 if let Some(last_entry) = entries.last() {
365 cursor.last_key = Some(last_entry.key.clone());
366 }
367 if !has_more {
368 cursor.exhausted = true;
369 }
370
371 Ok(RangeBatch {
372 entries,
373 has_more,
374 })
375 }
376
377 #[instrument(name = "store::multi::memory::range_rev_next", level = "trace", skip(self, cursor, start, end), fields(table = ?table, batch_size = batch_size, scope = ?scope))]
378 fn range_rev_next(
379 &self,
380 table: EntryKind,
381 cursor: &mut RangeCursor,
382 start: Bound<&[u8]>,
383 end: Bound<&[u8]>,
384 scope: MultiVersionScope,
385 batch_size: usize,
386 ) -> Result<RangeBatch> {
387 if cursor.exhausted {
388 return Ok(RangeBatch::empty());
389 }
390
391 let entry = match self.inner.entries.data.get(&table) {
392 Some(e) => e,
393 None => {
394 cursor.exhausted = true;
395 return Ok(RangeBatch::empty());
396 }
397 };
398
399 let cursor_key = cursor.last_key.clone();
400
401 let current = entry.current.read();
402 let historical = entry.historical.read();
403
404 let mut entries: Vec<RawEntry> = Vec::with_capacity(batch_size + 1);
405
406 let iter_start: Bound<&[u8]> = start;
407
408 let iter_end: Bound<&[u8]> = match &cursor_key {
409 Some(last) => Bound::Excluded(last.as_slice()),
410 None => end,
411 };
412
413 let mut cur_iter = current.range::<[u8], _>((iter_start, iter_end)).rev().peekable();
414 let mut hist_iter = historical.range::<[u8], _>((iter_start, iter_end)).rev().peekable();
415
416 while entries.len() <= batch_size {
417 let (take_cur, take_hist) = match (cur_iter.peek(), hist_iter.peek()) {
418 (None, None) => break,
419 (Some(_), None) => (true, false),
420 (None, Some(_)) => (false, true),
421 (Some((kc, _)), Some((kh, _))) => match kc.cmp(kh) {
422 Ordering::Greater => (true, false),
423 Ordering::Less => (false, true),
424 Ordering::Equal => (true, true),
425 },
426 };
427
428 if take_cur && take_hist {
429 let (key, (cur_version, cur_value)) = cur_iter.next().unwrap();
430 let (_, versions) = hist_iter.next().unwrap();
431 if scope.contains(*cur_version) {
432 entries.push(RawEntry {
433 key: key.clone(),
434 version: *cur_version,
435 value: cur_value.clone(),
436 });
437 } else if *cur_version > scope.read() {
438 for (Reverse(v), value) in versions.range(Reverse(scope.read())..) {
439 if scope.contains(*v) {
440 entries.push(RawEntry {
441 key: key.clone(),
442 version: *v,
443 value: value.clone(),
444 });
445 break;
446 }
447 if let MultiVersionScope::Between {
448 after,
449 ..
450 } = scope && *v <= after
451 {
452 break;
453 }
454 }
455 }
456 } else if take_cur {
457 let (key, (cur_version, cur_value)) = cur_iter.next().unwrap();
458 if scope.contains(*cur_version) {
459 entries.push(RawEntry {
460 key: key.clone(),
461 version: *cur_version,
462 value: cur_value.clone(),
463 });
464 }
465 } else {
466 let (key, versions) = hist_iter.next().unwrap();
467 for (Reverse(v), value) in versions.range(Reverse(scope.read())..) {
468 if scope.contains(*v) {
469 entries.push(RawEntry {
470 key: key.clone(),
471 version: *v,
472 value: value.clone(),
473 });
474 break;
475 }
476 if let MultiVersionScope::Between {
477 after,
478 ..
479 } = scope && *v <= after
480 {
481 break;
482 }
483 }
484 }
485 }
486
487 let has_more = entries.len() > batch_size;
488 if has_more {
489 entries.truncate(batch_size);
490 }
491
492 if let Some(last_entry) = entries.last() {
493 cursor.last_key = Some(last_entry.key.clone());
494 }
495 if !has_more {
496 cursor.exhausted = true;
497 }
498
499 Ok(RangeBatch {
500 entries,
501 has_more,
502 })
503 }
504
505 #[instrument(name = "store::multi::memory::ensure_table", level = "trace", skip(self), fields(table = ?table))]
506 fn ensure_table(&self, table: EntryKind) -> Result<()> {
507 let _ = self.get_or_create_table(table);
508 Ok(())
509 }
510
511 #[instrument(name = "store::multi::memory::clear_table", level = "debug", skip(self), fields(table = ?table))]
512 fn clear_table(&self, table: EntryKind) -> Result<()> {
513 if let Some(entry) = self.inner.entries.data.get(&table) {
514 *entry.current.write() = CurrentMap::new();
515 *entry.historical.write() = HistoricalMap::new();
516 }
517 Ok(())
518 }
519
520 #[instrument(name = "store::multi::memory::drop", level = "debug", skip(self, batches), fields(
521 table_count = batches.len(),
522 total_entry_count = field::Empty
523 ))]
524 fn drop(&self, batches: HashMap<EntryKind, Vec<(EncodedKey, CommitVersion)>>) -> Result<()> {
525 let total_entries: usize = batches.values().map(|v| v.len()).sum();
526
527 for (table, entries) in batches {
528 let table_entry = self.get_or_create_table(table);
529 let mut current = table_entry.current.write();
530 let mut historical = table_entry.historical.write();
531
532 let mut by_key: HashMap<EncodedKey, Vec<CommitVersion>> = HashMap::new();
533 for (key, version) in entries {
534 by_key.entry(key).or_default().push(version);
535 }
536
537 for (key, dropped_versions) in by_key {
538 let dropped_set: HashSet<CommitVersion> = dropped_versions.iter().copied().collect();
539
540 let cur_version = current.get(&key).map(|(v, _)| *v);
541 let stored_hist_covered = historical
542 .get(&key)
543 .map(|m| m.keys().all(|Reverse(v)| dropped_set.contains(v)))
544 .unwrap_or(true);
545 let stored_cur_covered = cur_version.is_none_or(|v| dropped_set.contains(&v));
546
547 if stored_cur_covered && stored_hist_covered {
548 current.remove(&key);
549 historical.remove(&key);
550 continue;
551 }
552
553 for version in dropped_versions {
554 let cur_matches = current.get(&key).map(|(v, _)| *v) == Some(version);
555 if cur_matches {
556 let popped = historical.get_mut(&key).and_then(|v| v.pop_first());
557 let now_empty = historical.get(&key).is_some_and(|v| v.is_empty());
558 if now_empty {
559 historical.remove(&key);
560 }
561 match popped {
562 Some((Reverse(promoted_v), promoted_value)) => {
563 current.insert(
564 key.clone(),
565 (promoted_v, promoted_value),
566 );
567 }
568 None => {
569 current.remove(&key);
570 }
571 }
572 } else {
573 let now_empty = if let Some(versions) = historical.get_mut(&key) {
574 versions.remove(&Reverse(version));
575 versions.is_empty()
576 } else {
577 false
578 };
579 if now_empty {
580 historical.remove(&key);
581 }
582 }
583 }
584 }
585 }
586
587 Span::current().record("total_entry_count", total_entries);
588 Ok(())
589 }
590
591 #[instrument(name = "store::multi::memory::get_all_versions", level = "trace", skip(self, key), fields(table = ?table, key_len = key.len()))]
592 fn get_all_versions(&self, table: EntryKind, key: &[u8]) -> Result<Vec<(CommitVersion, Option<CowVec<u8>>)>> {
593 let entry = match self.inner.entries.data.get(&table) {
594 Some(e) => e,
595 None => return Ok(Vec::new()),
596 };
597
598 let current = entry.current.read();
599 let current_hit = current.get(key).map(|(cur_version, value)| (*cur_version, value.clone()));
600 drop(current);
601
602 let historical = entry.historical.read();
603 let hist_versions = historical.get(key);
604
605 let mut versions: Vec<(CommitVersion, Option<CowVec<u8>>)> =
606 Vec::with_capacity(current_hit.is_some() as usize + hist_versions.map_or(0, |v| v.len()));
607 if let Some(hit) = current_hit {
608 versions.push(hit);
609 }
610 if let Some(hist_versions) = hist_versions {
611 for (Reverse(v), value) in hist_versions.iter() {
612 versions.push((*v, value.clone()));
613 }
614 }
615
616 versions.sort_by(|a, b| b.0.cmp(&a.0));
617
618 Ok(versions)
619 }
620
621 #[instrument(name = "store::multi::memory::scan_historical_below", level = "trace", skip(self, cursor), fields(table = ?table, cutoff = cutoff.0, batch_size = batch_size))]
622 fn scan_historical_below(
623 &self,
624 table: EntryKind,
625 cutoff: CommitVersion,
626 cursor: &mut HistoricalCursor,
627 batch_size: usize,
628 ) -> Result<Vec<(EncodedKey, CommitVersion)>> {
629 if cursor.exhausted || batch_size == 0 {
630 return Ok(Vec::new());
631 }
632
633 let entry = match self.inner.entries.data.get(&table) {
634 Some(e) => e,
635 None => {
636 cursor.exhausted = true;
637 return Ok(Vec::new());
638 }
639 };
640
641 let historical = entry.historical.read();
642
643 let mut collected: Vec<(EncodedKey, CommitVersion)> = Vec::new();
644 let mut over_limit = false;
645
646 for (key, versions) in historical.iter() {
647 match (cursor.last_key.as_ref(), cursor.last_version) {
648 (Some(lk), _) if key < lk => continue,
649 (Some(lk), Some(lv)) if key == lk => {
650 for (Reverse(v), _value) in versions.iter().rev() {
651 if *v <= lv {
652 continue;
653 }
654 if *v >= cutoff {
655 continue;
656 }
657 collected.push((key.clone(), *v));
658 if collected.len() > batch_size {
659 over_limit = true;
660 break;
661 }
662 }
663 }
664 _ => {
665 for (Reverse(v), _value) in versions.iter().rev() {
666 if *v >= cutoff {
667 continue;
668 }
669 collected.push((key.clone(), *v));
670 if collected.len() > batch_size {
671 over_limit = true;
672 break;
673 }
674 }
675 }
676 }
677
678 if over_limit {
679 break;
680 }
681 }
682
683 collected.sort_by(|a, b| a.0.as_slice().cmp(b.0.as_slice()).then(a.1.0.cmp(&b.1.0)));
684
685 let has_more = collected.len() > batch_size;
686 if has_more {
687 collected.truncate(batch_size);
688 }
689
690 if let Some(last) = collected.last() {
691 cursor.last_key = Some(last.0.clone());
692 cursor.last_version = Some(last.1);
693 }
694 if !has_more {
695 cursor.exhausted = true;
696 }
697
698 Ok(collected)
699 }
700}
701
702impl TierBackend for MemoryPrimitiveStorage {}
703
704#[cfg(test)]
705pub mod tests {
706 use reifydb_core::interface::catalog::{id::TableId, shape::ShapeId};
707
708 use super::*;
709
710 #[test]
711 fn test_basic_operations() {
712 let storage = MemoryPrimitiveStorage::new();
713
714 let key = EncodedKey::new(b"key1".to_vec());
715 let version = CommitVersion(1);
716
717 storage.set(
719 version,
720 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"value1".to_vec())))])]),
721 )
722 .unwrap();
723
724 let value = storage.get(EntryKind::Multi, &key, version).unwrap().value();
725 assert_eq!(value.as_deref(), Some(b"value1".as_slice()));
726
727 assert!(storage.contains(EntryKind::Multi, &key, version).unwrap());
729
730 assert!(!storage.contains(EntryKind::Multi, b"nonexistent", version).unwrap());
731
732 let version2 = CommitVersion(2);
734 storage.set(version2, HashMap::from([(EntryKind::Multi, vec![(key.clone(), None)])])).unwrap();
735 assert!(!storage.contains(EntryKind::Multi, &key, version2).unwrap());
736 }
737
738 #[test]
739 fn test_source_tables() {
740 let storage = MemoryPrimitiveStorage::new();
741
742 let source1 = ShapeId::Table(TableId(1));
743 let source2 = ShapeId::Table(TableId(2));
744
745 let key = EncodedKey::new(b"key".to_vec());
746 let version = CommitVersion(1);
747
748 storage.set(
749 version,
750 HashMap::from([(
751 EntryKind::Source(source1),
752 vec![(key.clone(), Some(CowVec::new(b"table1".to_vec())))],
753 )]),
754 )
755 .unwrap();
756 storage.set(
757 version,
758 HashMap::from([(
759 EntryKind::Source(source2),
760 vec![(key.clone(), Some(CowVec::new(b"table2".to_vec())))],
761 )]),
762 )
763 .unwrap();
764
765 assert_eq!(
766 storage.get(EntryKind::Source(source1), &key, version).unwrap().value().as_deref(),
767 Some(b"table1".as_slice())
768 );
769 assert_eq!(
770 storage.get(EntryKind::Source(source2), &key, version).unwrap().value().as_deref(),
771 Some(b"table2".as_slice())
772 );
773 }
774
775 #[test]
776 fn test_version_promotion_to_historical() {
777 let storage = MemoryPrimitiveStorage::new();
778
779 let key = EncodedKey::new(b"key1".to_vec());
780
781 storage.set(
783 CommitVersion(1),
784 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"v1".to_vec())))])]),
785 )
786 .unwrap();
787
788 storage.set(
790 CommitVersion(2),
791 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"v2".to_vec())))])]),
792 )
793 .unwrap();
794
795 storage.set(
797 CommitVersion(3),
798 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"v3".to_vec())))])]),
799 )
800 .unwrap();
801
802 assert_eq!(
804 storage.get(EntryKind::Multi, &key, CommitVersion(3)).unwrap().value().as_deref(),
805 Some(b"v3".as_slice())
806 );
807
808 assert_eq!(
810 storage.get(EntryKind::Multi, &key, CommitVersion(2)).unwrap().value().as_deref(),
811 Some(b"v2".as_slice())
812 );
813
814 assert_eq!(
816 storage.get(EntryKind::Multi, &key, CommitVersion(1)).unwrap().value().as_deref(),
817 Some(b"v1".as_slice())
818 );
819 }
820
821 #[test]
822 fn test_insert_older_version() {
823 let storage = MemoryPrimitiveStorage::new();
824
825 let key = EncodedKey::new(b"key1".to_vec());
826
827 storage.set(
829 CommitVersion(3),
830 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"v3".to_vec())))])]),
831 )
832 .unwrap();
833
834 storage.set(
836 CommitVersion(1),
837 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"v1".to_vec())))])]),
838 )
839 .unwrap();
840
841 assert_eq!(
843 storage.get(EntryKind::Multi, &key, CommitVersion(3)).unwrap().value().as_deref(),
844 Some(b"v3".as_slice())
845 );
846
847 assert_eq!(
849 storage.get(EntryKind::Multi, &key, CommitVersion(1)).unwrap().value().as_deref(),
850 Some(b"v1".as_slice())
851 );
852
853 assert_eq!(
855 storage.get(EntryKind::Multi, &key, CommitVersion(2)).unwrap().value().as_deref(),
856 Some(b"v1".as_slice())
857 );
858 }
859
860 #[test]
861 fn test_range_next() {
862 let storage = MemoryPrimitiveStorage::new();
863
864 let version = CommitVersion(1);
865 storage.set(
866 version,
867 HashMap::from([(
868 EntryKind::Multi,
869 vec![
870 (EncodedKey::new(b"a".to_vec()), Some(CowVec::new(b"1".to_vec()))),
871 (EncodedKey::new(b"b".to_vec()), Some(CowVec::new(b"2".to_vec()))),
872 (EncodedKey::new(b"c".to_vec()), Some(CowVec::new(b"3".to_vec()))),
873 ],
874 )]),
875 )
876 .unwrap();
877
878 let mut cursor = RangeCursor::new();
879 let batch = storage
880 .range_next(
881 EntryKind::Multi,
882 &mut cursor,
883 Bound::Unbounded,
884 Bound::Unbounded,
885 MultiVersionScope::AsOf {
886 read: version,
887 },
888 100,
889 )
890 .unwrap();
891
892 assert_eq!(batch.entries.len(), 3);
893 assert!(!batch.has_more);
894 assert!(cursor.exhausted);
895
896 assert_eq!(&*batch.entries[0].key, b"a");
898 assert_eq!(&*batch.entries[1].key, b"b");
899 assert_eq!(&*batch.entries[2].key, b"c");
900 }
901
902 #[test]
903 fn test_range_rev_next() {
904 let storage = MemoryPrimitiveStorage::new();
905
906 let version = CommitVersion(1);
907 storage.set(
908 version,
909 HashMap::from([(
910 EntryKind::Multi,
911 vec![
912 (EncodedKey::new(b"a".to_vec()), Some(CowVec::new(b"1".to_vec()))),
913 (EncodedKey::new(b"b".to_vec()), Some(CowVec::new(b"2".to_vec()))),
914 (EncodedKey::new(b"c".to_vec()), Some(CowVec::new(b"3".to_vec()))),
915 ],
916 )]),
917 )
918 .unwrap();
919
920 let mut cursor = RangeCursor::new();
921 let batch = storage
922 .range_rev_next(
923 EntryKind::Multi,
924 &mut cursor,
925 Bound::Unbounded,
926 Bound::Unbounded,
927 MultiVersionScope::AsOf {
928 read: version,
929 },
930 100,
931 )
932 .unwrap();
933
934 assert_eq!(batch.entries.len(), 3);
935 assert!(!batch.has_more);
936 assert!(cursor.exhausted);
937
938 assert_eq!(&*batch.entries[0].key, b"c");
940 assert_eq!(&*batch.entries[1].key, b"b");
941 assert_eq!(&*batch.entries[2].key, b"a");
942 }
943
944 #[test]
945 fn test_range_streaming_pagination() {
946 let storage = MemoryPrimitiveStorage::new();
947
948 let version = CommitVersion(1);
949
950 let entries: Vec<_> =
952 (0..10u8).map(|i| (EncodedKey::new(vec![i]), Some(CowVec::new(vec![i * 10])))).collect();
953 storage.set(version, HashMap::from([(EntryKind::Multi, entries)])).unwrap();
954
955 let mut cursor = RangeCursor::new();
957
958 let batch1 = storage
960 .range_next(
961 EntryKind::Multi,
962 &mut cursor,
963 Bound::Unbounded,
964 Bound::Unbounded,
965 MultiVersionScope::AsOf {
966 read: version,
967 },
968 3,
969 )
970 .unwrap();
971 assert_eq!(batch1.entries.len(), 3);
972 assert!(batch1.has_more);
973 assert!(!cursor.exhausted);
974
975 assert_eq!(&*batch1.entries[0].key, &[0]);
976 assert_eq!(&*batch1.entries[2].key, &[2]);
977
978 let batch2 = storage
980 .range_next(
981 EntryKind::Multi,
982 &mut cursor,
983 Bound::Unbounded,
984 Bound::Unbounded,
985 MultiVersionScope::AsOf {
986 read: version,
987 },
988 3,
989 )
990 .unwrap();
991 assert_eq!(batch2.entries.len(), 3);
992 assert!(batch2.has_more);
993 assert!(!cursor.exhausted);
994
995 assert_eq!(&*batch2.entries[0].key, &[3]);
996 assert_eq!(&*batch2.entries[2].key, &[5]);
997
998 let batch3 = storage
1000 .range_next(
1001 EntryKind::Multi,
1002 &mut cursor,
1003 Bound::Unbounded,
1004 Bound::Unbounded,
1005 MultiVersionScope::AsOf {
1006 read: version,
1007 },
1008 3,
1009 )
1010 .unwrap();
1011 assert_eq!(batch3.entries.len(), 3);
1012 assert!(batch3.has_more);
1013 assert!(!cursor.exhausted);
1014
1015 assert_eq!(&*batch3.entries[0].key, &[6]);
1016 assert_eq!(&*batch3.entries[2].key, &[8]);
1017
1018 let batch4 = storage
1020 .range_next(
1021 EntryKind::Multi,
1022 &mut cursor,
1023 Bound::Unbounded,
1024 Bound::Unbounded,
1025 MultiVersionScope::AsOf {
1026 read: version,
1027 },
1028 3,
1029 )
1030 .unwrap();
1031 assert_eq!(batch4.entries.len(), 1);
1032 assert!(!batch4.has_more);
1033 assert!(cursor.exhausted);
1034
1035 assert_eq!(&*batch4.entries[0].key, &[9]);
1036
1037 let batch5 = storage
1039 .range_next(
1040 EntryKind::Multi,
1041 &mut cursor,
1042 Bound::Unbounded,
1043 Bound::Unbounded,
1044 MultiVersionScope::AsOf {
1045 read: version,
1046 },
1047 3,
1048 )
1049 .unwrap();
1050 assert!(batch5.entries.is_empty());
1051 }
1052
1053 #[test]
1054 fn test_range_reving_pagination() {
1055 let storage = MemoryPrimitiveStorage::new();
1056
1057 let version = CommitVersion(1);
1058
1059 let entries: Vec<_> =
1061 (0..10u8).map(|i| (EncodedKey::new(vec![i]), Some(CowVec::new(vec![i * 10])))).collect();
1062 storage.set(version, HashMap::from([(EntryKind::Multi, entries)])).unwrap();
1063
1064 let mut cursor = RangeCursor::new();
1066
1067 let batch1 = storage
1069 .range_rev_next(
1070 EntryKind::Multi,
1071 &mut cursor,
1072 Bound::Unbounded,
1073 Bound::Unbounded,
1074 MultiVersionScope::AsOf {
1075 read: version,
1076 },
1077 3,
1078 )
1079 .unwrap();
1080 assert_eq!(batch1.entries.len(), 3);
1081 assert!(batch1.has_more);
1082 assert!(!cursor.exhausted);
1083
1084 assert_eq!(&*batch1.entries[0].key, &[9]);
1085 assert_eq!(&*batch1.entries[2].key, &[7]);
1086
1087 let batch2 = storage
1089 .range_rev_next(
1090 EntryKind::Multi,
1091 &mut cursor,
1092 Bound::Unbounded,
1093 Bound::Unbounded,
1094 MultiVersionScope::AsOf {
1095 read: version,
1096 },
1097 3,
1098 )
1099 .unwrap();
1100 assert_eq!(batch2.entries.len(), 3);
1101 assert!(batch2.has_more);
1102 assert!(!cursor.exhausted);
1103
1104 assert_eq!(&*batch2.entries[0].key, &[6]);
1105 assert_eq!(&*batch2.entries[2].key, &[4]);
1106 }
1107
1108 #[test]
1109 fn test_drop_from_historical() {
1110 let storage = MemoryPrimitiveStorage::new();
1111
1112 let key = EncodedKey::new(b"key1".to_vec());
1113
1114 for v in 1..=3u64 {
1116 storage.set(
1117 CommitVersion(v),
1118 HashMap::from([(
1119 EntryKind::Multi,
1120 vec![(key.clone(), Some(CowVec::new(format!("v{}", v).into_bytes())))],
1121 )]),
1122 )
1123 .unwrap();
1124 }
1125
1126 storage.drop(HashMap::from([(EntryKind::Multi, vec![(key.clone(), CommitVersion(1))])])).unwrap();
1129
1130 assert!(storage.get(EntryKind::Multi, &key, CommitVersion(1)).unwrap().value().is_none());
1132
1133 assert_eq!(
1135 storage.get(EntryKind::Multi, &key, CommitVersion(2)).unwrap().value().as_deref(),
1136 Some(b"v2".as_slice())
1137 );
1138 assert_eq!(
1139 storage.get(EntryKind::Multi, &key, CommitVersion(3)).unwrap().value().as_deref(),
1140 Some(b"v3".as_slice())
1141 );
1142 }
1143
1144 #[test]
1145 fn test_tombstones() {
1146 let storage = MemoryPrimitiveStorage::new();
1147
1148 let key = EncodedKey::new(b"key1".to_vec());
1149
1150 storage.set(
1152 CommitVersion(1),
1153 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"value".to_vec())))])]),
1154 )
1155 .unwrap();
1156
1157 storage.set(CommitVersion(2), HashMap::from([(EntryKind::Multi, vec![(key.clone(), None)])])).unwrap();
1159
1160 assert!(storage.get(EntryKind::Multi, &key, CommitVersion(2)).unwrap().value().is_none());
1162 assert!(!storage.contains(EntryKind::Multi, &key, CommitVersion(2)).unwrap());
1163
1164 assert_eq!(
1166 storage.get(EntryKind::Multi, &key, CommitVersion(1)).unwrap().value().as_deref(),
1167 Some(b"value".as_slice())
1168 );
1169 }
1170
1171 #[test]
1172 fn test_collect_evictable_below_keeps_versions_above_cutoff() {
1173 let storage = MemoryPrimitiveStorage::new();
1174 let key = EncodedKey::new(b"k".to_vec());
1175 for v in 1..=3u64 {
1176 storage.set(
1177 CommitVersion(v),
1178 HashMap::from([(
1179 EntryKind::Multi,
1180 vec![(key.clone(), Some(CowVec::new(format!("v{v}").into_bytes())))],
1181 )]),
1182 )
1183 .unwrap();
1184 }
1185
1186 let (to_persist, to_drop) = storage.collect_evictable_below(EntryKind::Multi, CommitVersion(2));
1189 assert_eq!(to_persist.len(), 1);
1190 assert_eq!(to_persist[0].0, key);
1191 assert_eq!(to_persist[0].1, CommitVersion(2));
1192 assert_eq!(to_persist[0].2.as_deref(), Some(b"v2".as_slice()));
1193 let dropped: HashSet<CommitVersion> = to_drop.iter().map(|(_, v)| *v).collect();
1194 assert_eq!(dropped, HashSet::from([CommitVersion(1), CommitVersion(2)]));
1195
1196 storage.drop(HashMap::from([(EntryKind::Multi, to_drop)])).unwrap();
1199 assert_eq!(
1200 storage.get(EntryKind::Multi, &key, CommitVersion(3)).unwrap().value().as_deref(),
1201 Some(b"v3".as_slice())
1202 );
1203 assert!(storage.get(EntryKind::Multi, &key, CommitVersion(2)).unwrap().value().is_none());
1204 assert!(storage.get(EntryKind::Multi, &key, CommitVersion(1)).unwrap().value().is_none());
1205 }
1206
1207 #[test]
1208 fn test_collect_evictable_below_empty_when_all_above_cutoff() {
1209 let storage = MemoryPrimitiveStorage::new();
1210 let key = EncodedKey::new(b"k".to_vec());
1211 storage.set(
1212 CommitVersion(5),
1213 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"v".to_vec())))])]),
1214 )
1215 .unwrap();
1216 let (to_persist, to_drop) = storage.collect_evictable_below(EntryKind::Multi, CommitVersion(3));
1217 assert!(to_persist.is_empty());
1218 assert!(to_drop.is_empty());
1219 }
1220
1221 #[test]
1222 fn test_collect_evictable_below_persists_exactly_one_value_per_key() {
1223 let storage = MemoryPrimitiveStorage::new();
1228 let key = EncodedKey::new(b"k".to_vec());
1229 for v in 1..=5u64 {
1230 storage.set(
1231 CommitVersion(v),
1232 HashMap::from([(
1233 EntryKind::Multi,
1234 vec![(key.clone(), Some(CowVec::new(format!("v{v}").into_bytes())))],
1235 )]),
1236 )
1237 .unwrap();
1238 }
1239
1240 let (to_persist, to_drop) = storage.collect_evictable_below(EntryKind::Multi, CommitVersion(4));
1242 assert_eq!(to_persist.len(), 1, "exactly one value persisted per key");
1243 assert_eq!(to_persist[0].1, CommitVersion(4), "the latest version <= cutoff");
1244 assert_eq!(to_persist[0].2.as_deref(), Some(b"v4".as_slice()));
1245
1246 let dropped: HashSet<CommitVersion> = to_drop.iter().map(|(_, v)| *v).collect();
1248 assert_eq!(
1249 dropped,
1250 HashSet::from([CommitVersion(1), CommitVersion(2), CommitVersion(3), CommitVersion(4)])
1251 );
1252 }
1253
1254 #[test]
1255 fn test_collect_evictable_below_persists_tombstone_when_it_is_the_latest() {
1256 let storage = MemoryPrimitiveStorage::new();
1260 let key = EncodedKey::new(b"k".to_vec());
1261 storage.set(
1262 CommitVersion(1),
1263 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"v1".to_vec())))])]),
1264 )
1265 .unwrap();
1266 storage.set(CommitVersion(2), HashMap::from([(EntryKind::Multi, vec![(key.clone(), None)])])).unwrap();
1267
1268 let (to_persist, to_drop) = storage.collect_evictable_below(EntryKind::Multi, CommitVersion(2));
1269 assert_eq!(to_persist.len(), 1);
1270 assert_eq!(to_persist[0].1, CommitVersion(2), "the tombstone is the latest version");
1271 assert!(to_persist[0].2.is_none(), "the persisted latest value must be the tombstone, not v1");
1272 assert_eq!(to_drop.len(), 2, "both v1 and the tombstone are dropped from the buffer");
1273 }
1274
1275 #[test]
1276 fn test_collect_evictable_below_only_drops_historical_when_current_is_above_cutoff() {
1277 let storage = MemoryPrimitiveStorage::new();
1281 let key = EncodedKey::new(b"k".to_vec());
1282 storage.set(
1283 CommitVersion(2),
1284 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"v2".to_vec())))])]),
1285 )
1286 .unwrap();
1287 storage.set(
1288 CommitVersion(5),
1289 HashMap::from([(EntryKind::Multi, vec![(key.clone(), Some(CowVec::new(b"v5".to_vec())))])]),
1290 )
1291 .unwrap();
1292
1293 let (to_persist, to_drop) = storage.collect_evictable_below(EntryKind::Multi, CommitVersion(3));
1294 assert_eq!(to_persist.len(), 1);
1295 assert_eq!(to_persist[0].1, CommitVersion(2), "only the aged-out historical version is persisted");
1296 assert_eq!(to_persist[0].2.as_deref(), Some(b"v2".as_slice()));
1297 let dropped: HashSet<CommitVersion> = to_drop.iter().map(|(_, v)| *v).collect();
1298 assert_eq!(dropped, HashSet::from([CommitVersion(2)]), "v5 (current, > cutoff) is never dropped");
1299
1300 storage.drop(HashMap::from([(EntryKind::Multi, to_drop)])).unwrap();
1302 assert_eq!(
1303 storage.get(EntryKind::Multi, &key, CommitVersion(5)).unwrap().value().as_deref(),
1304 Some(b"v5".as_slice())
1305 );
1306 assert!(
1307 storage.get(EntryKind::Multi, &key, CommitVersion(3)).unwrap().value().is_none(),
1308 "the v2 a reader at snapshot 3 used to see is gone from the buffer after eviction"
1309 );
1310 }
1311
1312 #[test]
1313 fn test_collect_evictable_below_handles_multiple_keys_independently() {
1314 let storage = MemoryPrimitiveStorage::new();
1318 let cold = EncodedKey::new(b"cold".to_vec());
1319 let hot = EncodedKey::new(b"hot".to_vec());
1320 storage.set(
1321 CommitVersion(1),
1322 HashMap::from([(EntryKind::Multi, vec![(cold.clone(), Some(CowVec::new(b"cold1".to_vec())))])]),
1323 )
1324 .unwrap();
1325 storage.set(
1326 CommitVersion(9),
1327 HashMap::from([(EntryKind::Multi, vec![(hot.clone(), Some(CowVec::new(b"hot9".to_vec())))])]),
1328 )
1329 .unwrap();
1330
1331 let (to_persist, to_drop) = storage.collect_evictable_below(EntryKind::Multi, CommitVersion(5));
1332 assert_eq!(to_persist.len(), 1, "only the cold key is evictable below the cutoff");
1333 assert_eq!(to_persist[0].0, cold);
1334 assert!(to_drop.iter().all(|(k, _)| *k == cold), "the hot key must not be scheduled for drop");
1335 }
1336}