icydb_core/db/diagnostics/snapshot/
mod.rs

1use crate::{
2    db::{
3        Db, EntityName,
4        data::{DataKey, StorageKey},
5        index::IndexKey,
6    },
7    error::InternalError,
8    traits::CanisterKind,
9    value::Value,
10};
11use candid::CandidType;
12use serde::{Deserialize, Serialize};
13use std::collections::BTreeMap;
14
15///
16/// StorageReport
17/// Live storage snapshot report
18///
19
20#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
21pub struct StorageReport {
22    pub storage_data: Vec<DataStoreSnapshot>,
23    pub storage_index: Vec<IndexStoreSnapshot>,
24    pub entity_storage: Vec<EntitySnapshot>,
25    pub corrupted_keys: u64,
26    pub corrupted_entries: u64,
27}
28
29///
30/// DataStoreSnapshot
31/// Store-level snapshot metrics.
32///
33
34#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
35pub struct DataStoreSnapshot {
36    pub path: String,
37    pub entries: u64,
38    pub memory_bytes: u64,
39}
40
41///
42/// IndexStoreSnapshot
43/// Index-store snapshot metrics
44///
45
46#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
47pub struct IndexStoreSnapshot {
48    pub path: String,
49    pub entries: u64,
50    pub user_entries: u64,
51    pub system_entries: u64,
52    pub memory_bytes: u64,
53}
54
55///
56/// EntitySnapshot
57/// Per-entity storage breakdown across stores
58///
59
60#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
61pub struct EntitySnapshot {
62    /// Store path (e.g., icydb_schema_tests::schema::TestDataStore)
63    pub store: String,
64
65    /// Entity path (e.g., icydb_schema_tests::canister::db::Index)
66    pub path: String,
67
68    /// Number of rows for this entity in the store
69    pub entries: u64,
70
71    /// Approximate bytes used (key + value)
72    pub memory_bytes: u64,
73
74    /// Minimum primary key for this entity (entity-local ordering)
75    pub min_key: Option<Value>,
76
77    /// Maximum primary key for this entity (entity-local ordering)
78    pub max_key: Option<Value>,
79}
80
81///
82/// EntityStats
83/// Internal struct for building per-entity stats before snapshotting.
84///
85
86#[derive(Default)]
87struct EntityStats {
88    entries: u64,
89    memory_bytes: u64,
90    min_key: Option<StorageKey>,
91    max_key: Option<StorageKey>,
92}
93
94impl EntityStats {
95    fn update(&mut self, dk: &DataKey, value_len: u64) {
96        self.entries = self.entries.saturating_add(1);
97        self.memory_bytes = self
98            .memory_bytes
99            .saturating_add(DataKey::entry_size_bytes(value_len));
100
101        let k = dk.storage_key();
102
103        match &mut self.min_key {
104            Some(min) if k < *min => *min = k,
105            None => self.min_key = Some(k),
106            _ => {}
107        }
108
109        match &mut self.max_key {
110            Some(max) if k > *max => *max = k,
111            None => self.max_key = Some(k),
112            _ => {}
113        }
114    }
115}
116
117/// Build storage snapshot and per-entity breakdown; enrich path names using name→path map
118pub(crate) fn storage_report<C: CanisterKind>(
119    db: &Db<C>,
120    name_to_path: &[(&'static str, &'static str)],
121) -> Result<StorageReport, InternalError> {
122    db.ensure_recovered_state()?;
123    // Build name→path map once, reuse across stores.
124    let name_map: BTreeMap<&'static str, &str> = name_to_path.iter().copied().collect();
125    let mut data = Vec::new();
126    let mut index = Vec::new();
127    let mut entity_storage: Vec<EntitySnapshot> = Vec::new();
128    let mut corrupted_keys = 0u64;
129    let mut corrupted_entries = 0u64;
130
131    db.with_store_registry(|reg| {
132        // Keep diagnostics snapshots deterministic by traversing stores in path order.
133        let mut stores = reg.iter().collect::<Vec<_>>();
134        stores.sort_by_key(|(path, _)| *path);
135
136        for (path, store_handle) in stores {
137            // Phase 1: collect data-store snapshots and per-entity stats.
138            store_handle.with_data(|store| {
139                data.push(DataStoreSnapshot {
140                    path: path.to_string(),
141                    entries: store.len(),
142                    memory_bytes: store.memory_bytes(),
143                });
144
145                // Track per-entity counts, memory, and min/max Keys (not DataKeys)
146                let mut by_entity: BTreeMap<EntityName, EntityStats> = BTreeMap::new();
147
148                for entry in store.iter() {
149                    let Ok(dk) = DataKey::try_from_raw(entry.key()) else {
150                        corrupted_keys = corrupted_keys.saturating_add(1);
151                        continue;
152                    };
153
154                    let value_len = entry.value().len() as u64;
155
156                    by_entity
157                        .entry(*dk.entity_name())
158                        .or_default()
159                        .update(&dk, value_len);
160                }
161
162                for (entity_name, stats) in by_entity {
163                    let path_name = name_map
164                        .get(entity_name.as_str())
165                        .copied()
166                        .unwrap_or(entity_name.as_str());
167                    entity_storage.push(EntitySnapshot {
168                        store: path.to_string(),
169                        path: path_name.to_string(),
170                        entries: stats.entries,
171                        memory_bytes: stats.memory_bytes,
172                        min_key: stats.min_key.map(|key| key.as_value()),
173                        max_key: stats.max_key.map(|key| key.as_value()),
174                    });
175                }
176            });
177
178            // Phase 2: collect index-store snapshots and integrity counters.
179            store_handle.with_index(|store| {
180                let mut user_entries = 0u64;
181                let mut system_entries = 0u64;
182
183                for (key, value) in store.entries() {
184                    let Ok(decoded_key) = IndexKey::try_from_raw(&key) else {
185                        corrupted_entries = corrupted_entries.saturating_add(1);
186                        continue;
187                    };
188
189                    if decoded_key.uses_system_namespace() {
190                        system_entries = system_entries.saturating_add(1);
191                    } else {
192                        user_entries = user_entries.saturating_add(1);
193                    }
194
195                    if value.validate().is_err() {
196                        corrupted_entries = corrupted_entries.saturating_add(1);
197                    }
198                }
199
200                index.push(IndexStoreSnapshot {
201                    path: path.to_string(),
202                    entries: store.len(),
203                    user_entries,
204                    system_entries,
205                    memory_bytes: store.memory_bytes(),
206                });
207            });
208        }
209    });
210
211    Ok(StorageReport {
212        storage_data: data,
213        storage_index: index,
214        entity_storage,
215        corrupted_keys,
216        corrupted_entries,
217    })
218}
219
220///
221/// TESTS
222///
223
224#[cfg(test)]
225mod tests {
226    use crate::{
227        db::{
228            Db,
229            commit::{ensure_recovered_for_write, init_commit_store_for_tests},
230            data::{DataKey, DataStore, RawDataKey, RawRow, StorageKey},
231            identity::{EntityName, IndexName},
232            index::{IndexId, IndexKey, IndexKeyKind, IndexStore, RawIndexEntry, RawIndexKey},
233            registry::StoreRegistry,
234        },
235        test_support::test_memory,
236        traits::Storable,
237    };
238    use std::{borrow::Cow, cell::RefCell};
239
240    use super::{StorageReport, storage_report};
241
242    crate::test_canister! {
243        ident = DiagnosticsCanister,
244    }
245
246    const STORE_Z_PATH: &str = "diagnostics_tests::z_store";
247    const STORE_A_PATH: &str = "diagnostics_tests::a_store";
248    const SINGLE_ENTITY_NAME: &str = "diag_single_entity";
249    const SINGLE_ENTITY_PATH: &str = "diagnostics_tests::entity::single";
250    const FIRST_ENTITY_NAME: &str = "diag_first_entity";
251    const FIRST_ENTITY_PATH: &str = "diagnostics_tests::entity::first";
252    const SECOND_ENTITY_NAME: &str = "diag_second_entity";
253    const SECOND_ENTITY_PATH: &str = "diagnostics_tests::entity::second";
254    const MINMAX_ENTITY_NAME: &str = "diag_minmax_entity";
255    const MINMAX_ENTITY_PATH: &str = "diagnostics_tests::entity::minmax";
256    const VALID_ENTITY_NAME: &str = "diag_valid_entity";
257    const VALID_ENTITY_PATH: &str = "diagnostics_tests::entity::valid";
258
259    thread_local! {
260        static STORE_Z_DATA: RefCell<DataStore> = RefCell::new(DataStore::init(test_memory(153)));
261        static STORE_Z_INDEX: RefCell<IndexStore> = RefCell::new(IndexStore::init(test_memory(154)));
262        static STORE_A_DATA: RefCell<DataStore> = RefCell::new(DataStore::init(test_memory(155)));
263        static STORE_A_INDEX: RefCell<IndexStore> = RefCell::new(IndexStore::init(test_memory(156)));
264        static DIAGNOSTICS_REGISTRY: StoreRegistry = {
265            let mut registry = StoreRegistry::new();
266            registry
267                .register_store(STORE_Z_PATH, &STORE_Z_DATA, &STORE_Z_INDEX)
268                .expect("diagnostics test z-store registration should succeed");
269            registry
270                .register_store(STORE_A_PATH, &STORE_A_DATA, &STORE_A_INDEX)
271                .expect("diagnostics test a-store registration should succeed");
272            registry
273        };
274    }
275
276    static DB: Db<DiagnosticsCanister> = Db::new(&DIAGNOSTICS_REGISTRY);
277
278    fn with_data_store_mut<R>(path: &'static str, f: impl FnOnce(&mut DataStore) -> R) -> R {
279        DB.with_store_registry(|registry| {
280            registry
281                .try_get_store(path)
282                .map(|store_handle| store_handle.with_data_mut(f))
283        })
284        .expect("data store lookup should succeed")
285    }
286
287    fn with_index_store_mut<R>(path: &'static str, f: impl FnOnce(&mut IndexStore) -> R) -> R {
288        DB.with_store_registry(|registry| {
289            registry
290                .try_get_store(path)
291                .map(|store_handle| store_handle.with_index_mut(f))
292        })
293        .expect("index store lookup should succeed")
294    }
295
296    fn reset_stores() {
297        init_commit_store_for_tests().expect("commit store init should succeed");
298        ensure_recovered_for_write(&DB).expect("write-side recovery should succeed");
299        DB.with_store_registry(|registry| {
300            for (_, store_handle) in registry.iter() {
301                store_handle.with_data_mut(DataStore::clear);
302                store_handle.with_index_mut(IndexStore::clear);
303            }
304        });
305    }
306
307    fn insert_data_row(path: &'static str, entity_name: &str, key: StorageKey, row_len: usize) {
308        let entity =
309            EntityName::try_from_str(entity_name).expect("diagnostics test entity name is valid");
310        let raw_key = DataKey::raw_from_parts(entity, key)
311            .expect("diagnostics test data key should encode from valid parts");
312        let row_bytes = vec![0xAB; row_len.max(1)];
313        let raw_row = RawRow::try_new(row_bytes).expect("diagnostics test row should encode");
314
315        with_data_store_mut(path, |store| {
316            store.insert(raw_key, raw_row);
317        });
318    }
319
320    fn insert_corrupted_data_key(path: &'static str) {
321        let valid = DataKey::raw_from_parts(
322            EntityName::try_from_str(VALID_ENTITY_NAME).expect("valid test entity name"),
323            StorageKey::Int(1),
324        )
325        .expect("valid data key should encode");
326
327        let mut corrupted_bytes = valid.as_bytes().to_vec();
328        corrupted_bytes[0] = 0;
329        let corrupted_key = <RawDataKey as Storable>::from_bytes(Cow::Owned(corrupted_bytes));
330        let raw_row = RawRow::try_new(vec![0xCD]).expect("diagnostics test row should encode");
331
332        with_data_store_mut(path, |store| {
333            store.insert(corrupted_key, raw_row);
334        });
335    }
336
337    fn index_id(entity_name: &str, field: &str) -> IndexId {
338        let entity =
339            EntityName::try_from_str(entity_name).expect("diagnostics test entity name is valid");
340        let name = IndexName::try_from_parts(&entity, &[field])
341            .expect("diagnostics test index name should encode");
342
343        IndexId(name)
344    }
345
346    fn index_key(kind: IndexKeyKind, entity_name: &str, field: &str) -> RawIndexKey {
347        let id = index_id(entity_name, field);
348        IndexKey::empty_with_kind(&id, kind).to_raw()
349    }
350
351    fn insert_index_entry(path: &'static str, key: RawIndexKey, entry: RawIndexEntry) {
352        with_index_store_mut(path, |store| {
353            store.insert(key, entry);
354        });
355    }
356
357    fn diagnostics_report(name_to_path: &[(&'static str, &'static str)]) -> StorageReport {
358        storage_report(&DB, name_to_path).expect("diagnostics snapshot should succeed")
359    }
360
361    fn data_paths(report: &StorageReport) -> Vec<&str> {
362        report
363            .storage_data
364            .iter()
365            .map(|snapshot| snapshot.path.as_str())
366            .collect()
367    }
368
369    fn index_paths(report: &StorageReport) -> Vec<&str> {
370        report
371            .storage_index
372            .iter()
373            .map(|snapshot| snapshot.path.as_str())
374            .collect()
375    }
376
377    #[test]
378    fn storage_report_empty_store_snapshot() {
379        reset_stores();
380
381        let report = diagnostics_report(&[]);
382
383        assert_eq!(report.corrupted_keys, 0);
384        assert_eq!(report.corrupted_entries, 0);
385        assert!(report.entity_storage.is_empty());
386
387        assert_eq!(data_paths(&report), vec![STORE_A_PATH, STORE_Z_PATH]);
388        assert_eq!(index_paths(&report), vec![STORE_A_PATH, STORE_Z_PATH]);
389        assert!(
390            report
391                .storage_data
392                .iter()
393                .all(|snapshot| snapshot.entries == 0)
394        );
395        assert!(
396            report
397                .storage_index
398                .iter()
399                .all(|snapshot| snapshot.entries == 0)
400        );
401    }
402
403    #[test]
404    fn storage_report_single_entity_multiple_rows() {
405        reset_stores();
406
407        insert_data_row(STORE_A_PATH, SINGLE_ENTITY_NAME, StorageKey::Int(3), 3);
408        insert_data_row(STORE_A_PATH, SINGLE_ENTITY_NAME, StorageKey::Int(1), 1);
409        insert_data_row(STORE_A_PATH, SINGLE_ENTITY_NAME, StorageKey::Int(2), 2);
410
411        let report = diagnostics_report(&[(SINGLE_ENTITY_NAME, SINGLE_ENTITY_PATH)]);
412        let entity_snapshot = report
413            .entity_storage
414            .iter()
415            .find(|snapshot| snapshot.store == STORE_A_PATH && snapshot.path == SINGLE_ENTITY_PATH)
416            .expect("single-entity snapshot should exist");
417
418        assert_eq!(entity_snapshot.entries, 3);
419    }
420
421    #[test]
422    fn storage_report_multiple_entities_in_same_store() {
423        reset_stores();
424
425        insert_data_row(STORE_A_PATH, FIRST_ENTITY_NAME, StorageKey::Int(10), 1);
426        insert_data_row(STORE_A_PATH, FIRST_ENTITY_NAME, StorageKey::Int(11), 1);
427        insert_data_row(STORE_A_PATH, SECOND_ENTITY_NAME, StorageKey::Int(20), 1);
428
429        let report = diagnostics_report(&[
430            (FIRST_ENTITY_NAME, FIRST_ENTITY_PATH),
431            (SECOND_ENTITY_NAME, SECOND_ENTITY_PATH),
432        ]);
433
434        let first = report
435            .entity_storage
436            .iter()
437            .find(|snapshot| snapshot.store == STORE_A_PATH && snapshot.path == FIRST_ENTITY_PATH)
438            .expect("first-entity snapshot should exist");
439        let second = report
440            .entity_storage
441            .iter()
442            .find(|snapshot| snapshot.store == STORE_A_PATH && snapshot.path == SECOND_ENTITY_PATH)
443            .expect("second-entity snapshot should exist");
444
445        assert_eq!(first.entries, 2);
446        assert_eq!(second.entries, 1);
447    }
448
449    #[test]
450    fn storage_report_min_max_key_correctness() {
451        reset_stores();
452
453        insert_data_row(STORE_A_PATH, MINMAX_ENTITY_NAME, StorageKey::Int(9), 1);
454        insert_data_row(STORE_A_PATH, MINMAX_ENTITY_NAME, StorageKey::Int(-5), 1);
455        insert_data_row(STORE_A_PATH, MINMAX_ENTITY_NAME, StorageKey::Int(3), 1);
456
457        let report = diagnostics_report(&[(MINMAX_ENTITY_NAME, MINMAX_ENTITY_PATH)]);
458        let entity_snapshot = report
459            .entity_storage
460            .iter()
461            .find(|snapshot| snapshot.store == STORE_A_PATH && snapshot.path == MINMAX_ENTITY_PATH)
462            .expect("min/max snapshot should exist");
463
464        assert_eq!(
465            entity_snapshot.min_key,
466            Some(StorageKey::Int(-5).as_value())
467        );
468        assert_eq!(entity_snapshot.max_key, Some(StorageKey::Int(9).as_value()));
469    }
470
471    #[test]
472    fn storage_report_corrupted_key_detection() {
473        reset_stores();
474
475        insert_data_row(STORE_A_PATH, VALID_ENTITY_NAME, StorageKey::Int(7), 1);
476        insert_corrupted_data_key(STORE_A_PATH);
477
478        let report = diagnostics_report(&[(VALID_ENTITY_NAME, VALID_ENTITY_PATH)]);
479
480        assert_eq!(report.corrupted_keys, 1);
481        let entity_snapshot = report
482            .entity_storage
483            .iter()
484            .find(|snapshot| snapshot.store == STORE_A_PATH && snapshot.path == VALID_ENTITY_PATH)
485            .expect("valid-entity snapshot should exist");
486        assert_eq!(entity_snapshot.entries, 1);
487    }
488
489    #[test]
490    fn storage_report_corrupted_index_value_detection() {
491        reset_stores();
492
493        let key = index_key(IndexKeyKind::User, "diag_index_entity", "email");
494        let corrupted_entry = <RawIndexEntry as Storable>::from_bytes(Cow::Owned(vec![0, 0, 0, 0]));
495        insert_index_entry(STORE_A_PATH, key, corrupted_entry);
496
497        let report = diagnostics_report(&[]);
498        let index_snapshot = report
499            .storage_index
500            .iter()
501            .find(|snapshot| snapshot.path == STORE_A_PATH)
502            .expect("index snapshot should exist");
503
504        assert_eq!(report.corrupted_entries, 1);
505        assert_eq!(index_snapshot.entries, 1);
506        assert_eq!(index_snapshot.user_entries, 1);
507        assert_eq!(index_snapshot.system_entries, 0);
508    }
509
510    #[test]
511    fn storage_report_system_vs_user_namespace_split() {
512        reset_stores();
513
514        let user_key = index_key(IndexKeyKind::User, "diag_namespace_entity", "email");
515        let system_key = index_key(IndexKeyKind::System, "diag_namespace_entity", "email");
516        let user_entry =
517            RawIndexEntry::try_from_keys([StorageKey::Int(1)]).expect("user entry should encode");
518        let system_entry =
519            RawIndexEntry::try_from_keys([StorageKey::Int(2)]).expect("system entry should encode");
520        insert_index_entry(STORE_A_PATH, user_key, user_entry);
521        insert_index_entry(STORE_A_PATH, system_key, system_entry);
522
523        let report = diagnostics_report(&[]);
524        let index_snapshot = report
525            .storage_index
526            .iter()
527            .find(|snapshot| snapshot.path == STORE_A_PATH)
528            .expect("index snapshot should exist");
529
530        assert_eq!(report.corrupted_entries, 0);
531        assert_eq!(index_snapshot.entries, 2);
532        assert_eq!(index_snapshot.user_entries, 1);
533        assert_eq!(index_snapshot.system_entries, 1);
534    }
535}
icydb_core/db/diagnostics/snapshot/mod.rs

icydb_core/db/diagnostics/snapshot/
mod.rs
