icydb_core/obs/snapshot/
mod.rs

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

icydb_core/obs/snapshot/
mod.rs
