Skip to content

Commit ac1eb1b

Browse files
arya2arya2
committed
refactors add subtrees format upgrade to use new trait
1 parent 2e3e177 commit ac1eb1b

File tree

5 files changed

+148
-144
lines changed

5 files changed

+148
-144
lines changed

‎zebra-state/src/constants.rs

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -78,13 +78,6 @@ pub fn state_database_format_version_in_code() -> Version {
7878
}
7979
}
8080

81-
/// Returns the highest database version that modifies the subtree index format.
82-
///
83-
/// This version is used by tests to wait for the subtree upgrade to finish.
84-
pub fn latest_version_for_adding_subtrees() -> Version {
85-
Version::parse("25.2.2").expect("Hardcoded version string should be valid.")
86-
}
87-
8881
/// The name of the file containing the minor and patch database versions.
8982
///
9083
/// Use [`Config::version_file_path()`] to get the path to this file.

‎zebra-state/src/lib.rs

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,9 +81,6 @@ pub use service::{
8181
init_test, init_test_services,
8282
};
8383

84-
#[cfg(any(test, feature = "proptest-impl"))]
85-
pub use constants::latest_version_for_adding_subtrees;
86-
8784
#[cfg(any(test, feature = "proptest-impl"))]
8885
pub use config::hidden::{
8986
write_database_format_version_to_disk, write_state_database_format_version_to_disk,

‎zebra-state/src/service/finalized_state/disk_format/upgrade.rs

Lines changed: 16 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use zebra_chain::{
2020

2121
use DbFormatChange::*;
2222

23-
use crate::{constants::latest_version_for_adding_subtrees, service::finalized_state::ZebraDb};
23+
use crate::service::finalized_state::ZebraDb;
2424

2525
pub(crate) mod add_subtrees;
2626
pub(crate) mod cache_genesis_roots;
@@ -47,7 +47,7 @@ pub trait DiskFormatUpgrade {
4747
initial_tip_height: Height,
4848
db: &ZebraDb,
4949
cancel_receiver: &Receiver<CancelFormatChange>,
50-
);
50+
) -> Result<(), CancelFormatChange>;
5151

5252
/// Check that state has been upgraded to this format correctly.
5353
///
@@ -61,13 +61,24 @@ pub trait DiskFormatUpgrade {
6161
) -> Result<Result<(), String>, CancelFormatChange> {
6262
Ok(Ok(()))
6363
}
64+
65+
/// Prepare for disk format upgrade.
66+
fn prepare(
67+
&self,
68+
_initial_tip_height: Height,
69+
_upgrade_db: &ZebraDb,
70+
_cancel_receiver: &Receiver<CancelFormatChange>,
71+
_older_disk_version: &Version,
72+
) -> Result<(), CancelFormatChange> {
73+
Ok(())
74+
}
6475
}
6576

6677
fn format_upgrades() -> Vec<Box<dyn DiskFormatUpgrade>> {
6778
vec![
6879
Box::new(prune_trees::PruneTrees),
80+
Box::new(add_subtrees::AddSubtrees),
6981
// TODO:
70-
// Box::new(add_subtrees::AddSubtrees),
7182
// Box::new(cache_genesis_roots::CacheGenesisRoots),
7283
// Box::new(fix_tree_key_type::FixTreeKeyType),
7384
]
@@ -523,7 +534,8 @@ impl DbFormatChange {
523534

524535
let timer = CodeTimer::start();
525536

526-
upgrade.run(initial_tip_height, db, cancel_receiver);
537+
upgrade.prepare(initial_tip_height, db, cancel_receiver, older_disk_version)?;
538+
upgrade.run(initial_tip_height, db, cancel_receiver)?;
527539

528540
// Before marking the state as upgraded, check that the upgrade completed successfully.
529541
upgrade
@@ -541,34 +553,6 @@ impl DbFormatChange {
541553
timer.finish(module_path!(), line!(), upgrade.description());
542554
}
543555

544-
// Note commitment subtree creation database upgrade task.
545-
546-
let latest_version_for_adding_subtrees = latest_version_for_adding_subtrees();
547-
let first_version_for_adding_subtrees =
548-
Version::parse("25.2.0").expect("Hardcoded version string should be valid.");
549-
550-
// Check if we need to add or fix note commitment subtrees in the database.
551-
if older_disk_version < &latest_version_for_adding_subtrees {
552-
let timer = CodeTimer::start();
553-
554-
if older_disk_version >= &first_version_for_adding_subtrees {
555-
// Clear previous upgrade data, because it was incorrect.
556-
add_subtrees::reset(initial_tip_height, db, cancel_receiver)?;
557-
}
558-
559-
add_subtrees::run(initial_tip_height, db, cancel_receiver)?;
560-
561-
// Before marking the state as upgraded, check that the upgrade completed successfully.
562-
add_subtrees::subtree_format_validity_checks_detailed(db, cancel_receiver)?
563-
.expect("database format is valid after upgrade");
564-
565-
// Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
566-
// database is marked, so the upgrade MUST be complete at this point.
567-
Self::mark_as_upgraded_to(db, &latest_version_for_adding_subtrees);
568-
569-
timer.finish(module_path!(), line!(), "add subtrees upgrade");
570-
}
571-
572556
// Sprout & history tree key formats, and cached genesis tree roots database upgrades.
573557

574558
let version_for_tree_keys_and_caches =
@@ -666,10 +650,6 @@ impl DbFormatChange {
666650
results.push(upgrade.validate(db, cancel_receiver)?);
667651
}
668652

669-
results.push(add_subtrees::subtree_format_validity_checks_detailed(
670-
db,
671-
cancel_receiver,
672-
)?);
673653
results.push(cache_genesis_roots::detailed_check(db, cancel_receiver)?);
674654
results.push(fix_tree_key_type::detailed_check(db, cancel_receiver)?);
675655

‎zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs

Lines changed: 126 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use std::sync::Arc;
55
use crossbeam_channel::{Receiver, TryRecvError};
66
use hex_literal::hex;
77
use itertools::Itertools;
8+
use semver::Version;
89
use tracing::instrument;
910

1011
use zebra_chain::{
@@ -17,90 +18,144 @@ use zebra_chain::{
1718
};
1819

1920
use crate::service::finalized_state::{
20-
disk_format::upgrade::CancelFormatChange, DiskWriteBatch, ZebraDb,
21+
disk_format::upgrade::{CancelFormatChange, DiskFormatUpgrade},
22+
DiskWriteBatch, ZebraDb,
2123
};
2224

23-
/// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
24-
///
25-
/// Trees are added to the database in reverse height order, so that wallets can sync correctly
26-
/// while the upgrade is running.
27-
///
28-
/// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
29-
#[allow(clippy::unwrap_in_result)]
30-
#[instrument(skip(upgrade_db, cancel_receiver))]
31-
pub fn run(
32-
initial_tip_height: Height,
33-
upgrade_db: &ZebraDb,
34-
cancel_receiver: &Receiver<CancelFormatChange>,
35-
) -> Result<(), CancelFormatChange> {
36-
// # Consensus
37-
//
38-
// Zebra stores exactly one note commitment tree for every block with sapling notes.
39-
// (It also stores the empty note commitment tree for the genesis block, but we skip that.)
40-
//
41-
// The consensus rules limit blocks to less than 2^16 sapling and 2^16 orchard outputs. So a
42-
// block can't complete multiple level 16 subtrees (or complete an entire subtree by itself).
43-
// Currently, with 2MB blocks and v4/v5 sapling and orchard output sizes, the subtree index can
44-
// increase by at most 1 every ~20 blocks.
45-
//
46-
// # Compatibility
47-
//
48-
// Because wallets search backwards from the chain tip, subtrees need to be added to the
49-
// database in reverse height order. (Tip first, genesis last.)
50-
//
51-
// Otherwise, wallets that sync during the upgrade will be missing some notes.
25+
/// Implements [`DiskFormatUpgrade`] for populating Sapling and Orchard note commitment subtrees.
26+
pub struct AddSubtrees;
5227

53-
// Generate a list of sapling subtree inputs: previous and current trees, and their end heights.
54-
let subtrees = upgrade_db
55-
.sapling_tree_by_reversed_height_range(..=initial_tip_height)
56-
// We need both the tree and its previous tree for each shielded block.
57-
.tuple_windows()
58-
// Because the iterator is reversed, the larger tree is first.
59-
.map(|((end_height, tree), (prev_end_height, prev_tree))| {
60-
(prev_end_height, prev_tree, end_height, tree)
61-
})
62-
// Find new subtrees.
63-
.filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
64-
tree.contains_new_subtree(prev_tree)
65-
});
28+
impl DiskFormatUpgrade for AddSubtrees {
29+
fn version(&self) -> Version {
30+
Version::new(25, 2, 2)
31+
}
6632

67-
for (prev_end_height, prev_tree, end_height, tree) in subtrees {
68-
// Return early if the upgrade is cancelled.
69-
if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
70-
return Err(CancelFormatChange);
33+
fn description(&self) -> &'static str {
34+
"add subtrees upgrade"
35+
}
36+
37+
fn prepare(
38+
&self,
39+
initial_tip_height: Height,
40+
upgrade_db: &ZebraDb,
41+
cancel_receiver: &Receiver<CancelFormatChange>,
42+
older_disk_version: &Version,
43+
) -> Result<(), CancelFormatChange> {
44+
let first_version_for_adding_subtrees = Version::new(25, 2, 0);
45+
if older_disk_version >= &first_version_for_adding_subtrees {
46+
// Clear previous upgrade data, because it was incorrect.
47+
reset(initial_tip_height, upgrade_db, cancel_receiver)?;
7148
}
7249

73-
let subtree =
74-
calculate_sapling_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
75-
write_sapling_subtree(upgrade_db, subtree);
50+
Ok(())
7651
}
7752

78-
// Generate a list of orchard subtree inputs: previous and current trees, and their end heights.
79-
let subtrees = upgrade_db
80-
.orchard_tree_by_reversed_height_range(..=initial_tip_height)
81-
// We need both the tree and its previous tree for each shielded block.
82-
.tuple_windows()
83-
// Because the iterator is reversed, the larger tree is first.
84-
.map(|((end_height, tree), (prev_end_height, prev_tree))| {
85-
(prev_end_height, prev_tree, end_height, tree)
86-
})
87-
// Find new subtrees.
88-
.filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
89-
tree.contains_new_subtree(prev_tree)
90-
});
53+
/// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
54+
///
55+
/// Trees are added to the database in reverse height order, so that wallets can sync correctly
56+
/// while the upgrade is running.
57+
///
58+
/// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
59+
fn run(
60+
&self,
61+
initial_tip_height: Height,
62+
upgrade_db: &ZebraDb,
63+
cancel_receiver: &Receiver<CancelFormatChange>,
64+
) -> Result<(), CancelFormatChange> {
65+
// # Consensus
66+
//
67+
// Zebra stores exactly one note commitment tree for every block with sapling notes.
68+
// (It also stores the empty note commitment tree for the genesis block, but we skip that.)
69+
//
70+
// The consensus rules limit blocks to less than 2^16 sapling and 2^16 orchard outputs. So a
71+
// block can't complete multiple level 16 subtrees (or complete an entire subtree by itself).
72+
// Currently, with 2MB blocks and v4/v5 sapling and orchard output sizes, the subtree index can
73+
// increase by at most 1 every ~20 blocks.
74+
//
75+
// # Compatibility
76+
//
77+
// Because wallets search backwards from the chain tip, subtrees need to be added to the
78+
// database in reverse height order. (Tip first, genesis last.)
79+
//
80+
// Otherwise, wallets that sync during the upgrade will be missing some notes.
81+
82+
// Generate a list of sapling subtree inputs: previous and current trees, and their end heights.
83+
let subtrees = upgrade_db
84+
.sapling_tree_by_reversed_height_range(..=initial_tip_height)
85+
// We need both the tree and its previous tree for each shielded block.
86+
.tuple_windows()
87+
// Because the iterator is reversed, the larger tree is first.
88+
.map(|((end_height, tree), (prev_end_height, prev_tree))| {
89+
(prev_end_height, prev_tree, end_height, tree)
90+
})
91+
// Find new subtrees.
92+
.filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
93+
tree.contains_new_subtree(prev_tree)
94+
});
95+
96+
for (prev_end_height, prev_tree, end_height, tree) in subtrees {
97+
// Return early if the upgrade is cancelled.
98+
if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
99+
return Err(CancelFormatChange);
100+
}
91101

92-
for (prev_end_height, prev_tree, end_height, tree) in subtrees {
93-
// Return early if the upgrade is cancelled.
94-
if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
95-
return Err(CancelFormatChange);
102+
let subtree =
103+
calculate_sapling_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
104+
write_sapling_subtree(upgrade_db, subtree);
96105
}
97106

98-
let subtree =
99-
calculate_orchard_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
100-
write_orchard_subtree(upgrade_db, subtree);
107+
// Generate a list of orchard subtree inputs: previous and current trees, and their end heights.
108+
let subtrees = upgrade_db
109+
.orchard_tree_by_reversed_height_range(..=initial_tip_height)
110+
// We need both the tree and its previous tree for each shielded block.
111+
.tuple_windows()
112+
// Because the iterator is reversed, the larger tree is first.
113+
.map(|((end_height, tree), (prev_end_height, prev_tree))| {
114+
(prev_end_height, prev_tree, end_height, tree)
115+
})
116+
// Find new subtrees.
117+
.filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
118+
tree.contains_new_subtree(prev_tree)
119+
});
120+
121+
for (prev_end_height, prev_tree, end_height, tree) in subtrees {
122+
// Return early if the upgrade is cancelled.
123+
if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
124+
return Err(CancelFormatChange);
125+
}
126+
127+
let subtree =
128+
calculate_orchard_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
129+
write_orchard_subtree(upgrade_db, subtree);
130+
}
131+
132+
Ok(())
101133
}
102134

103-
Ok(())
135+
#[allow(clippy::unwrap_in_result)]
136+
fn validate(
137+
&self,
138+
db: &ZebraDb,
139+
cancel_receiver: &Receiver<CancelFormatChange>,
140+
) -> Result<Result<(), String>, CancelFormatChange> {
141+
// This is redundant in some code paths, but not in others. But it's quick anyway.
142+
let quick_result = subtree_format_calculation_pre_checks(db);
143+
144+
// Check the entire format before returning any errors.
145+
let sapling_result = check_sapling_subtrees(db, cancel_receiver)?;
146+
let orchard_result = check_orchard_subtrees(db, cancel_receiver)?;
147+
148+
if quick_result.is_err() || sapling_result.is_err() || orchard_result.is_err() {
149+
let err = Err(format!(
150+
"missing or invalid subtree(s): \
151+
quick: {quick_result:?}, sapling: {sapling_result:?}, orchard: {orchard_result:?}"
152+
));
153+
warn!(?err);
154+
return Ok(err);
155+
}
156+
157+
Ok(Ok(()))
158+
}
104159
}
105160

106161
/// Reset data from previous upgrades. This data can be complete or incomplete.
@@ -304,30 +359,6 @@ fn quick_check_orchard_subtrees(db: &ZebraDb) -> Result<(), &'static str> {
304359
Ok(())
305360
}
306361

307-
/// Check that note commitment subtrees were correctly added.
308-
pub fn subtree_format_validity_checks_detailed(
309-
db: &ZebraDb,
310-
cancel_receiver: &Receiver<CancelFormatChange>,
311-
) -> Result<Result<(), String>, CancelFormatChange> {
312-
// This is redundant in some code paths, but not in others. But it's quick anyway.
313-
let quick_result = subtree_format_calculation_pre_checks(db);
314-
315-
// Check the entire format before returning any errors.
316-
let sapling_result = check_sapling_subtrees(db, cancel_receiver)?;
317-
let orchard_result = check_orchard_subtrees(db, cancel_receiver)?;
318-
319-
if quick_result.is_err() || sapling_result.is_err() || orchard_result.is_err() {
320-
let err = Err(format!(
321-
"missing or invalid subtree(s): \
322-
quick: {quick_result:?}, sapling: {sapling_result:?}, orchard: {orchard_result:?}"
323-
));
324-
warn!(?err);
325-
return Ok(err);
326-
}
327-
328-
Ok(Ok(()))
329-
}
330-
331362
/// Check that Sapling note commitment subtrees were correctly added.
332363
///
333364
/// Returns an error if a note commitment subtree is missing or incorrect.

0 commit comments

Comments
 (0)