icydb_core/obs/snapshot/
mod.rs

1use crate::{
2    db::{
3        Db, ensure_recovered,
4        identity::EntityName,
5        index::IndexKey,
6        store::{DataKey, StorageKey},
7    },
8    error::InternalError,
9    traits::CanisterKind,
10    value::Value,
11};
12use candid::CandidType;
13use serde::{Deserialize, Serialize};
14use std::collections::BTreeMap;
15
16///
17/// StorageReport
18/// Live storage snapshot report
19///
20
21#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
22pub struct StorageReport {
23    pub storage_data: Vec<DataStoreSnapshot>,
24    pub storage_index: Vec<IndexStoreSnapshot>,
25    pub entity_storage: Vec<EntitySnapshot>,
26    pub corrupted_keys: u64,
27    pub corrupted_entries: u64,
28}
29
30///
31/// DataStoreSnapshot
32/// Store-level snapshot metrics.
33///
34
35#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
36pub struct DataStoreSnapshot {
37    pub path: String,
38    pub entries: u64,
39    pub memory_bytes: u64,
40}
41
42///
43/// IndexStoreSnapshot
44/// Index-store snapshot metrics
45///
46
47#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
48pub struct IndexStoreSnapshot {
49    pub path: String,
50    pub entries: u64,
51    pub memory_bytes: u64,
52}
53
54///
55/// EntitySnapshot
56/// Per-entity storage breakdown across stores
57///
58
59#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
60pub struct EntitySnapshot {
61    /// Store path (e.g., icydb_schema_tests::schema::TestDataStore)
62    pub store: String,
63
64    /// Entity path (e.g., icydb_schema_tests::canister::db::Index)
65    pub path: String,
66
67    /// Number of rows for this entity in the store
68    pub entries: u64,
69
70    /// Approximate bytes used (key + value)
71    pub memory_bytes: u64,
72
73    /// Minimum primary key for this entity (entity-local ordering)
74    pub min_key: Option<Value>,
75
76    /// Maximum primary key for this entity (entity-local ordering)
77    pub max_key: Option<Value>,
78}
79
80///
81/// EntityStats
82/// Internal struct for building per-entity stats before snapshotting.
83///
84
85#[derive(Default)]
86struct EntityStats {
87    entries: u64,
88    memory_bytes: u64,
89    min_key: Option<StorageKey>,
90    max_key: Option<StorageKey>,
91}
92
93impl EntityStats {
94    fn update(&mut self, dk: &DataKey, value_len: u64) {
95        self.entries = self.entries.saturating_add(1);
96        self.memory_bytes = self
97            .memory_bytes
98            .saturating_add(DataKey::entry_size_bytes(value_len));
99
100        let k = dk.storage_key();
101
102        match &mut self.min_key {
103            Some(min) if k < *min => *min = k,
104            None => self.min_key = Some(k),
105            _ => {}
106        }
107
108        match &mut self.max_key {
109            Some(max) if k > *max => *max = k,
110            None => self.max_key = Some(k),
111            _ => {}
112        }
113    }
114}
115
116/// Build storage snapshot and per-entity breakdown; enrich path names using name→path map
117pub fn storage_report<C: CanisterKind>(
118    db: &Db<C>,
119    name_to_path: &[(&'static str, &'static str)],
120) -> Result<StorageReport, InternalError> {
121    ensure_recovered(db)?;
122    // Build name→path map once, reuse across stores
123    let name_map: BTreeMap<&'static str, &str> = name_to_path.iter().copied().collect();
124    let mut data = Vec::new();
125    let mut index = Vec::new();
126    let mut entity_storage: Vec<EntitySnapshot> = Vec::new();
127    let mut corrupted_keys = 0u64;
128    let mut corrupted_entries = 0u64;
129
130    db.with_data(|reg| {
131        reg.for_each(|path, store| {
132            data.push(DataStoreSnapshot {
133                path: path.to_string(),
134                entries: store.len(),
135                memory_bytes: store.memory_bytes(),
136            });
137
138            // Track per-entity counts, memory, and min/max Keys (not DataKeys)
139            let mut by_entity: BTreeMap<EntityName, EntityStats> = BTreeMap::new();
140
141            for entry in store.iter() {
142                let Ok(dk) = DataKey::try_from_raw(entry.key()) else {
143                    corrupted_keys = corrupted_keys.saturating_add(1);
144                    continue;
145                };
146
147                let value_len = entry.value().len() as u64;
148
149                by_entity
150                    .entry(*dk.entity_name())
151                    .or_default()
152                    .update(&dk, value_len);
153            }
154
155            for (entity_name, stats) in by_entity {
156                let path_name = name_map.get(entity_name.as_str()).copied().unwrap_or("");
157                entity_storage.push(EntitySnapshot {
158                    store: path.to_string(),
159                    path: path_name.to_string(),
160                    entries: stats.entries,
161                    memory_bytes: stats.memory_bytes,
162                    min_key: stats.min_key.map(|key| key.as_value()),
163                    max_key: stats.max_key.map(|key| key.as_value()),
164                });
165            }
166        });
167    });
168
169    db.with_index(|reg| {
170        reg.for_each(|path, store| {
171            index.push(IndexStoreSnapshot {
172                path: path.to_string(),
173                entries: store.len(),
174                memory_bytes: store.memory_bytes(),
175            });
176
177            for (key, value) in store.entries() {
178                if IndexKey::try_from_raw(&key).is_err() {
179                    corrupted_entries = corrupted_entries.saturating_add(1);
180                    continue;
181                }
182                if value.validate().is_err() {
183                    corrupted_entries = corrupted_entries.saturating_add(1);
184                }
185            }
186        });
187    });
188
189    Ok(StorageReport {
190        storage_data: data,
191        storage_index: index,
192        entity_storage,
193        corrupted_keys,
194        corrupted_entries,
195    })
196}
197
198#[cfg(test)]
199mod tests {
200    use crate::{
201        db::{
202            Db,
203            index::{IndexId, IndexKey, IndexStore, IndexStoreRegistry, RawIndexEntry},
204            init_commit_store_for_tests,
205            store::{DataKey, DataStore, DataStoreRegistry, RawRow},
206        },
207        obs::snapshot::storage_report,
208        test_support::test_memory,
209        traits::{CanisterKind, Path, Storable},
210    };
211    use std::{borrow::Cow, cell::RefCell};
212
213    const DATA_STORE_PATH: &str = "snapshot_tests::DataStore";
214    const INDEX_STORE_PATH: &str = "snapshot_tests::IndexStore";
215
216    struct SnapshotTestCanister;
217
218    impl Path for SnapshotTestCanister {
219        const PATH: &'static str = "snapshot_tests::Canister";
220    }
221
222    impl CanisterKind for SnapshotTestCanister {}
223
224    thread_local! {
225        static SNAPSHOT_DATA_STORE: RefCell<DataStore> = RefCell::new(DataStore::init(test_memory(101)));
226        static SNAPSHOT_INDEX_STORE: RefCell<IndexStore> = RefCell::new(IndexStore::init(
227            test_memory(102),
228            test_memory(103),
229        ));
230        static SNAPSHOT_DATA_REGISTRY: DataStoreRegistry = {
231            let mut reg = DataStoreRegistry::new();
232            reg.register(DATA_STORE_PATH, &SNAPSHOT_DATA_STORE);
233            reg
234        };
235        static SNAPSHOT_INDEX_REGISTRY: IndexStoreRegistry = {
236            let mut reg = IndexStoreRegistry::new();
237            reg.register(INDEX_STORE_PATH, &SNAPSHOT_INDEX_STORE);
238            reg
239        };
240    }
241
242    static DB: Db<SnapshotTestCanister> =
243        Db::new(&SNAPSHOT_DATA_REGISTRY, &SNAPSHOT_INDEX_REGISTRY);
244
245    fn reset_snapshot_state() {
246        init_commit_store_for_tests().expect("commit store init should succeed");
247
248        DB.with_data(|reg| reg.with_store_mut(DATA_STORE_PATH, DataStore::clear))
249            .expect("data store reset should succeed");
250        DB.with_index(|reg| reg.with_store_mut(INDEX_STORE_PATH, IndexStore::clear))
251            .expect("index store reset should succeed");
252    }
253
254    #[test]
255    fn storage_report_lists_registered_store_snapshots() {
256        reset_snapshot_state();
257
258        let report = storage_report(&DB, &[]).expect("storage report should succeed");
259        assert_eq!(report.storage_data.len(), 1);
260        assert_eq!(report.storage_data[0].path, DATA_STORE_PATH);
261        assert_eq!(report.storage_data[0].entries, 0);
262        assert_eq!(report.storage_index.len(), 1);
263        assert_eq!(report.storage_index[0].path, INDEX_STORE_PATH);
264        assert_eq!(report.storage_index[0].entries, 0);
265        assert!(report.entity_storage.is_empty());
266        assert_eq!(report.corrupted_keys, 0);
267        assert_eq!(report.corrupted_entries, 0);
268    }
269
270    #[test]
271    fn storage_report_counts_entity_rows_and_corrupted_index_entries() {
272        reset_snapshot_state();
273
274        let data_key = DataKey::max_storable()
275            .to_raw()
276            .expect("max storable data key should encode");
277        let row = RawRow::try_new(vec![1, 2, 3]).expect("row bytes should be valid");
278        DB.with_data(|reg| {
279            reg.with_store_mut(DATA_STORE_PATH, |store| {
280                store.insert(data_key, row);
281            })
282        })
283        .expect("data insert should succeed");
284
285        let index_key = IndexKey::empty(IndexId::max_storable()).to_raw();
286        let malformed_index_entry = RawIndexEntry::from_bytes(Cow::Owned(vec![0, 0, 0, 0]));
287        DB.with_index(|reg| {
288            reg.with_store_mut(INDEX_STORE_PATH, |store| {
289                store.insert(index_key, malformed_index_entry);
290            })
291        })
292        .expect("index insert should succeed");
293
294        let report = storage_report(&DB, &[]).expect("storage report should succeed");
295        assert_eq!(report.storage_data[0].entries, 1);
296        assert_eq!(report.storage_index[0].entries, 1);
297        assert_eq!(report.entity_storage.len(), 1);
298        assert_eq!(report.entity_storage[0].path, "");
299        assert_eq!(report.entity_storage[0].entries, 1);
300        assert!(report.entity_storage[0].min_key.is_some());
301        assert!(report.entity_storage[0].max_key.is_some());
302        assert_eq!(report.corrupted_entries, 1);
303        assert_eq!(report.corrupted_keys, 0);
304    }
305}
icydb_core/obs/snapshot/mod.rs

icydb_core/obs/snapshot/
mod.rs
