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    /// Entity path (e.g., icydb_schema_tests::canister::db::Index)
64    pub path: String,
65    /// Number of rows for this entity in the store
66    pub entries: u64,
67    /// Approximate bytes used (key + value)
68    pub memory_bytes: u64,
69    /// Minimum primary key for this entity (entity-local ordering)
70    pub min_key: Option<Value>,
71    /// Maximum primary key for this entity (entity-local ordering)
72    pub max_key: Option<Value>,
73}
74
75///
76/// EntityStats
77/// Internal struct for building per-entity stats before snapshotting.
78///
79
80#[derive(Default)]
81struct EntityStats {
82    entries: u64,
83    memory_bytes: u64,
84    min_key: Option<StorageKey>,
85    max_key: Option<StorageKey>,
86}
87
88impl EntityStats {
89    fn update(&mut self, dk: &DataKey, value_len: u64) {
90        self.entries = self.entries.saturating_add(1);
91        self.memory_bytes = self
92            .memory_bytes
93            .saturating_add(DataKey::entry_size_bytes(value_len));
94
95        let k = dk.storage_key();
96
97        match &mut self.min_key {
98            Some(min) if k < *min => *min = k,
99            None => self.min_key = Some(k),
100            _ => {}
101        }
102
103        match &mut self.max_key {
104            Some(max) if k > *max => *max = k,
105            None => self.max_key = Some(k),
106            _ => {}
107        }
108    }
109}
110
111/// Build storage snapshot and per-entity breakdown; enrich path names using name→path map
112pub fn storage_report<C: CanisterKind>(
113    db: &Db<C>,
114    name_to_path: &[(&'static str, &'static str)],
115) -> Result<StorageReport, InternalError> {
116    ensure_recovered(db)?;
117    // Build name→path map once, reuse across stores
118    let name_map: BTreeMap<&'static str, &str> = name_to_path.iter().copied().collect();
119    let mut data = Vec::new();
120    let mut index = Vec::new();
121    let mut entity_storage: Vec<EntitySnapshot> = Vec::new();
122    let mut corrupted_keys = 0u64;
123    let mut corrupted_entries = 0u64;
124
125    db.with_data(|reg| {
126        reg.for_each(|path, store| {
127            data.push(DataStoreSnapshot {
128                path: path.to_string(),
129                entries: store.len(),
130                memory_bytes: store.memory_bytes(),
131            });
132
133            // Track per-entity counts, memory, and min/max Keys (not DataKeys)
134            let mut by_entity: BTreeMap<EntityName, EntityStats> = BTreeMap::new();
135
136            for entry in store.iter() {
137                let Ok(dk) = DataKey::try_from_raw(entry.key()) else {
138                    corrupted_keys = corrupted_keys.saturating_add(1);
139                    continue;
140                };
141
142                let value_len = entry.value().len() as u64;
143
144                by_entity
145                    .entry(*dk.entity_name())
146                    .or_default()
147                    .update(&dk, value_len);
148            }
149
150            for (entity_name, stats) in by_entity {
151                let path_name = name_map.get(entity_name.as_str()).copied().unwrap_or("");
152                entity_storage.push(EntitySnapshot {
153                    store: path.to_string(),
154                    path: path_name.to_string(),
155                    entries: stats.entries,
156                    memory_bytes: stats.memory_bytes,
157                    min_key: stats.min_key.map(|key| key.as_value()),
158                    max_key: stats.max_key.map(|key| key.as_value()),
159                });
160            }
161        });
162    });
163
164    db.with_index(|reg| {
165        reg.for_each(|path, store| {
166            index.push(IndexStoreSnapshot {
167                path: path.to_string(),
168                entries: store.len(),
169                memory_bytes: store.memory_bytes(),
170            });
171
172            for (key, value) in store.entries() {
173                if IndexKey::try_from_raw(&key).is_err() {
174                    corrupted_entries = corrupted_entries.saturating_add(1);
175                    continue;
176                }
177                if value.validate().is_err() {
178                    corrupted_entries = corrupted_entries.saturating_add(1);
179                }
180            }
181        });
182    });
183
184    Ok(StorageReport {
185        storage_data: data,
186        storage_index: index,
187        entity_storage,
188        corrupted_keys,
189        corrupted_entries,
190    })
191}
icydb_core/obs/snapshot/mod.rs

icydb_core/obs/snapshot/
mod.rs
