Skip to main content

icydb_core/db/diagnostics/
mod.rs

1//! Module: diagnostics
2//! Responsibility: read-only storage footprint and integrity snapshots.
3//! Does not own: recovery, write-path mutation, or query planning semantics.
4//! Boundary: consumes `Db`/store read APIs and returns DTO snapshots.
5
6mod execution_trace;
7#[cfg(test)]
8mod tests;
9
10use crate::{
11    db::{
12        Db, EntityRuntimeHooks,
13        commit::CommitRowOp,
14        data::{DataKey, StorageKey, decode_structural_row_cbor},
15        index::IndexKey,
16        registry::StoreHandle,
17    },
18    error::{ErrorClass, InternalError},
19    traits::{CanisterKind, Repr},
20    types::EntityTag,
21};
22use candid::CandidType;
23use serde::Deserialize;
24use std::collections::{BTreeMap, BTreeSet};
25
26pub use execution_trace::{
27    ExecutionAccessPathVariant, ExecutionMetrics, ExecutionOptimization, ExecutionTrace,
28};
29
30#[cfg_attr(doc, doc = "StorageReport\n\nLive storage snapshot payload.")]
31#[derive(CandidType, Clone, Debug, Default, Deserialize)]
32pub struct StorageReport {
33    pub(crate) storage_data: Vec<DataStoreSnapshot>,
34    pub(crate) storage_index: Vec<IndexStoreSnapshot>,
35    pub(crate) entity_storage: Vec<EntitySnapshot>,
36    pub(crate) corrupted_keys: u64,
37    pub(crate) corrupted_entries: u64,
38}
39
40#[cfg_attr(
41    doc,
42    doc = "IntegrityTotals\n\nAggregated integrity-scan counters across all stores."
43)]
44#[derive(CandidType, Clone, Debug, Default, Deserialize)]
45pub struct IntegrityTotals {
46    pub(crate) data_rows_scanned: u64,
47    pub(crate) index_entries_scanned: u64,
48    pub(crate) corrupted_data_keys: u64,
49    pub(crate) corrupted_data_rows: u64,
50    pub(crate) corrupted_index_keys: u64,
51    pub(crate) corrupted_index_entries: u64,
52    pub(crate) missing_index_entries: u64,
53    pub(crate) divergent_index_entries: u64,
54    pub(crate) orphan_index_references: u64,
55    pub(crate) compatibility_findings: u64,
56    pub(crate) misuse_findings: u64,
57}
58
59impl IntegrityTotals {
60    const fn add_store_snapshot(&mut self, store: &IntegrityStoreSnapshot) {
61        self.data_rows_scanned = self
62            .data_rows_scanned
63            .saturating_add(store.data_rows_scanned);
64        self.index_entries_scanned = self
65            .index_entries_scanned
66            .saturating_add(store.index_entries_scanned);
67        self.corrupted_data_keys = self
68            .corrupted_data_keys
69            .saturating_add(store.corrupted_data_keys);
70        self.corrupted_data_rows = self
71            .corrupted_data_rows
72            .saturating_add(store.corrupted_data_rows);
73        self.corrupted_index_keys = self
74            .corrupted_index_keys
75            .saturating_add(store.corrupted_index_keys);
76        self.corrupted_index_entries = self
77            .corrupted_index_entries
78            .saturating_add(store.corrupted_index_entries);
79        self.missing_index_entries = self
80            .missing_index_entries
81            .saturating_add(store.missing_index_entries);
82        self.divergent_index_entries = self
83            .divergent_index_entries
84            .saturating_add(store.divergent_index_entries);
85        self.orphan_index_references = self
86            .orphan_index_references
87            .saturating_add(store.orphan_index_references);
88        self.compatibility_findings = self
89            .compatibility_findings
90            .saturating_add(store.compatibility_findings);
91        self.misuse_findings = self.misuse_findings.saturating_add(store.misuse_findings);
92    }
93
94    /// Return total number of data rows scanned.
95    #[must_use]
96    pub const fn data_rows_scanned(&self) -> u64 {
97        self.data_rows_scanned
98    }
99
100    /// Return total number of index entries scanned.
101    #[must_use]
102    pub const fn index_entries_scanned(&self) -> u64 {
103        self.index_entries_scanned
104    }
105
106    /// Return total number of corrupted data-key findings.
107    #[must_use]
108    pub const fn corrupted_data_keys(&self) -> u64 {
109        self.corrupted_data_keys
110    }
111
112    /// Return total number of corrupted data-row findings.
113    #[must_use]
114    pub const fn corrupted_data_rows(&self) -> u64 {
115        self.corrupted_data_rows
116    }
117
118    /// Return total number of corrupted index-key findings.
119    #[must_use]
120    pub const fn corrupted_index_keys(&self) -> u64 {
121        self.corrupted_index_keys
122    }
123
124    /// Return total number of corrupted index-entry findings.
125    #[must_use]
126    pub const fn corrupted_index_entries(&self) -> u64 {
127        self.corrupted_index_entries
128    }
129
130    /// Return total number of missing index-entry findings.
131    #[must_use]
132    pub const fn missing_index_entries(&self) -> u64 {
133        self.missing_index_entries
134    }
135
136    /// Return total number of divergent index-entry findings.
137    #[must_use]
138    pub const fn divergent_index_entries(&self) -> u64 {
139        self.divergent_index_entries
140    }
141
142    /// Return total number of orphan index-reference findings.
143    #[must_use]
144    pub const fn orphan_index_references(&self) -> u64 {
145        self.orphan_index_references
146    }
147
148    /// Return total number of compatibility findings.
149    #[must_use]
150    pub const fn compatibility_findings(&self) -> u64 {
151        self.compatibility_findings
152    }
153
154    /// Return total number of misuse findings.
155    #[must_use]
156    pub const fn misuse_findings(&self) -> u64 {
157        self.misuse_findings
158    }
159}
160
161#[cfg_attr(
162    doc,
163    doc = "IntegrityStoreSnapshot\n\nPer-store integrity findings and scan counters."
164)]
165#[derive(CandidType, Clone, Debug, Default, Deserialize)]
166pub struct IntegrityStoreSnapshot {
167    pub(crate) path: String,
168    pub(crate) data_rows_scanned: u64,
169    pub(crate) index_entries_scanned: u64,
170    pub(crate) corrupted_data_keys: u64,
171    pub(crate) corrupted_data_rows: u64,
172    pub(crate) corrupted_index_keys: u64,
173    pub(crate) corrupted_index_entries: u64,
174    pub(crate) missing_index_entries: u64,
175    pub(crate) divergent_index_entries: u64,
176    pub(crate) orphan_index_references: u64,
177    pub(crate) compatibility_findings: u64,
178    pub(crate) misuse_findings: u64,
179}
180
181impl IntegrityStoreSnapshot {
182    /// Construct one empty store-level integrity snapshot.
183    #[must_use]
184    pub fn new(path: String) -> Self {
185        Self {
186            path,
187            ..Self::default()
188        }
189    }
190
191    /// Borrow store path.
192    #[must_use]
193    pub const fn path(&self) -> &str {
194        self.path.as_str()
195    }
196
197    /// Return number of scanned data rows.
198    #[must_use]
199    pub const fn data_rows_scanned(&self) -> u64 {
200        self.data_rows_scanned
201    }
202
203    /// Return number of scanned index entries.
204    #[must_use]
205    pub const fn index_entries_scanned(&self) -> u64 {
206        self.index_entries_scanned
207    }
208
209    /// Return number of corrupted data-key findings.
210    #[must_use]
211    pub const fn corrupted_data_keys(&self) -> u64 {
212        self.corrupted_data_keys
213    }
214
215    /// Return number of corrupted data-row findings.
216    #[must_use]
217    pub const fn corrupted_data_rows(&self) -> u64 {
218        self.corrupted_data_rows
219    }
220
221    /// Return number of corrupted index-key findings.
222    #[must_use]
223    pub const fn corrupted_index_keys(&self) -> u64 {
224        self.corrupted_index_keys
225    }
226
227    /// Return number of corrupted index-entry findings.
228    #[must_use]
229    pub const fn corrupted_index_entries(&self) -> u64 {
230        self.corrupted_index_entries
231    }
232
233    /// Return number of missing index-entry findings.
234    #[must_use]
235    pub const fn missing_index_entries(&self) -> u64 {
236        self.missing_index_entries
237    }
238
239    /// Return number of divergent index-entry findings.
240    #[must_use]
241    pub const fn divergent_index_entries(&self) -> u64 {
242        self.divergent_index_entries
243    }
244
245    /// Return number of orphan index-reference findings.
246    #[must_use]
247    pub const fn orphan_index_references(&self) -> u64 {
248        self.orphan_index_references
249    }
250
251    /// Return number of compatibility findings.
252    #[must_use]
253    pub const fn compatibility_findings(&self) -> u64 {
254        self.compatibility_findings
255    }
256
257    /// Return number of misuse findings.
258    #[must_use]
259    pub const fn misuse_findings(&self) -> u64 {
260        self.misuse_findings
261    }
262}
263
264#[cfg_attr(
265    doc,
266    doc = "IntegrityReport\n\nFull integrity-scan output across all registered stores."
267)]
268#[derive(CandidType, Clone, Debug, Default, Deserialize)]
269pub struct IntegrityReport {
270    pub(crate) stores: Vec<IntegrityStoreSnapshot>,
271    pub(crate) totals: IntegrityTotals,
272}
273
274impl IntegrityReport {
275    /// Construct one integrity report payload.
276    #[must_use]
277    pub const fn new(stores: Vec<IntegrityStoreSnapshot>, totals: IntegrityTotals) -> Self {
278        Self { stores, totals }
279    }
280
281    /// Borrow per-store integrity snapshots.
282    #[must_use]
283    pub const fn stores(&self) -> &[IntegrityStoreSnapshot] {
284        self.stores.as_slice()
285    }
286
287    /// Borrow aggregated integrity totals.
288    #[must_use]
289    pub const fn totals(&self) -> &IntegrityTotals {
290        &self.totals
291    }
292}
293
294impl StorageReport {
295    /// Construct one storage report payload.
296    #[must_use]
297    pub const fn new(
298        storage_data: Vec<DataStoreSnapshot>,
299        storage_index: Vec<IndexStoreSnapshot>,
300        entity_storage: Vec<EntitySnapshot>,
301        corrupted_keys: u64,
302        corrupted_entries: u64,
303    ) -> Self {
304        Self {
305            storage_data,
306            storage_index,
307            entity_storage,
308            corrupted_keys,
309            corrupted_entries,
310        }
311    }
312
313    /// Borrow data-store snapshots.
314    #[must_use]
315    pub const fn storage_data(&self) -> &[DataStoreSnapshot] {
316        self.storage_data.as_slice()
317    }
318
319    /// Borrow index-store snapshots.
320    #[must_use]
321    pub const fn storage_index(&self) -> &[IndexStoreSnapshot] {
322        self.storage_index.as_slice()
323    }
324
325    /// Borrow entity-level storage snapshots.
326    #[must_use]
327    pub const fn entity_storage(&self) -> &[EntitySnapshot] {
328        self.entity_storage.as_slice()
329    }
330
331    /// Return count of corrupted decoded data keys.
332    #[must_use]
333    pub const fn corrupted_keys(&self) -> u64 {
334        self.corrupted_keys
335    }
336
337    /// Return count of corrupted index entries.
338    #[must_use]
339    pub const fn corrupted_entries(&self) -> u64 {
340        self.corrupted_entries
341    }
342}
343
344#[cfg_attr(doc, doc = "DataStoreSnapshot\n\nData-store snapshot row.")]
345#[derive(CandidType, Clone, Debug, Default, Deserialize)]
346pub struct DataStoreSnapshot {
347    pub(crate) path: String,
348    pub(crate) entries: u64,
349    pub(crate) memory_bytes: u64,
350}
351
352impl DataStoreSnapshot {
353    /// Construct one data-store snapshot row.
354    #[must_use]
355    pub const fn new(path: String, entries: u64, memory_bytes: u64) -> Self {
356        Self {
357            path,
358            entries,
359            memory_bytes,
360        }
361    }
362
363    /// Borrow store path.
364    #[must_use]
365    pub const fn path(&self) -> &str {
366        self.path.as_str()
367    }
368
369    /// Return row count.
370    #[must_use]
371    pub const fn entries(&self) -> u64 {
372        self.entries
373    }
374
375    /// Return memory usage in bytes.
376    #[must_use]
377    pub const fn memory_bytes(&self) -> u64 {
378        self.memory_bytes
379    }
380}
381
382#[cfg_attr(doc, doc = "IndexStoreSnapshot\n\nIndex-store snapshot row.")]
383#[derive(CandidType, Clone, Debug, Default, Deserialize)]
384pub struct IndexStoreSnapshot {
385    pub(crate) path: String,
386    pub(crate) entries: u64,
387    pub(crate) user_entries: u64,
388    pub(crate) system_entries: u64,
389    pub(crate) memory_bytes: u64,
390}
391
392impl IndexStoreSnapshot {
393    /// Construct one index-store snapshot row.
394    #[must_use]
395    pub const fn new(
396        path: String,
397        entries: u64,
398        user_entries: u64,
399        system_entries: u64,
400        memory_bytes: u64,
401    ) -> Self {
402        Self {
403            path,
404            entries,
405            user_entries,
406            system_entries,
407            memory_bytes,
408        }
409    }
410
411    /// Borrow store path.
412    #[must_use]
413    pub const fn path(&self) -> &str {
414        self.path.as_str()
415    }
416
417    /// Return total entry count.
418    #[must_use]
419    pub const fn entries(&self) -> u64 {
420        self.entries
421    }
422
423    /// Return user-namespace entry count.
424    #[must_use]
425    pub const fn user_entries(&self) -> u64 {
426        self.user_entries
427    }
428
429    /// Return system-namespace entry count.
430    #[must_use]
431    pub const fn system_entries(&self) -> u64 {
432        self.system_entries
433    }
434
435    /// Return memory usage in bytes.
436    #[must_use]
437    pub const fn memory_bytes(&self) -> u64 {
438        self.memory_bytes
439    }
440}
441
442#[cfg_attr(doc, doc = "EntitySnapshot\n\nPer-entity storage snapshot row.")]
443#[derive(CandidType, Clone, Debug, Default, Deserialize)]
444pub struct EntitySnapshot {
445    pub(crate) store: String,
446
447    pub(crate) path: String,
448
449    pub(crate) entries: u64,
450
451    pub(crate) memory_bytes: u64,
452
453    pub(crate) min_key: Option<String>,
454
455    pub(crate) max_key: Option<String>,
456}
457
458impl EntitySnapshot {
459    /// Construct one entity-storage snapshot row.
460    #[must_use]
461    pub fn new(
462        store: String,
463        path: String,
464        entries: u64,
465        memory_bytes: u64,
466        min_key: Option<StorageKey>,
467        max_key: Option<StorageKey>,
468    ) -> Self {
469        Self {
470            store,
471            path,
472            entries,
473            memory_bytes,
474            min_key: min_key.map(Self::storage_key_text),
475            max_key: max_key.map(Self::storage_key_text),
476        }
477    }
478
479    // Keep snapshot key rendering local to the diagnostics contract so the
480    // canister DTO does not retain the full `Value` Candid surface.
481    fn storage_key_text(key: StorageKey) -> String {
482        match key {
483            StorageKey::Account(value) => value.to_string(),
484            StorageKey::Int(value) => value.to_string(),
485            StorageKey::Principal(value) => value.to_string(),
486            StorageKey::Subaccount(value) => value.to_string(),
487            StorageKey::Timestamp(value) => value.repr().to_string(),
488            StorageKey::Uint(value) => value.to_string(),
489            StorageKey::Ulid(value) => value.to_string(),
490            StorageKey::Unit => "()".to_string(),
491        }
492    }
493
494    /// Borrow store path.
495    #[must_use]
496    pub const fn store(&self) -> &str {
497        self.store.as_str()
498    }
499
500    /// Borrow entity path.
501    #[must_use]
502    pub const fn path(&self) -> &str {
503        self.path.as_str()
504    }
505
506    /// Return row count.
507    #[must_use]
508    pub const fn entries(&self) -> u64 {
509        self.entries
510    }
511
512    /// Return memory usage in bytes.
513    #[must_use]
514    pub const fn memory_bytes(&self) -> u64 {
515        self.memory_bytes
516    }
517
518    /// Borrow optional minimum primary key.
519    #[must_use]
520    pub fn min_key(&self) -> Option<&str> {
521        self.min_key.as_deref()
522    }
523
524    /// Borrow optional maximum primary key.
525    #[must_use]
526    pub fn max_key(&self) -> Option<&str> {
527        self.max_key.as_deref()
528    }
529}
530
531#[cfg_attr(
532    doc,
533    doc = "EntityStats\n\nInternal struct for building per-entity stats before snapshotting."
534)]
535#[derive(Default)]
536struct EntityStats {
537    entries: u64,
538    memory_bytes: u64,
539    min_key: Option<StorageKey>,
540    max_key: Option<StorageKey>,
541}
542
543impl EntityStats {
544    // Accumulate per-entity counters and keep min/max over entity-local storage keys.
545    fn update(&mut self, dk: &DataKey, value_len: u64) {
546        self.entries = self.entries.saturating_add(1);
547        self.memory_bytes = self
548            .memory_bytes
549            .saturating_add(DataKey::entry_size_bytes(value_len));
550
551        let k = dk.storage_key();
552
553        match &mut self.min_key {
554            Some(min) if k < *min => *min = k,
555            None => self.min_key = Some(k),
556            _ => {}
557        }
558
559        match &mut self.max_key {
560            Some(max) if k > *max => *max = k,
561            None => self.max_key = Some(k),
562            _ => {}
563        }
564    }
565}
566
567fn storage_report_name_for_hook<'a, C: CanisterKind>(
568    name_map: &BTreeMap<&'static str, &'a str>,
569    hooks: &EntityRuntimeHooks<C>,
570) -> &'a str {
571    name_map
572        .get(hooks.entity_path)
573        .copied()
574        .or_else(|| name_map.get(hooks.model.name()).copied())
575        .unwrap_or(hooks.entity_path)
576}
577
578#[cfg_attr(
579    doc,
580    doc = "Build one deterministic storage snapshot with per-entity rollups.\n\nThis path is read-only and fail-closed on decode/validation errors by counting corrupted keys/entries instead of panicking."
581)]
582pub(crate) fn storage_report<C: CanisterKind>(
583    db: &Db<C>,
584    name_to_path: &[(&'static str, &'static str)],
585) -> Result<StorageReport, InternalError> {
586    db.ensure_recovered_state()?;
587    // Build one optional alias map once, then resolve report names from the
588    // runtime hook table so entity tags keep distinct path identity even when
589    // multiple hooks intentionally share the same model name.
590    let name_map: BTreeMap<&'static str, &str> = name_to_path.iter().copied().collect();
591    let mut tag_name_map = BTreeMap::<EntityTag, &str>::new();
592    for hooks in db.entity_runtime_hooks {
593        tag_name_map
594            .entry(hooks.entity_tag)
595            .or_insert_with(|| storage_report_name_for_hook(&name_map, hooks));
596    }
597    let mut data = Vec::new();
598    let mut index = Vec::new();
599    let mut entity_storage: Vec<EntitySnapshot> = Vec::new();
600    let mut corrupted_keys = 0u64;
601    let mut corrupted_entries = 0u64;
602
603    db.with_store_registry(|reg| {
604        // Keep diagnostics snapshots deterministic by traversing stores in path order.
605        let mut stores = reg.iter().collect::<Vec<_>>();
606        stores.sort_by_key(|(path, _)| *path);
607
608        for (path, store_handle) in stores {
609            // Phase 1: collect data-store snapshots and per-entity stats.
610            store_handle.with_data(|store| {
611                data.push(DataStoreSnapshot::new(
612                    path.to_string(),
613                    store.len(),
614                    store.memory_bytes(),
615                ));
616
617                // Track per-entity counts, memory, and min/max Keys (not DataKeys)
618                let mut by_entity: BTreeMap<EntityTag, EntityStats> = BTreeMap::new();
619
620                for entry in store.iter() {
621                    let Ok(dk) = DataKey::try_from_raw(entry.key()) else {
622                        corrupted_keys = corrupted_keys.saturating_add(1);
623                        continue;
624                    };
625
626                    let value_len = entry.value().len() as u64;
627
628                    by_entity
629                        .entry(dk.entity_tag())
630                        .or_default()
631                        .update(&dk, value_len);
632                }
633
634                for (entity_tag, stats) in by_entity {
635                    let path_name = tag_name_map
636                        .get(&entity_tag)
637                        .copied()
638                        .map(str::to_string)
639                        .or_else(|| {
640                            db.runtime_hook_for_entity_tag(entity_tag)
641                                .ok()
642                                .map(|hooks| {
643                                    storage_report_name_for_hook(&name_map, hooks).to_string()
644                                })
645                        })
646                        .unwrap_or_else(|| format!("#{}", entity_tag.value()));
647                    entity_storage.push(EntitySnapshot::new(
648                        path.to_string(),
649                        path_name,
650                        stats.entries,
651                        stats.memory_bytes,
652                        stats.min_key,
653                        stats.max_key,
654                    ));
655                }
656            });
657
658            // Phase 2: collect index-store snapshots and integrity counters.
659            store_handle.with_index(|store| {
660                let mut user_entries = 0u64;
661                let mut system_entries = 0u64;
662
663                for (key, value) in store.entries() {
664                    let Ok(decoded_key) = IndexKey::try_from_raw(&key) else {
665                        corrupted_entries = corrupted_entries.saturating_add(1);
666                        continue;
667                    };
668
669                    if decoded_key.uses_system_namespace() {
670                        system_entries = system_entries.saturating_add(1);
671                    } else {
672                        user_entries = user_entries.saturating_add(1);
673                    }
674
675                    if value.validate().is_err() {
676                        corrupted_entries = corrupted_entries.saturating_add(1);
677                    }
678                }
679
680                index.push(IndexStoreSnapshot::new(
681                    path.to_string(),
682                    store.len(),
683                    user_entries,
684                    system_entries,
685                    store.memory_bytes(),
686                ));
687            });
688        }
689    });
690
691    // Phase 3: enforce deterministic entity snapshot emission order.
692    // This remains stable even if outer store traversal internals change.
693    entity_storage
694        .sort_by(|left, right| (left.store(), left.path()).cmp(&(right.store(), right.path())));
695
696    Ok(StorageReport::new(
697        data,
698        index,
699        entity_storage,
700        corrupted_keys,
701        corrupted_entries,
702    ))
703}
704
705#[cfg_attr(
706    doc,
707    doc = "Build one deterministic integrity scan over all registered stores.\n\nThis scan is read-only and classifies findings as:\n- corruption: malformed persisted bytes or inconsistent structural links\n- compatibility: persisted payloads outside decode compatibility windows\n- misuse: unsupported runtime wiring (for example missing entity hooks)"
708)]
709pub(crate) fn integrity_report<C: CanisterKind>(
710    db: &Db<C>,
711) -> Result<IntegrityReport, InternalError> {
712    db.ensure_recovered_state()?;
713
714    integrity_report_after_recovery(db)
715}
716
717#[cfg_attr(
718    doc,
719    doc = "Build one deterministic integrity scan after recovery has already completed.\n\nCallers running inside recovery flow should use this variant to avoid recursive recovery gating."
720)]
721pub(in crate::db) fn integrity_report_after_recovery<C: CanisterKind>(
722    db: &Db<C>,
723) -> Result<IntegrityReport, InternalError> {
724    build_integrity_report(db)
725}
726
727fn build_integrity_report<C: CanisterKind>(db: &Db<C>) -> Result<IntegrityReport, InternalError> {
728    let mut stores = Vec::new();
729    let mut totals = IntegrityTotals::default();
730    let global_live_keys_by_entity = collect_global_live_keys_by_entity(db)?;
731
732    db.with_store_registry(|reg| {
733        // Keep deterministic output order across registry traversal implementations.
734        let mut store_entries = reg.iter().collect::<Vec<_>>();
735        store_entries.sort_by_key(|(path, _)| *path);
736
737        for (path, store_handle) in store_entries {
738            let mut snapshot = IntegrityStoreSnapshot::new(path.to_string());
739            scan_store_forward_integrity(db, store_handle, &mut snapshot)?;
740            scan_store_reverse_integrity(store_handle, &global_live_keys_by_entity, &mut snapshot);
741
742            totals.add_store_snapshot(&snapshot);
743            stores.push(snapshot);
744        }
745
746        Ok::<(), InternalError>(())
747    })?;
748
749    Ok(IntegrityReport::new(stores, totals))
750}
751
752// Build one global map of live data keys grouped by entity across all stores.
753fn collect_global_live_keys_by_entity<C: CanisterKind>(
754    db: &Db<C>,
755) -> Result<BTreeMap<EntityTag, BTreeSet<StorageKey>>, InternalError> {
756    let mut keys = BTreeMap::<EntityTag, BTreeSet<StorageKey>>::new();
757
758    db.with_store_registry(|reg| {
759        for (_, store_handle) in reg.iter() {
760            store_handle.with_data(|data_store| {
761                for entry in data_store.iter() {
762                    if let Ok(data_key) = DataKey::try_from_raw(entry.key()) {
763                        keys.entry(data_key.entity_tag())
764                            .or_default()
765                            .insert(data_key.storage_key());
766                    }
767                }
768            });
769        }
770
771        Ok::<(), InternalError>(())
772    })?;
773
774    Ok(keys)
775}
776
777// Run forward (data -> index) integrity checks for one store.
778fn scan_store_forward_integrity<C: CanisterKind>(
779    db: &Db<C>,
780    store_handle: StoreHandle,
781    snapshot: &mut IntegrityStoreSnapshot,
782) -> Result<(), InternalError> {
783    store_handle.with_data(|data_store| {
784        for entry in data_store.iter() {
785            snapshot.data_rows_scanned = snapshot.data_rows_scanned.saturating_add(1);
786
787            let raw_key = *entry.key();
788
789            let Ok(data_key) = DataKey::try_from_raw(&raw_key) else {
790                snapshot.corrupted_data_keys = snapshot.corrupted_data_keys.saturating_add(1);
791                continue;
792            };
793
794            let hooks = match db.runtime_hook_for_entity_tag(data_key.entity_tag()) {
795                Ok(hooks) => hooks,
796                Err(err) => {
797                    classify_scan_error(err, snapshot)?;
798                    continue;
799                }
800            };
801
802            let marker_row = CommitRowOp::new(
803                hooks.entity_path,
804                raw_key.as_bytes().to_vec(),
805                None,
806                Some(entry.value().as_bytes().to_vec()),
807                crate::db::schema::commit_schema_fingerprint_for_model(
808                    hooks.entity_path,
809                    hooks.model,
810                ),
811            );
812
813            // Validate envelope compatibility before typed preparation so
814            // incompatible persisted formats remain compatibility-classified.
815            if let Err(err) = decode_structural_row_cbor(&entry.value()) {
816                classify_scan_error(err, snapshot)?;
817                continue;
818            }
819
820            let prepared = match db.prepare_row_commit_op(&marker_row) {
821                Ok(prepared) => prepared,
822                Err(err) => {
823                    classify_scan_error(err, snapshot)?;
824                    continue;
825                }
826            };
827
828            for index_op in prepared.index_ops {
829                let Some(expected_value) = index_op.value else {
830                    continue;
831                };
832
833                let actual = index_op
834                    .store
835                    .with_borrow(|index_store| index_store.get(&index_op.key));
836                match actual {
837                    Some(actual_value) if actual_value == expected_value => {}
838                    Some(_) => {
839                        snapshot.divergent_index_entries =
840                            snapshot.divergent_index_entries.saturating_add(1);
841                    }
842                    None => {
843                        snapshot.missing_index_entries =
844                            snapshot.missing_index_entries.saturating_add(1);
845                    }
846                }
847            }
848        }
849
850        Ok::<(), InternalError>(())
851    })
852}
853
854// Run reverse (index -> data) integrity checks for one store.
855fn scan_store_reverse_integrity(
856    store_handle: StoreHandle,
857    live_keys_by_entity: &BTreeMap<EntityTag, BTreeSet<StorageKey>>,
858    snapshot: &mut IntegrityStoreSnapshot,
859) {
860    store_handle.with_index(|index_store| {
861        for (raw_index_key, raw_index_entry) in index_store.entries() {
862            snapshot.index_entries_scanned = snapshot.index_entries_scanned.saturating_add(1);
863
864            let Ok(decoded_index_key) = IndexKey::try_from_raw(&raw_index_key) else {
865                snapshot.corrupted_index_keys = snapshot.corrupted_index_keys.saturating_add(1);
866                continue;
867            };
868
869            let index_entity_tag = data_entity_tag_for_index_key(&decoded_index_key);
870
871            let Ok(indexed_primary_keys) = raw_index_entry.decode_keys() else {
872                snapshot.corrupted_index_entries =
873                    snapshot.corrupted_index_entries.saturating_add(1);
874                continue;
875            };
876
877            for primary_key in indexed_primary_keys {
878                let exists = live_keys_by_entity
879                    .get(&index_entity_tag)
880                    .is_some_and(|entity_keys| entity_keys.contains(&primary_key));
881                if !exists {
882                    snapshot.orphan_index_references =
883                        snapshot.orphan_index_references.saturating_add(1);
884                }
885            }
886        }
887    });
888}
889
890// Map scan-time errors into explicit integrity classification buckets.
891fn classify_scan_error(
892    err: InternalError,
893    snapshot: &mut IntegrityStoreSnapshot,
894) -> Result<(), InternalError> {
895    match err.class() {
896        ErrorClass::Corruption => {
897            snapshot.corrupted_data_rows = snapshot.corrupted_data_rows.saturating_add(1);
898            Ok(())
899        }
900        ErrorClass::IncompatiblePersistedFormat => {
901            snapshot.compatibility_findings = snapshot.compatibility_findings.saturating_add(1);
902            Ok(())
903        }
904        ErrorClass::Unsupported | ErrorClass::NotFound | ErrorClass::Conflict => {
905            snapshot.misuse_findings = snapshot.misuse_findings.saturating_add(1);
906            Ok(())
907        }
908        ErrorClass::Internal | ErrorClass::InvariantViolation => Err(err),
909    }
910}
911
912// Parse the data-entity identity from one decoded index key.
913const fn data_entity_tag_for_index_key(index_key: &IndexKey) -> EntityTag {
914    index_key.index_id().entity_tag
915}