icydb_core/obs/snapshot/
mod.rs

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