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(all(feature = "sql", 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(all(feature = "sql", feature = "diagnostics"))]
24thread_local! {
25    static DATA_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
26}
27
28#[cfg(all(feature = "sql", 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(all(feature = "sql", 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    pub(in crate::db) fn visit_range_rev<E>(
416        &self,
417        key_range: impl RangeBounds<RawDataStoreKey>,
418        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
419    ) -> Result<(), E> {
420        let bounds = Self::owned_range_bounds(&key_range);
421        match &self.backend {
422            DataStoreBackend::Heap(map) => {
423                for (key, row) in map.range((bounds.0.clone(), bounds.1)).rev() {
424                    if visitor(key, row)?.should_stop() {
425                        break;
426                    }
427                }
428            }
429            DataStoreBackend::Journaled {
430                canonical: _,
431                live: _,
432                tombstones: _,
433            } => Self::visit_journaled_entries_in_bounds(&self.backend, bounds, true, visitor)?,
434        }
435
436        Ok(())
437    }
438
439    /// Sum of bytes used by all stored rows.
440    pub(in crate::db) fn memory_bytes(&self) -> u64 {
441        // Report map footprint as key bytes + row bytes per entry.
442        let mut bytes = 0u64;
443        let _: Result<(), Infallible> = self.visit_entries(|key, row| {
444            bytes = bytes.saturating_add(key.as_bytes().len() as u64 + row.len() as u64);
445            Ok(StoreVisit::Continue)
446        });
447        bytes
448    }
449
450    const fn bump_generation(&mut self) {
451        self.generation = self.generation.saturating_add(1);
452    }
453
454    fn rebuild_entity_cardinality_from_entries(&mut self) {
455        let mut cardinality = EntityCardinality::empty();
456        let _: Result<(), Infallible> = self.visit_entries(|key, _row| {
457            cardinality.apply_present_key(key);
458            Ok(StoreVisit::Continue)
459        });
460        self.entity_cardinality = cardinality;
461    }
462
463    /// Return the monotonic perf-only count of stable row fetches seen by this process.
464    #[cfg(all(feature = "sql", feature = "diagnostics"))]
465    pub(in crate::db) fn current_get_call_count() -> u64 {
466        DATA_STORE_GET_CALL_COUNT.with(Cell::get)
467    }
468
469    fn journaled_get_raw(backend: &DataStoreBackend, key: &RawDataStoreKey) -> Option<RawRow> {
470        let DataStoreBackend::Journaled {
471            canonical,
472            live,
473            tombstones,
474        } = backend
475        else {
476            return None;
477        };
478
479        if tombstones.contains(key) {
480            return None;
481        }
482        live.get(key).cloned().or_else(|| canonical.get(key))
483    }
484
485    fn owned_range_bounds(
486        key_range: &impl RangeBounds<RawDataStoreKey>,
487    ) -> (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>) {
488        let lower = match key_range.start_bound() {
489            Bound::Included(key) => Bound::Included(key.clone()),
490            Bound::Excluded(key) => Bound::Excluded(key.clone()),
491            Bound::Unbounded => Bound::Unbounded,
492        };
493        let upper = match key_range.end_bound() {
494            Bound::Included(key) => Bound::Included(key.clone()),
495            Bound::Excluded(key) => Bound::Excluded(key.clone()),
496            Bound::Unbounded => Bound::Unbounded,
497        };
498
499        (lower, upper)
500    }
501
502    fn visit_journaled_entries_in_bounds<E>(
503        backend: &DataStoreBackend,
504        bounds: (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>),
505        reverse: bool,
506        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
507    ) -> Result<(), E> {
508        let DataStoreBackend::Journaled {
509            canonical,
510            live,
511            tombstones,
512        } = backend
513        else {
514            return Ok(());
515        };
516
517        if canonical.is_empty() {
518            if reverse {
519                for (key, row) in live.range(bounds).rev() {
520                    if visitor(key, row)?.should_stop() {
521                        return Ok(());
522                    }
523                }
524            } else {
525                for (key, row) in live.range(bounds) {
526                    if visitor(key, row)?.should_stop() {
527                        return Ok(());
528                    }
529                }
530            }
531            return Ok(());
532        }
533
534        if live.is_empty() && tombstones.is_empty() {
535            if reverse {
536                for entry in canonical.range(bounds).rev() {
537                    if visitor(entry.key(), &entry.value())?.should_stop() {
538                        return Ok(());
539                    }
540                }
541            } else {
542                for entry in canonical.range(bounds) {
543                    if visitor(entry.key(), &entry.value())?.should_stop() {
544                        return Ok(());
545                    }
546                }
547            }
548            return Ok(());
549        }
550
551        match if reverse {
552            Direction::Desc
553        } else {
554            Direction::Asc
555        } {
556            Direction::Asc => visit_ordered_overlay(
557                canonical.range((bounds.0.clone(), bounds.1.clone())),
558                live.range((bounds.0, bounds.1)),
559                Direction::Asc,
560                |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
561                |canonical_entry| !tombstones.contains(canonical_entry.key()),
562                |live_entry| !tombstones.contains(live_entry.0),
563                |entry| {
564                    let visit = match entry {
565                        OrderedOverlayEntry::Canonical(canonical_entry) => {
566                            visitor(canonical_entry.key(), &canonical_entry.value())?
567                        }
568                        OrderedOverlayEntry::Live((key, row)) => visitor(key, row)?,
569                    };
570                    Ok(if visit.should_stop() {
571                        OrderedOverlayVisit::Stop
572                    } else {
573                        OrderedOverlayVisit::Continue
574                    })
575                },
576            ),
577            Direction::Desc => visit_ordered_overlay(
578                canonical.range((bounds.0.clone(), bounds.1.clone())).rev(),
579                live.range((bounds.0, bounds.1)).rev(),
580                Direction::Desc,
581                |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
582                |canonical_entry| !tombstones.contains(canonical_entry.key()),
583                |live_entry| !tombstones.contains(live_entry.0),
584                |entry| {
585                    let visit = match entry {
586                        OrderedOverlayEntry::Canonical(canonical_entry) => {
587                            visitor(canonical_entry.key(), &canonical_entry.value())?
588                        }
589                        OrderedOverlayEntry::Live((key, row)) => visitor(key, row)?,
590                    };
591                    Ok(if visit.should_stop() {
592                        OrderedOverlayVisit::Stop
593                    } else {
594                        OrderedOverlayVisit::Continue
595                    })
596                },
597            ),
598        }
599    }
600}
601
602#[derive(Clone, Debug)]
603struct EntityCardinality {
604    counts: HeapBTreeMap<EntityTag, u64>,
605    decodable: bool,
606}
607
608impl EntityCardinality {
609    const fn empty() -> Self {
610        Self {
611            counts: HeapBTreeMap::new(),
612            decodable: true,
613        }
614    }
615
616    fn exact_count(&self, entity: EntityTag) -> Option<u64> {
617        self.decodable
618            .then(|| self.counts.get(&entity).copied().unwrap_or(0))
619    }
620
621    fn apply_insert(&mut self, key: &RawDataStoreKey, previous: Option<&RawRow>) {
622        if previous.is_some() {
623            return;
624        }
625        self.apply_present_key(key);
626    }
627
628    fn apply_remove(&mut self, key: &RawDataStoreKey, previous: Option<&RawRow>) {
629        if previous.is_none() {
630            return;
631        }
632        self.apply_removed_key(key);
633    }
634
635    fn apply_present_key(&mut self, key: &RawDataStoreKey) {
636        if !self.decodable {
637            return;
638        }
639        let Some(entity) = key.entity_tag_prefix() else {
640            self.invalidate();
641            return;
642        };
643
644        let count = self.counts.entry(entity).or_insert(0);
645        *count = count.saturating_add(1);
646    }
647
648    fn apply_removed_key(&mut self, key: &RawDataStoreKey) {
649        if !self.decodable {
650            return;
651        }
652        let Some(entity) = key.entity_tag_prefix() else {
653            self.invalidate();
654            return;
655        };
656
657        if let Some(count) = self.counts.get_mut(&entity) {
658            *count = count.saturating_sub(1);
659            if *count == 0 {
660                self.counts.remove(&entity);
661            }
662        }
663    }
664
665    fn invalidate(&mut self) {
666        self.counts.clear();
667        self.decodable = false;
668    }
669}
670
671#[cfg(test)]
672mod tests;
icydb_core/db/data/store.rs

icydb_core/db/data/
store.rs
