1//! Prunes duplicate Sapling and Orchard note commitment trees from database
23use crossbeam_channel::{Receiver, TryRecvError};
45use semver::Version;
6use zebra_chain::block::Height;
78use crate::service::finalized_state::{DiskWriteBatch, ZebraDb};
910use super::{CancelFormatChange, DiskFormatUpgrade};
1112/// Implements [`DiskFormatUpgrade`] for pruning duplicate Sapling and Orchard note commitment trees from database
13pub struct PruneTrees;
1415impl DiskFormatUpgrade for PruneTrees {
16fn version(&self) -> Version {
17 Version::new(25, 1, 1)
18 }
1920fn description(&self) -> &'static str {
21"deduplicate trees upgrade"
22}
2324#[allow(clippy::unwrap_in_result)]
25fn 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.
3233 // The last tree we checked.
34let mut last_tree = db
35 .sapling_tree_by_height(&Height(0))
36 .expect("Checked above that the genesis block is in the database.");
3738// Run through all the possible duplicate trees in the finalized chain.
39 // The block after genesis is the first possible duplicate.
40for (height, tree) in db.sapling_tree_by_height_range(Height(1)..=initial_tip_height) {
41// Return early if there is a cancel signal.
42if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
43return Err(CancelFormatChange);
44 }
4546// Delete any duplicate trees.
47if tree == last_tree {
48let 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 }
5354// Compare against the last tree to find unique trees.
55last_tree = tree;
56 }
5758// Prune duplicate Orchard note commitment trees.
5960 // The last tree we checked.
61let mut last_tree = db
62 .orchard_tree_by_height(&Height(0))
63 .expect("Checked above that the genesis block is in the database.");
6465// Run through all the possible duplicate trees in the finalized chain.
66 // The block after genesis is the first possible duplicate.
67for (height, tree) in db.orchard_tree_by_height_range(Height(1)..=initial_tip_height) {
68// Return early if there is a cancel signal.
69if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
70return Err(CancelFormatChange);
71 }
7273// Delete any duplicate trees.
74if tree == last_tree {
75let 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 }
8081// Compare against the last tree to find unique trees.
82last_tree = tree;
83 }
8485Ok(())
86 }
8788/// Check that note commitment trees were correctly de-duplicated.
89#[allow(clippy::unwrap_in_result)]
90fn 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.
97let mut result = Ok(());
9899let mut prev_height = None;
100let mut prev_tree = None;
101for (height, tree) in db.sapling_tree_by_height_range(..) {
102// Return early if the format check is cancelled.
103if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
104return Err(CancelFormatChange);
105 }
106107if 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 ));
114error!(?result);
115 }
116117 prev_height = Some(height);
118 prev_tree = Some(tree);
119 }
120121let mut prev_height = None;
122let mut prev_tree = None;
123for (height, tree) in db.orchard_tree_by_height_range(..) {
124// Return early if the format check is cancelled.
125if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
126return Err(CancelFormatChange);
127 }
128129if 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 ));
136error!(?result);
137 }
138139 prev_height = Some(height);
140 prev_tree = Some(tree);
141 }
142143Ok(result)
144 }
145}