icydb_core/obs/snapshot/
mod.rs

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

icydb_core/obs/snapshot/
mod.rs
