icydb_core/obs/snapshot/
mod.rs

1use crate::{
2    db::{Db, identity::EntityName, index::IndexKey, store::DataKey},
3    key::Key,
4    traits::CanisterKind,
5};
6use candid::CandidType;
7use serde::{Deserialize, Serialize};
8use std::collections::BTreeMap;
9
10///
11/// StorageReport
12/// Live storage snapshot report
13///
14
15#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
16pub struct StorageReport {
17    pub storage_data: Vec<DataStoreSnapshot>,
18    pub storage_index: Vec<IndexStoreSnapshot>,
19    pub entity_storage: Vec<EntitySnapshot>,
20    pub corrupted_keys: u64,
21    pub corrupted_entries: u64,
22}
23
24///
25/// DataStoreSnapshot
26/// Store-level snapshot metrics.
27///
28
29#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
30pub struct DataStoreSnapshot {
31    pub path: String,
32    pub entries: u64,
33    pub memory_bytes: u64,
34}
35
36///
37/// IndexStoreSnapshot
38/// Index-store snapshot metrics
39///
40
41#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
42pub struct IndexStoreSnapshot {
43    pub path: String,
44    pub entries: u64,
45    pub memory_bytes: u64,
46}
47
48///
49/// EntitySnapshot
50/// Per-entity storage breakdown across stores
51///
52
53#[derive(CandidType, Clone, Debug, Default, Deserialize, Serialize)]
54pub struct EntitySnapshot {
55    /// Store path (e.g., test_design::schema::TestDataStore)
56    pub store: String,
57    /// Entity path (e.g., test_design::canister::db::Index)
58    pub path: String,
59    /// Number of rows for this entity in the store
60    pub entries: u64,
61    /// Approximate bytes used (key + value)
62    pub memory_bytes: u64,
63    /// Minimum primary key for this entity (entity-local ordering)
64    pub min_key: Option<Key>,
65    /// Maximum primary key for this entity (entity-local ordering)
66    pub max_key: Option<Key>,
67}
68
69///
70/// EntityStats
71/// Internal struct for building per-entity stats before snapshotting.
72///
73
74#[derive(Default)]
75struct EntityStats {
76    entries: u64,
77    memory_bytes: u64,
78    min_key: Option<Key>,
79    max_key: Option<Key>,
80}
81
82impl EntityStats {
83    fn update(&mut self, dk: &DataKey, value_len: u64) {
84        self.entries = self.entries.saturating_add(1);
85        self.memory_bytes = self
86            .memory_bytes
87            .saturating_add(DataKey::entry_size_bytes(value_len));
88
89        let k = dk.key();
90
91        match &mut self.min_key {
92            Some(min) if k < *min => *min = k,
93            None => self.min_key = Some(k),
94            _ => {}
95        }
96
97        match &mut self.max_key {
98            Some(max) if k > *max => *max = k,
99            None => self.max_key = Some(k),
100            _ => {}
101        }
102    }
103}
104
105/// Build storage snapshot and per-entity breakdown; enrich path names using name→path map
106#[must_use]
107pub fn storage_report<C: CanisterKind>(
108    db: &Db<C>,
109    name_to_path: &[(&'static str, &'static str)],
110) -> StorageReport {
111    // Build name→path map once, reuse across stores
112    let name_map: BTreeMap<&'static str, &str> = name_to_path.iter().copied().collect();
113    let mut data = Vec::new();
114    let mut index = Vec::new();
115    let mut entity_storage: Vec<EntitySnapshot> = Vec::new();
116    let mut corrupted_keys = 0u64;
117    let mut corrupted_entries = 0u64;
118
119    db.with_data(|reg| {
120        reg.for_each(|path, store| {
121            data.push(DataStoreSnapshot {
122                path: path.to_string(),
123                entries: store.len(),
124                memory_bytes: store.memory_bytes(),
125            });
126
127            // Track per-entity counts, memory, and min/max Keys (not DataKeys)
128            let mut by_entity: BTreeMap<EntityName, EntityStats> = BTreeMap::new();
129
130            for entry in store.iter() {
131                let Ok(dk) = DataKey::try_from_raw(entry.key()) else {
132                    corrupted_keys = corrupted_keys.saturating_add(1);
133                    continue;
134                };
135
136                let value_len = entry.value().len() as u64;
137
138                by_entity
139                    .entry(*dk.entity_name())
140                    .or_default()
141                    .update(&dk, value_len);
142            }
143
144            for (entity_name, stats) in by_entity {
145                let path_name = name_map.get(entity_name.as_str()).copied().unwrap_or("");
146                entity_storage.push(EntitySnapshot {
147                    store: path.to_string(),
148                    path: path_name.to_string(),
149                    entries: stats.entries,
150                    memory_bytes: stats.memory_bytes,
151                    min_key: stats.min_key,
152                    max_key: stats.max_key,
153                });
154            }
155        });
156    });
157
158    db.with_index(|reg| {
159        reg.for_each(|path, store| {
160            index.push(IndexStoreSnapshot {
161                path: path.to_string(),
162                entries: store.len(),
163                memory_bytes: store.memory_bytes(),
164            });
165
166            for entry in store.iter() {
167                if IndexKey::try_from_raw(entry.key()).is_err() {
168                    corrupted_entries = corrupted_entries.saturating_add(1);
169                    continue;
170                }
171                if entry.value().try_decode().is_err() {
172                    corrupted_entries = corrupted_entries.saturating_add(1);
173                }
174            }
175        });
176    });
177
178    StorageReport {
179        storage_data: data,
180        storage_index: index,
181        entity_storage,
182        corrupted_keys,
183        corrupted_entries,
184    }
185}
icydb_core/obs/snapshot/mod.rs

icydb_core/obs/snapshot/
mod.rs
