icydb_core/obs/snapshot/
mod.rs

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

icydb_core/obs/snapshot/
mod.rs
