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_store_registry(|reg| {
131        reg.iter().for_each(|(path, store_handle)| {
132            store_handle.with_data(|store| {
133                data.push(DataStoreSnapshot {
134                    path: path.to_string(),
135                    entries: store.len(),
136                    memory_bytes: store.memory_bytes(),
137                });
138
139                // Track per-entity counts, memory, and min/max Keys (not DataKeys)
140                let mut by_entity: BTreeMap<EntityName, EntityStats> = BTreeMap::new();
141
142                for entry in store.iter() {
143                    let Ok(dk) = DataKey::try_from_raw(entry.key()) else {
144                        corrupted_keys = corrupted_keys.saturating_add(1);
145                        continue;
146                    };
147
148                    let value_len = entry.value().len() as u64;
149
150                    by_entity
151                        .entry(*dk.entity_name())
152                        .or_default()
153                        .update(&dk, value_len);
154                }
155
156                for (entity_name, stats) in by_entity {
157                    let path_name = name_map.get(entity_name.as_str()).copied().unwrap_or("");
158                    entity_storage.push(EntitySnapshot {
159                        store: path.to_string(),
160                        path: path_name.to_string(),
161                        entries: stats.entries,
162                        memory_bytes: stats.memory_bytes,
163                        min_key: stats.min_key.map(|key| key.as_value()),
164                        max_key: stats.max_key.map(|key| key.as_value()),
165                    });
166                }
167            });
168        });
169    });
170
171    db.with_store_registry(|reg| {
172        reg.iter().for_each(|(path, store_handle)| {
173            store_handle.with_index(|store| {
174                index.push(IndexStoreSnapshot {
175                    path: path.to_string(),
176                    entries: store.len(),
177                    memory_bytes: store.memory_bytes(),
178                });
179
180                for (key, value) in store.entries() {
181                    if IndexKey::try_from_raw(&key).is_err() {
182                        corrupted_entries = corrupted_entries.saturating_add(1);
183                        continue;
184                    }
185                    if value.validate().is_err() {
186                        corrupted_entries = corrupted_entries.saturating_add(1);
187                    }
188                }
189            });
190        });
191    });
192
193    Ok(StorageReport {
194        storage_data: data,
195        storage_index: index,
196        entity_storage,
197        corrupted_keys,
198        corrupted_entries,
199    })
200}
201
202#[cfg(test)]
203mod tests {
204    use crate::{
205        db::{
206            Db,
207            index::{IndexId, IndexKey, IndexStore, RawIndexEntry},
208            init_commit_store_for_tests,
209            store::{DataKey, DataStore, RawRow, StoreRegistry},
210        },
211        obs::snapshot::storage_report,
212        test_support::test_memory,
213        traits::{CanisterKind, Path, Storable},
214    };
215    use std::{borrow::Cow, cell::RefCell};
216
217    const STORE_PATH: &str = "snapshot_tests::Store";
218
219    struct SnapshotTestCanister;
220
221    impl Path for SnapshotTestCanister {
222        const PATH: &'static str = "snapshot_tests::Canister";
223    }
224
225    impl CanisterKind for SnapshotTestCanister {}
226
227    thread_local! {
228        static SNAPSHOT_DATA_STORE: RefCell<DataStore> = RefCell::new(DataStore::init(test_memory(101)));
229        static SNAPSHOT_INDEX_STORE: RefCell<IndexStore> =
230            RefCell::new(IndexStore::init(test_memory(102)));
231        static SNAPSHOT_STORE_REGISTRY: StoreRegistry = {
232            let mut reg = StoreRegistry::new();
233            reg.register_store(STORE_PATH, &SNAPSHOT_DATA_STORE, &SNAPSHOT_INDEX_STORE)
234                .expect("snapshot store registration should succeed");
235            reg
236        };
237    }
238
239    static DB: Db<SnapshotTestCanister> = Db::new(&SNAPSHOT_STORE_REGISTRY);
240
241    fn with_snapshot_store<R>(f: impl FnOnce(crate::db::store::StoreHandle) -> R) -> R {
242        DB.with_store_registry(|reg| reg.try_get_store(STORE_PATH).map(f))
243            .expect("snapshot store access should succeed")
244    }
245
246    fn reset_snapshot_state() {
247        init_commit_store_for_tests().expect("commit store init should succeed");
248
249        with_snapshot_store(|store| {
250            store.with_data_mut(DataStore::clear);
251            store.with_index_mut(IndexStore::clear);
252        });
253    }
254
255    #[test]
256    fn storage_report_lists_registered_store_snapshots() {
257        reset_snapshot_state();
258
259        let report = storage_report(&DB, &[]).expect("storage report should succeed");
260        assert_eq!(report.storage_data.len(), 1);
261        assert_eq!(report.storage_data[0].path, STORE_PATH);
262        assert_eq!(report.storage_data[0].entries, 0);
263        assert_eq!(report.storage_index.len(), 1);
264        assert_eq!(report.storage_index[0].path, STORE_PATH);
265        assert_eq!(report.storage_index[0].entries, 0);
266        assert!(report.entity_storage.is_empty());
267        assert_eq!(report.corrupted_keys, 0);
268        assert_eq!(report.corrupted_entries, 0);
269    }
270
271    #[test]
272    fn storage_report_counts_entity_rows_and_corrupted_index_entries() {
273        reset_snapshot_state();
274
275        let data_key = DataKey::max_storable()
276            .to_raw()
277            .expect("max storable data key should encode");
278        let row = RawRow::try_new(vec![1, 2, 3]).expect("row bytes should be valid");
279        with_snapshot_store(|store| {
280            store.with_data_mut(|data_store| {
281                data_store.insert(data_key, row);
282            });
283        });
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        with_snapshot_store(|store| {
288            store.with_index_mut(|index_store| {
289                index_store.insert(index_key, malformed_index_entry);
290            });
291        });
292
293        let report = storage_report(&DB, &[]).expect("storage report should succeed");
294        assert_eq!(report.storage_data[0].entries, 1);
295        assert_eq!(report.storage_index[0].entries, 1);
296        assert_eq!(report.entity_storage.len(), 1);
297        assert_eq!(report.entity_storage[0].path, "");
298        assert_eq!(report.entity_storage[0].entries, 1);
299        assert!(report.entity_storage[0].min_key.is_some());
300        assert!(report.entity_storage[0].max_key.is_some());
301        assert_eq!(report.corrupted_entries, 1);
302        assert_eq!(report.corrupted_keys, 0);
303    }
304}
icydb_core/obs/snapshot/mod.rs

icydb_core/obs/snapshot/
mod.rs
