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

1//! Updating the sprout and history tree key type from `Height` to the empty key `()`.
2//!
3//! This avoids a potential concurrency bug, and a known database performance issue.
4
5use std::sync::Arc;
6
7use crossbeam_channel::{Receiver, TryRecvError};
8use zebra_chain::{block::Height, history_tree::HistoryTree, sprout};
9
10use crate::service::finalized_state::{
11    disk_db::DiskWriteBatch, disk_format::MAX_ON_DISK_HEIGHT, ZebraDb,
12};
13
14use super::CancelFormatChange;
15
16/// Runs disk format upgrade for changing the sprout and history tree key types.
17///
18/// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
19#[allow(clippy::unwrap_in_result)]
20#[instrument(skip(upgrade_db, cancel_receiver))]
21pub fn run(
22    _initial_tip_height: Height,
23    upgrade_db: &ZebraDb,
24    cancel_receiver: &Receiver<CancelFormatChange>,
25) -> Result<(), CancelFormatChange> {
26    let sprout_tip_tree = upgrade_db.sprout_tree_for_tip();
27    let history_tip_tree = upgrade_db.history_tree();
28
29    // Writing the trees back to the database automatically updates their format.
30    let mut batch = DiskWriteBatch::new();
31
32    // Update the sprout tip key format in the database.
33    batch.update_sprout_tree(upgrade_db, &sprout_tip_tree);
34    batch.update_history_tree(upgrade_db, &history_tip_tree);
35
36    // Return before we write if the upgrade is cancelled.
37    if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
38        return Err(CancelFormatChange);
39    }
40
41    upgrade_db
42        .write_batch(batch)
43        .expect("updating tree key formats should always succeed");
44
45    // The deletes below can be slow due to tombstones for previously deleted keys,
46    // so we do it in a separate batch to avoid data races with syncing (#7961).
47    let mut batch = DiskWriteBatch::new();
48
49    // Delete the previous `Height` tip key format, which is now a duplicate.
50    // This doesn't delete the new `()` key format, because it serializes to an empty array.
51    batch.delete_range_sprout_tree(upgrade_db, &Height(0), &MAX_ON_DISK_HEIGHT);
52    batch.delete_range_history_tree(upgrade_db, &Height(0), &MAX_ON_DISK_HEIGHT);
53
54    // Return before we write if the upgrade is cancelled.
55    if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
56        return Err(CancelFormatChange);
57    }
58
59    upgrade_db
60        .write_batch(batch)
61        .expect("cleaning up old tree key formats should always succeed");
62
63    Ok(())
64}
65
66/// Quickly check that the sprout and history tip trees have updated key formats.
67///
68/// # Panics
69///
70/// If the state is empty.
71pub fn quick_check(db: &ZebraDb) -> Result<(), String> {
72    // Check the entire format before returning any errors.
73    let mut result = Ok(());
74
75    let mut prev_key = None;
76    let mut prev_tree: Option<Arc<sprout::tree::NoteCommitmentTree>> = None;
77
78    for (key, tree) in db.sprout_trees_full_tip() {
79        // The tip tree should be indexed by `()` (which serializes to an empty array).
80        if !key.raw_bytes().is_empty() {
81            result = Err(format!(
82                "found incorrect sprout tree key format after running key format upgrade \
83                 key: {key:?}, tree: {:?}",
84                tree.root()
85            ));
86            error!(?result);
87        }
88
89        // There should only be one tip tree in this column family.
90        if prev_tree.is_some() {
91            result = Err(format!(
92                "found duplicate sprout trees after running key format upgrade\n\
93                 key: {key:?}, tree: {:?}\n\
94                 prev key: {prev_key:?}, prev_tree: {:?}\n\
95                 ",
96                tree.root(),
97                prev_tree.unwrap().root(),
98            ));
99            error!(?result);
100        }
101
102        prev_key = Some(key);
103        prev_tree = Some(tree);
104    }
105
106    let mut prev_key = None;
107    let mut prev_tree: Option<Arc<HistoryTree>> = None;
108
109    for (key, tree) in db.history_trees_full_tip() {
110        // The tip tree should be indexed by `()` (which serializes to an empty array).
111        if !key.raw_bytes().is_empty() {
112            result = Err(format!(
113                "found incorrect history tree key format after running key format upgrade \
114                 key: {key:?}, tree: {:?}",
115                tree.hash()
116            ));
117            error!(?result);
118        }
119
120        // There should only be one tip tree in this column family.
121        if prev_tree.is_some() {
122            result = Err(format!(
123                "found duplicate history trees after running key format upgrade\n\
124                 key: {key:?}, tree: {:?}\n\
125                 prev key: {prev_key:?}, prev_tree: {:?}\n\
126                 ",
127                tree.hash(),
128                prev_tree.unwrap().hash(),
129            ));
130            error!(?result);
131        }
132
133        prev_key = Some(key);
134        prev_tree = Some(tree);
135    }
136
137    result
138}
139
140/// Detailed check that the sprout and history tip trees have updated key formats.
141/// This is currently the same as the quick check.
142///
143/// # Panics
144///
145/// If the state is empty.
146pub fn detailed_check(
147    db: &ZebraDb,
148    _cancel_receiver: &Receiver<CancelFormatChange>,
149) -> Result<Result<(), String>, CancelFormatChange> {
150    // This upgrade only changes two key-value pairs, so checking it is always quick.
151    Ok(quick_check(db))
152}