zebra_state/service/finalized_state/disk_format/upgrade/
prune_trees.rs

1//! Prunes duplicate Sapling and Orchard note commitment trees from database
2
3use crossbeam_channel::{Receiver, TryRecvError};
4
5use semver::Version;
6use zebra_chain::block::Height;
7
8use crate::service::finalized_state::{DiskWriteBatch, ZebraDb};
9
10use super::{CancelFormatChange, DiskFormatUpgrade};
11
12/// Implements [`DiskFormatUpgrade`] for pruning duplicate Sapling and Orchard note commitment trees from database
13pub struct PruneTrees;
14
15impl DiskFormatUpgrade for PruneTrees {
16    fn version(&self) -> Version {
17        Version::new(25, 1, 1)
18    }
19
20    fn description(&self) -> &'static str {
21        "deduplicate trees upgrade"
22    }
23
24    #[allow(clippy::unwrap_in_result)]
25    fn run(
26        &self,
27        initial_tip_height: Height,
28        db: &ZebraDb,
29        cancel_receiver: &Receiver<CancelFormatChange>,
30    ) -> Result<(), CancelFormatChange> {
31        // Prune duplicate Sapling note commitment trees.
32
33        // The last tree we checked.
34        let mut last_tree = db
35            .sapling_tree_by_height(&Height(0))
36            .expect("Checked above that the genesis block is in the database.");
37
38        // Run through all the possible duplicate trees in the finalized chain.
39        // The block after genesis is the first possible duplicate.
40        for (height, tree) in db.sapling_tree_by_height_range(Height(1)..=initial_tip_height) {
41            // Return early if there is a cancel signal.
42            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
43                return Err(CancelFormatChange);
44            }
45
46            // Delete any duplicate trees.
47            if tree == last_tree {
48                let mut batch = DiskWriteBatch::new();
49                batch.delete_sapling_tree(db, &height);
50                db.write_batch(batch)
51                    .expect("Deleting Sapling note commitment trees should always succeed.");
52            }
53
54            // Compare against the last tree to find unique trees.
55            last_tree = tree;
56        }
57
58        // Prune duplicate Orchard note commitment trees.
59
60        // The last tree we checked.
61        let mut last_tree = db
62            .orchard_tree_by_height(&Height(0))
63            .expect("Checked above that the genesis block is in the database.");
64
65        // Run through all the possible duplicate trees in the finalized chain.
66        // The block after genesis is the first possible duplicate.
67        for (height, tree) in db.orchard_tree_by_height_range(Height(1)..=initial_tip_height) {
68            // Return early if there is a cancel signal.
69            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
70                return Err(CancelFormatChange);
71            }
72
73            // Delete any duplicate trees.
74            if tree == last_tree {
75                let mut batch = DiskWriteBatch::new();
76                batch.delete_orchard_tree(db, &height);
77                db.write_batch(batch)
78                    .expect("Deleting Orchard note commitment trees should always succeed.");
79            }
80
81            // Compare against the last tree to find unique trees.
82            last_tree = tree;
83        }
84
85        Ok(())
86    }
87
88    /// Check that note commitment trees were correctly de-duplicated.
89    #[allow(clippy::unwrap_in_result)]
90    fn validate(
91        &self,
92        db: &ZebraDb,
93        cancel_receiver: &Receiver<CancelFormatChange>,
94    ) -> Result<Result<(), String>, CancelFormatChange> {
95        // Runtime test: make sure we removed all duplicates.
96        // We always run this test, even if the state has supposedly been upgraded.
97        let mut result = Ok(());
98
99        let mut prev_height = None;
100        let mut prev_tree = None;
101        for (height, tree) in db.sapling_tree_by_height_range(..) {
102            // Return early if the format check is cancelled.
103            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
104                return Err(CancelFormatChange);
105            }
106
107            if prev_tree == Some(tree.clone()) {
108                result = Err(format!(
109                    "found duplicate sapling trees after running de-duplicate tree upgrade:\
110                     height: {height:?}, previous height: {:?}, tree root: {:?}",
111                    prev_height.unwrap(),
112                    tree.root()
113                ));
114                error!(?result);
115            }
116
117            prev_height = Some(height);
118            prev_tree = Some(tree);
119        }
120
121        let mut prev_height = None;
122        let mut prev_tree = None;
123        for (height, tree) in db.orchard_tree_by_height_range(..) {
124            // Return early if the format check is cancelled.
125            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
126                return Err(CancelFormatChange);
127            }
128
129            if prev_tree == Some(tree.clone()) {
130                result = Err(format!(
131                    "found duplicate orchard trees after running de-duplicate tree upgrade:\
132                     height: {height:?}, previous height: {:?}, tree root: {:?}",
133                    prev_height.unwrap(),
134                    tree.root()
135                ));
136                error!(?result);
137            }
138
139            prev_height = Some(height);
140            prev_tree = Some(tree);
141        }
142
143        Ok(result)
144    }
145}
zebra_state/service/finalized_state/disk_format/upgrade/prune_trees.rs

zebra_state/service/finalized_state/disk_format/upgrade/
prune_trees.rs
