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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
//! Prunes duplicate Sapling and Orchard note commitment trees from database

use crossbeam_channel::{Receiver, TryRecvError};

use semver::Version;
use zebra_chain::block::Height;

use crate::service::finalized_state::{DiskWriteBatch, ZebraDb};

use super::{CancelFormatChange, DiskFormatUpgrade};

/// Implements [`DiskFormatUpgrade`] for pruning duplicate Sapling and Orchard note commitment trees from database
pub struct PruneTrees;

impl DiskFormatUpgrade for PruneTrees {
    fn version(&self) -> Version {
        Version::new(25, 1, 1)
    }

    fn description(&self) -> &'static str {
        "deduplicate trees upgrade"
    }

    #[allow(clippy::unwrap_in_result)]
    fn run(
        &self,
        initial_tip_height: Height,
        db: &ZebraDb,
        cancel_receiver: &Receiver<CancelFormatChange>,
    ) -> Result<(), CancelFormatChange> {
        // Prune duplicate Sapling note commitment trees.

        // The last tree we checked.
        let mut last_tree = db
            .sapling_tree_by_height(&Height(0))
            .expect("Checked above that the genesis block is in the database.");

        // Run through all the possible duplicate trees in the finalized chain.
        // The block after genesis is the first possible duplicate.
        for (height, tree) in db.sapling_tree_by_height_range(Height(1)..=initial_tip_height) {
            // Return early if there is a cancel signal.
            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
                return Err(CancelFormatChange);
            }

            // Delete any duplicate trees.
            if tree == last_tree {
                let mut batch = DiskWriteBatch::new();
                batch.delete_sapling_tree(db, &height);
                db.write_batch(batch)
                    .expect("Deleting Sapling note commitment trees should always succeed.");
            }

            // Compare against the last tree to find unique trees.
            last_tree = tree;
        }

        // Prune duplicate Orchard note commitment trees.

        // The last tree we checked.
        let mut last_tree = db
            .orchard_tree_by_height(&Height(0))
            .expect("Checked above that the genesis block is in the database.");

        // Run through all the possible duplicate trees in the finalized chain.
        // The block after genesis is the first possible duplicate.
        for (height, tree) in db.orchard_tree_by_height_range(Height(1)..=initial_tip_height) {
            // Return early if there is a cancel signal.
            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
                return Err(CancelFormatChange);
            }

            // Delete any duplicate trees.
            if tree == last_tree {
                let mut batch = DiskWriteBatch::new();
                batch.delete_orchard_tree(db, &height);
                db.write_batch(batch)
                    .expect("Deleting Orchard note commitment trees should always succeed.");
            }

            // Compare against the last tree to find unique trees.
            last_tree = tree;
        }

        Ok(())
    }

    /// Check that note commitment trees were correctly de-duplicated.
    #[allow(clippy::unwrap_in_result)]
    fn validate(
        &self,
        db: &ZebraDb,
        cancel_receiver: &Receiver<CancelFormatChange>,
    ) -> Result<Result<(), String>, CancelFormatChange> {
        // Runtime test: make sure we removed all duplicates.
        // We always run this test, even if the state has supposedly been upgraded.
        let mut result = Ok(());

        let mut prev_height = None;
        let mut prev_tree = None;
        for (height, tree) in db.sapling_tree_by_height_range(..) {
            // Return early if the format check is cancelled.
            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
                return Err(CancelFormatChange);
            }

            if prev_tree == Some(tree.clone()) {
                result = Err(format!(
                    "found duplicate sapling trees after running de-duplicate tree upgrade:\
                     height: {height:?}, previous height: {:?}, tree root: {:?}",
                    prev_height.unwrap(),
                    tree.root()
                ));
                error!(?result);
            }

            prev_height = Some(height);
            prev_tree = Some(tree);
        }

        let mut prev_height = None;
        let mut prev_tree = None;
        for (height, tree) in db.orchard_tree_by_height_range(..) {
            // Return early if the format check is cancelled.
            if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
                return Err(CancelFormatChange);
            }

            if prev_tree == Some(tree.clone()) {
                result = Err(format!(
                    "found duplicate orchard trees after running de-duplicate tree upgrade:\
                     height: {height:?}, previous height: {:?}, tree root: {:?}",
                    prev_height.unwrap(),
                    tree.root()
                ));
                error!(?result);
            }

            prev_height = Some(height);
            prev_tree = Some(tree);
        }

        Ok(result)
    }
}