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