icydb_core/db/data/
store.rs

1//! Module: data::store
2//! Responsibility: journaled-or-heap row storage behind the data-store boundary.
3//! Does not own: key/row validation policy beyond type boundaries.
4//! Boundary: commit/executor call into this layer after prevalidation.
5
6use crate::{
7    db::{
8        data::{CanonicalRow, RawDataStoreKey, RawRow},
9        direction::Direction,
10        ordered_overlay::{OrderedOverlayEntry, OrderedOverlayVisit, visit_ordered_overlay},
11    },
12    types::EntityTag,
13};
14use ic_stable_structures::{
15    BTreeMap as StableBTreeMap, DefaultMemoryImpl, memory_manager::VirtualMemory,
16};
17#[cfg(feature = "diagnostics")]
18use std::cell::Cell;
19use std::collections::{BTreeMap as HeapBTreeMap, BTreeSet};
20use std::convert::Infallible;
21use std::ops::{Bound, RangeBounds};
22
23#[cfg(feature = "diagnostics")]
24thread_local! {
25    static DATA_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
26}
27
28#[cfg(feature = "diagnostics")]
29fn record_data_store_get_call() {
30    DATA_STORE_GET_CALL_COUNT.with(|count| {
31        count.set(count.get().saturating_add(1));
32    });
33}
34
35///
36/// DataStore
37///
38/// Thin persistence wrapper over one journaled or heap BTreeMap.
39///
40/// Invariant: callers provide already-validated `RawDataStoreKey` and canonical row bytes.
41/// This type intentionally does not enforce commit-phase ordering.
42///
43
44pub struct DataStore {
45    backend: DataStoreBackend,
46    generation: u64,
47    entity_cardinality: EntityCardinality,
48}
49
50enum DataStoreBackend {
51    Heap(HeapBTreeMap<RawDataStoreKey, RawRow>),
52    Journaled {
53        canonical: StableBTreeMap<RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>>,
54        live: HeapBTreeMap<RawDataStoreKey, RawRow>,
55        tombstones: BTreeSet<RawDataStoreKey>,
56    },
57}
58
59/// Control-flow result for store traversal visitors.
60#[derive(Clone, Copy, Debug, Eq, PartialEq)]
61pub(in crate::db) enum StoreVisit {
62    Continue,
63    Stop,
64}
65
66impl StoreVisit {
67    const fn should_stop(self) -> bool {
68        matches!(self, Self::Stop)
69    }
70}
71
72impl DataStore {
73    /// Initialize a volatile heap-backed data store.
74    #[must_use]
75    pub const fn init_heap() -> Self {
76        Self {
77            backend: DataStoreBackend::Heap(HeapBTreeMap::new()),
78            generation: 0,
79            entity_cardinality: EntityCardinality::empty(),
80        }
81    }
82
83    /// Initialize a journaled cached-stable data store.
84    ///
85    /// Normal writes update only the live projection. The canonical stable map
86    /// is the future fold target and is not mutated by this wrapper's write
87    /// methods.
88    #[must_use]
89    pub fn init_journaled(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
90        let mut store = Self {
91            backend: DataStoreBackend::Journaled {
92                canonical: StableBTreeMap::init(memory),
93                live: HeapBTreeMap::new(),
94                tombstones: BTreeSet::new(),
95            },
96            generation: 0,
97            entity_cardinality: EntityCardinality::empty(),
98        };
99        store.rebuild_entity_cardinality_from_entries();
100        store
101    }
102
103    /// Insert or replace one row by raw key.
104    pub(in crate::db) fn insert(
105        &mut self,
106        key: RawDataStoreKey,
107        row: CanonicalRow,
108    ) -> Option<RawRow> {
109        let row = row.into_raw_row();
110        let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
111            self.get(&key)
112        } else {
113            None
114        };
115        let cardinality_key = key.clone();
116        let previous = match &mut self.backend {
117            DataStoreBackend::Heap(map) => map.insert(key, row),
118            DataStoreBackend::Journaled {
119                live, tombstones, ..
120            } => {
121                tombstones.remove(&key);
122                live.insert(key, row);
123                previous_journaled
124            }
125        };
126        self.entity_cardinality
127            .apply_insert(&cardinality_key, previous.as_ref());
128        self.bump_generation();
129        previous
130    }
131
132    /// Insert one raw row directly for corruption-focused test setup only.
133    #[cfg(test)]
134    pub(in crate::db) fn insert_raw_for_test(
135        &mut self,
136        key: RawDataStoreKey,
137        row: RawRow,
138    ) -> Option<RawRow> {
139        let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
140            self.get(&key)
141        } else {
142            None
143        };
144        let cardinality_key = key.clone();
145        let previous = match &mut self.backend {
146            DataStoreBackend::Heap(map) => map.insert(key, row),
147            DataStoreBackend::Journaled {
148                live, tombstones, ..
149            } => {
150                tombstones.remove(&key);
151                live.insert(key, row);
152                previous_journaled
153            }
154        };
155        self.entity_cardinality
156            .apply_insert(&cardinality_key, previous.as_ref());
157        self.bump_generation();
158        previous
159    }
160
161    /// Remove one row by raw key.
162    pub(in crate::db) fn remove(&mut self, key: &RawDataStoreKey) -> Option<RawRow> {
163        let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
164            self.get(key)
165        } else {
166            None
167        };
168        let previous = match &mut self.backend {
169            DataStoreBackend::Heap(map) => map.remove(key),
170            DataStoreBackend::Journaled {
171                live, tombstones, ..
172            } => {
173                live.remove(key);
174                tombstones.insert(key.clone());
175                previous_journaled
176            }
177        };
178        self.entity_cardinality.apply_remove(key, previous.as_ref());
179        self.bump_generation();
180        previous
181    }
182
183    /// Reset the volatile projection for journaled recovery without mutating
184    /// the canonical stable base.
185    pub(in crate::db) fn reset_journaled_live_projection(
186        &mut self,
187    ) -> Result<(), crate::error::InternalError> {
188        let DataStoreBackend::Journaled {
189            live, tombstones, ..
190        } = &mut self.backend
191        else {
192            return Err(crate::error::InternalError::store_invariant());
193        };
194
195        live.clear();
196        tombstones.clear();
197        self.rebuild_entity_cardinality_from_entries();
198        self.bump_generation();
199
200        Ok(())
201    }
202
203    /// Apply one recovered journal row put into the volatile projection.
204    pub(in crate::db) fn apply_recovered_journal_put(
205        &mut self,
206        key: RawDataStoreKey,
207        row: RawRow,
208    ) -> Result<Option<RawRow>, crate::error::InternalError> {
209        let DataStoreBackend::Journaled {
210            canonical,
211            live,
212            tombstones,
213        } = &mut self.backend
214        else {
215            return Err(crate::error::InternalError::store_invariant());
216        };
217
218        let previous = if tombstones.contains(&key) {
219            None
220        } else {
221            live.get(&key).cloned().or_else(|| canonical.get(&key))
222        };
223        tombstones.remove(&key);
224        let cardinality_key = key.clone();
225        live.insert(key, row);
226        self.entity_cardinality
227            .apply_insert(&cardinality_key, previous.as_ref());
228        self.bump_generation();
229
230        Ok(previous)
231    }
232
233    /// Apply one recovered journal row delete into the volatile projection.
234    pub(in crate::db) fn apply_recovered_journal_delete(
235        &mut self,
236        key: &RawDataStoreKey,
237    ) -> Result<Option<RawRow>, crate::error::InternalError> {
238        let DataStoreBackend::Journaled {
239            canonical,
240            live,
241            tombstones,
242        } = &mut self.backend
243        else {
244            return Err(crate::error::InternalError::store_invariant());
245        };
246
247        let previous = if tombstones.contains(key) {
248            None
249        } else {
250            live.get(key).cloned().or_else(|| canonical.get(key))
251        };
252        live.remove(key);
253        tombstones.insert(key.clone());
254        self.entity_cardinality.apply_remove(key, previous.as_ref());
255        self.bump_generation();
256
257        Ok(previous)
258    }
259
260    /// Apply one folded journal row put into the canonical stable base.
261    pub(in crate::db) fn fold_recovered_journal_put(
262        &mut self,
263        key: RawDataStoreKey,
264        row: RawRow,
265    ) -> Result<Option<RawRow>, crate::error::InternalError> {
266        let DataStoreBackend::Journaled {
267            canonical,
268            live,
269            tombstones,
270        } = &mut self.backend
271        else {
272            return Err(crate::error::InternalError::store_invariant());
273        };
274
275        let visible = !live.contains_key(&key) && !tombstones.contains(&key);
276        let cardinality_key = key.clone();
277        let previous = canonical.insert(key, row);
278        if visible {
279            self.entity_cardinality
280                .apply_insert(&cardinality_key, previous.as_ref());
281        }
282        self.bump_generation();
283
284        Ok(previous)
285    }
286
287    /// Apply one folded journal row delete into the canonical stable base.
288    pub(in crate::db) fn fold_recovered_journal_delete(
289        &mut self,
290        key: &RawDataStoreKey,
291    ) -> Result<Option<RawRow>, crate::error::InternalError> {
292        let DataStoreBackend::Journaled {
293            canonical,
294            live,
295            tombstones,
296        } = &mut self.backend
297        else {
298            return Err(crate::error::InternalError::store_invariant());
299        };
300
301        let visible = !live.contains_key(key) && !tombstones.contains(key);
302        let previous = canonical.remove(key);
303        if visible {
304            self.entity_cardinality.apply_remove(key, previous.as_ref());
305        }
306        self.bump_generation();
307
308        Ok(previous)
309    }
310
311    /// Load one row by raw key.
312    pub(in crate::db) fn get(&self, key: &RawDataStoreKey) -> Option<RawRow> {
313        #[cfg(feature = "diagnostics")]
314        record_data_store_get_call();
315
316        match &self.backend {
317            DataStoreBackend::Heap(map) => map.get(key).cloned(),
318            DataStoreBackend::Journaled { .. } => Self::journaled_get_raw(&self.backend, key),
319        }
320    }
321
322    /// Return whether one raw key exists without cloning the row payload.
323    #[must_use]
324    pub(in crate::db) fn contains(&self, key: &RawDataStoreKey) -> bool {
325        match &self.backend {
326            DataStoreBackend::Heap(map) => map.contains_key(key),
327            DataStoreBackend::Journaled { .. } => {
328                Self::journaled_get_raw(&self.backend, key).is_some()
329            }
330        }
331    }
332
333    /// Return the number of stored rows without exposing the backing map.
334    #[must_use]
335    pub(in crate::db) fn len(&self) -> u64 {
336        match &self.backend {
337            DataStoreBackend::Heap(map) => u64::try_from(map.len()).unwrap_or(u64::MAX),
338            DataStoreBackend::Journaled { .. } => {
339                let mut count = 0_u64;
340                let _: Result<(), Infallible> = self.visit_entries(|_key, _row| {
341                    count = count.saturating_add(1);
342                    Ok(StoreVisit::Continue)
343                });
344                count
345            }
346        }
347    }
348
349    /// Return the row-store generation used to prove index metadata freshness.
350    #[must_use]
351    pub(in crate::db) const fn generation(&self) -> u64 {
352        self.generation
353    }
354
355    /// Return an exact current row count for one entity when store metadata is valid.
356    #[must_use]
357    pub(in crate::db) fn exact_entity_count(&self, entity: EntityTag) -> Option<u64> {
358        self.entity_cardinality.exact_count(entity)
359    }
360
361    /// Visit raw row entries in canonical storage order.
362    pub(in crate::db) fn visit_entries<E>(
363        &self,
364        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
365    ) -> Result<(), E> {
366        match &self.backend {
367            DataStoreBackend::Heap(map) => {
368                for (key, row) in map {
369                    if visitor(key, row)?.should_stop() {
370                        break;
371                    }
372                }
373            }
374            DataStoreBackend::Journaled {
375                canonical: _,
376                live: _,
377                tombstones: _,
378            } => Self::visit_journaled_entries_in_bounds(
379                &self.backend,
380                (Bound::Unbounded, Bound::Unbounded),
381                false,
382                visitor,
383            )?,
384        }
385
386        Ok(())
387    }
388
389    /// Visit raw row entries whose keys belong to the provided storage range.
390    pub(in crate::db) fn visit_range<E>(
391        &self,
392        key_range: impl RangeBounds<RawDataStoreKey>,
393        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
394    ) -> Result<(), E> {
395        let bounds = Self::owned_range_bounds(&key_range);
396        match &self.backend {
397            DataStoreBackend::Heap(map) => {
398                for (key, row) in map.range((bounds.0.clone(), bounds.1)) {
399                    if visitor(key, row)?.should_stop() {
400                        break;
401                    }
402                }
403            }
404            DataStoreBackend::Journaled {
405                canonical: _,
406                live: _,
407                tombstones: _,
408            } => Self::visit_journaled_entries_in_bounds(&self.backend, bounds, false, visitor)?,
409        }
410
411        Ok(())
412    }
413
414    /// Visit raw row entries in reverse order whose keys belong to the provided storage range.
415    #[cfg(any(test, feature = "query"))]
416    pub(in crate::db) fn visit_range_rev<E>(
417        &self,
418        key_range: impl RangeBounds<RawDataStoreKey>,
419        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
420    ) -> Result<(), E> {
421        let bounds = Self::owned_range_bounds(&key_range);
422        match &self.backend {
423            DataStoreBackend::Heap(map) => {
424                for (key, row) in map.range((bounds.0.clone(), bounds.1)).rev() {
425                    if visitor(key, row)?.should_stop() {
426                        break;
427                    }
428                }
429            }
430            DataStoreBackend::Journaled {
431                canonical: _,
432                live: _,
433                tombstones: _,
434            } => Self::visit_journaled_entries_in_bounds(&self.backend, bounds, true, visitor)?,
435        }
436
437        Ok(())
438    }
439
440    /// Sum of bytes used by all stored rows.
441    pub(in crate::db) fn memory_bytes(&self) -> u64 {
442        // Report map footprint as key bytes + row bytes per entry.
443        let mut bytes = 0u64;
444        let _: Result<(), Infallible> = self.visit_entries(|key, row| {
445            bytes = bytes.saturating_add(key.as_bytes().len() as u64 + row.len() as u64);
446            Ok(StoreVisit::Continue)
447        });
448        bytes
449    }
450
451    const fn bump_generation(&mut self) {
452        self.generation = self.generation.saturating_add(1);
453    }
454
455    fn rebuild_entity_cardinality_from_entries(&mut self) {
456        let mut cardinality = EntityCardinality::empty();
457        let _: Result<(), Infallible> = self.visit_entries(|key, _row| {
458            cardinality.apply_present_key(key);
459            Ok(StoreVisit::Continue)
460        });
461        self.entity_cardinality = cardinality;
462    }
463
464    /// Return the monotonic perf-only count of stable row fetches seen by this process.
465    #[cfg(all(feature = "diagnostics", any(test, feature = "query")))]
466    pub(in crate::db) fn current_get_call_count() -> u64 {
467        DATA_STORE_GET_CALL_COUNT.with(Cell::get)
468    }
469
470    fn journaled_get_raw(backend: &DataStoreBackend, key: &RawDataStoreKey) -> Option<RawRow> {
471        let DataStoreBackend::Journaled {
472            canonical,
473            live,
474            tombstones,
475        } = backend
476        else {
477            return None;
478        };
479
480        if tombstones.contains(key) {
481            return None;
482        }
483        live.get(key).cloned().or_else(|| canonical.get(key))
484    }
485
486    fn owned_range_bounds(
487        key_range: &impl RangeBounds<RawDataStoreKey>,
488    ) -> (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>) {
489        let lower = match key_range.start_bound() {
490            Bound::Included(key) => Bound::Included(key.clone()),
491            Bound::Excluded(key) => Bound::Excluded(key.clone()),
492            Bound::Unbounded => Bound::Unbounded,
493        };
494        let upper = match key_range.end_bound() {
495            Bound::Included(key) => Bound::Included(key.clone()),
496            Bound::Excluded(key) => Bound::Excluded(key.clone()),
497            Bound::Unbounded => Bound::Unbounded,
498        };
499
500        (lower, upper)
501    }
502
503    fn visit_journaled_entries_in_bounds<E>(
504        backend: &DataStoreBackend,
505        bounds: (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>),
506        reverse: bool,
507        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
508    ) -> Result<(), E> {
509        let DataStoreBackend::Journaled {
510            canonical,
511            live,
512            tombstones,
513        } = backend
514        else {
515            return Ok(());
516        };
517
518        if canonical.is_empty() {
519            if reverse {
520                for (key, row) in live.range(bounds).rev() {
521                    if visitor(key, row)?.should_stop() {
522                        return Ok(());
523                    }
524                }
525            } else {
526                for (key, row) in live.range(bounds) {
527                    if visitor(key, row)?.should_stop() {
528                        return Ok(());
529                    }
530                }
531            }
532            return Ok(());
533        }
534
535        if live.is_empty() && tombstones.is_empty() {
536            if reverse {
537                for entry in canonical.range(bounds).rev() {
538                    if visitor(entry.key(), &entry.value())?.should_stop() {
539                        return Ok(());
540                    }
541                }
542            } else {
543                for entry in canonical.range(bounds) {
544                    if visitor(entry.key(), &entry.value())?.should_stop() {
545                        return Ok(());
546                    }
547                }
548            }
549            return Ok(());
550        }
551
552        match if reverse {
553            Direction::Desc
554        } else {
555            Direction::Asc
556        } {
557            Direction::Asc => visit_ordered_overlay(
558                canonical.range((bounds.0.clone(), bounds.1.clone())),
559                live.range((bounds.0, bounds.1)),
560                Direction::Asc,
561                |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
562                |canonical_entry| !tombstones.contains(canonical_entry.key()),
563                |live_entry| !tombstones.contains(live_entry.0),
564                |entry| {
565                    let visit = match entry {
566                        OrderedOverlayEntry::Canonical(canonical_entry) => {
567                            visitor(canonical_entry.key(), &canonical_entry.value())?
568                        }
569                        OrderedOverlayEntry::Live((key, row)) => visitor(key, row)?,
570                    };
571                    Ok(if visit.should_stop() {
572                        OrderedOverlayVisit::Stop
573                    } else {
574                        OrderedOverlayVisit::Continue
575                    })
576                },
577            ),
578            Direction::Desc => visit_ordered_overlay(
579                canonical.range((bounds.0.clone(), bounds.1.clone())).rev(),
580                live.range((bounds.0, bounds.1)).rev(),
581                Direction::Desc,
582                |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
583                |canonical_entry| !tombstones.contains(canonical_entry.key()),
584                |live_entry| !tombstones.contains(live_entry.0),
585                |entry| {
586                    let visit = match entry {
587                        OrderedOverlayEntry::Canonical(canonical_entry) => {
588                            visitor(canonical_entry.key(), &canonical_entry.value())?
589                        }
590                        OrderedOverlayEntry::Live((key, row)) => visitor(key, row)?,
591                    };
592                    Ok(if visit.should_stop() {
593                        OrderedOverlayVisit::Stop
594                    } else {
595                        OrderedOverlayVisit::Continue
596                    })
597                },
598            ),
599        }
600    }
601}
602
603#[derive(Clone, Debug)]
604struct EntityCardinality {
605    counts: HeapBTreeMap<EntityTag, u64>,
606    decodable: bool,
607}
608
609impl EntityCardinality {
610    const fn empty() -> Self {
611        Self {
612            counts: HeapBTreeMap::new(),
613            decodable: true,
614        }
615    }
616
617    fn exact_count(&self, entity: EntityTag) -> Option<u64> {
618        self.decodable
619            .then(|| self.counts.get(&entity).copied().unwrap_or(0))
620    }
621
622    fn apply_insert(&mut self, key: &RawDataStoreKey, previous: Option<&RawRow>) {
623        if previous.is_some() {
624            return;
625        }
626        self.apply_present_key(key);
627    }
628
629    fn apply_remove(&mut self, key: &RawDataStoreKey, previous: Option<&RawRow>) {
630        if previous.is_none() {
631            return;
632        }
633        self.apply_removed_key(key);
634    }
635
636    fn apply_present_key(&mut self, key: &RawDataStoreKey) {
637        if !self.decodable {
638            return;
639        }
640        let Some(entity) = key.entity_tag_prefix() else {
641            self.invalidate();
642            return;
643        };
644
645        let count = self.counts.entry(entity).or_insert(0);
646        *count = count.saturating_add(1);
647    }
648
649    fn apply_removed_key(&mut self, key: &RawDataStoreKey) {
650        if !self.decodable {
651            return;
652        }
653        let Some(entity) = key.entity_tag_prefix() else {
654            self.invalidate();
655            return;
656        };
657
658        if let Some(count) = self.counts.get_mut(&entity) {
659            *count = count.saturating_sub(1);
660            if *count == 0 {
661                self.counts.remove(&entity);
662            }
663        }
664    }
665
666    fn invalidate(&mut self) {
667        self.counts.clear();
668        self.decodable = false;
669    }
670}
671
672#[cfg(test)]
673mod tests;
icydb_core/db/data/store.rs

icydb_core/db/data/
store.rs
