icydb_core/db/data/
store.rs

1//! Module: data::store
2//! Responsibility: stable BTreeMap-backed row persistence.
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        key_taxonomy::RawDataStoreKeyRange,
10    },
11    types::EntityTag,
12};
13use ic_memory::stable_structures::{
14    BTreeMap, DefaultMemoryImpl, btreemap::Iter, memory_manager::VirtualMemory,
15};
16#[cfg(feature = "diagnostics")]
17use std::cell::Cell;
18use std::ops::RangeBounds;
19
20#[cfg(feature = "diagnostics")]
21thread_local! {
22    static DATA_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
23}
24
25#[cfg(feature = "diagnostics")]
26fn record_data_store_get_call() {
27    DATA_STORE_GET_CALL_COUNT.with(|count| {
28        count.set(count.get().saturating_add(1));
29    });
30}
31
32///
33/// DataStore
34///
35/// Thin persistence wrapper over one stable BTreeMap.
36///
37/// Invariant: callers provide already-validated `RawDataStoreKey` and canonical row bytes.
38/// This type intentionally does not enforce commit-phase ordering.
39///
40
41pub struct DataStore {
42    map: BTreeMap<RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>>,
43}
44
45impl DataStore {
46    /// Initialize a data store with the provided backing memory.
47    #[must_use]
48    pub fn init(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
49        Self {
50            map: BTreeMap::init(memory),
51        }
52    }
53
54    /// Insert or replace one row by raw key.
55    pub(in crate::db) fn insert(
56        &mut self,
57        key: RawDataStoreKey,
58        row: CanonicalRow,
59    ) -> Option<RawRow> {
60        self.map.insert(key, row.into_raw_row())
61    }
62
63    /// Insert one raw row directly for corruption-focused test setup only.
64    #[cfg(test)]
65    pub(in crate::db) fn insert_raw_for_test(
66        &mut self,
67        key: RawDataStoreKey,
68        row: RawRow,
69    ) -> Option<RawRow> {
70        self.map.insert(key, row)
71    }
72
73    /// Remove one row by raw key.
74    pub(in crate::db) fn remove(&mut self, key: &RawDataStoreKey) -> Option<RawRow> {
75        self.map.remove(key)
76    }
77
78    /// Load one row by raw key.
79    pub(in crate::db) fn get(&self, key: &RawDataStoreKey) -> Option<RawRow> {
80        #[cfg(feature = "diagnostics")]
81        record_data_store_get_call();
82
83        self.map.get(key)
84    }
85
86    /// Return whether one raw key exists without cloning the row payload.
87    #[must_use]
88    pub(in crate::db) fn contains(&self, key: &RawDataStoreKey) -> bool {
89        self.map.contains_key(key)
90    }
91
92    /// Clear all stored rows from the data store.
93    #[cfg(test)]
94    pub(in crate::db) fn clear(&mut self) {
95        self.map.clear_new();
96    }
97
98    /// Return the number of stored rows without exposing the backing map.
99    #[must_use]
100    pub(in crate::db) fn len(&self) -> u64 {
101        self.map.len()
102    }
103
104    /// Return whether the data store currently contains no rows.
105    #[must_use]
106    #[cfg(test)]
107    pub(in crate::db) fn is_empty(&self) -> bool {
108        self.map.is_empty()
109    }
110
111    /// Return raw row entries in canonical storage order.
112    pub(in crate::db) fn entries(
113        &self,
114    ) -> Iter<'_, RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
115        self.map.iter()
116    }
117
118    /// Iterate over raw row entries whose keys belong to the provided storage range.
119    pub(in crate::db) fn range(
120        &self,
121        key_range: impl RangeBounds<RawDataStoreKey>,
122    ) -> Iter<'_, RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
123        self.map.range(key_range)
124    }
125
126    /// Iterate over raw row entries for one entity using compact prefix bounds.
127    pub(in crate::db) fn range_for_entity(
128        &self,
129        entity: EntityTag,
130    ) -> Iter<'_, RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
131        let range = RawDataStoreKeyRange::entity_prefix(entity);
132        self.map.range(RawDataStoreKey::store_range_bounds(&range))
133    }
134
135    /// Sum of bytes used by all stored rows.
136    pub(in crate::db) fn memory_bytes(&self) -> u64 {
137        // Report map footprint as key bytes + row bytes per entry.
138        self.entries()
139            .map(|entry| entry.key().as_bytes().len() as u64 + entry.value().len() as u64)
140            .sum()
141    }
142
143    /// Return the monotonic perf-only count of stable row fetches seen by this process.
144    #[cfg(feature = "diagnostics")]
145    pub(in crate::db) fn current_get_call_count() -> u64 {
146        DATA_STORE_GET_CALL_COUNT.with(Cell::get)
147    }
148}
icydb_core/db/data/store.rs

icydb_core/db/data/
store.rs
