zecfoundation
/
zebra
mirror of https://github.com/ZcashFoundation/zebra.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

change(state): Refactor format upgrades into trait (#9263) 

* Adds a new trait for disk format upgrades, implements in on a new struct, `PruneTrees`, and moves the logic for tree deduplication to the trait impl

* refactors add subtrees format upgrade to use new trait

* refactors fix tree keys, cache genesis roots, and value balance upgrades to use new trait

* Applies suggestions from code review:
- Avoids duplicate validation of format upgrades at startup when db is already upgraded,
- Minor refactors
- Doc fixes and cleanups

* Applies suggestions from code review

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Arya 2025-03-28 06:52:05 -04:00 committed by GitHub
parent 26c569e29f
commit 49011f8460
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 478 additions and 323 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
zebra-state/src/constants.rs
View File

@ -78,13 +78,6 @@ pub fn state_database_format_version_in_code() -> Version {
}
}
/// Returns the highest database version that modifies the subtree index format.
///
/// This version is used by tests to wait for the subtree upgrade to finish.
pub fn latest_version_for_adding_subtrees() -> Version {
Version::parse("25.2.2").expect("Hardcoded version string should be valid.")
}
/// The name of the file containing the minor and patch database versions.
///
/// Use [`Config::version_file_path()`] to get the path to this file.

3
zebra-state/src/lib.rs
View File

@ -81,9 +81,6 @@ pub use service::{
init_test, init_test_services,
};
#[cfg(any(test, feature = "proptest-impl"))]
pub use constants::latest_version_for_adding_subtrees;
#[cfg(any(test, feature = "proptest-impl"))]
pub use config::hidden::{
write_database_format_version_to_disk, write_state_database_format_version_to_disk,

307
zebra-state/src/service/finalized_state/disk_format/upgrade.rs
View File

@ -6,7 +6,7 @@ use std::{
thread::{self, JoinHandle},
};
use crossbeam_channel::{bounded, Receiver, RecvTimeoutError, Sender, TryRecvError};
use crossbeam_channel::{bounded, Receiver, RecvTimeoutError, Sender};
use semver::Version;
use tracing::Span;
@ -20,14 +20,14 @@ use zebra_chain::{
use DbFormatChange::*;
use crate::{
constants::latest_version_for_adding_subtrees,
service::finalized_state::{DiskWriteBatch, ZebraDb},
};
use crate::service::finalized_state::ZebraDb;
pub(crate) mod add_subtrees;
pub(crate) mod cache_genesis_roots;
pub(crate) mod fix_tree_key_type;
pub(crate) mod no_migration;
pub(crate) mod prune_trees;
pub(crate) mod tree_keys_and_caches_upgrade;
#[cfg(not(feature = "indexer"))]
pub(crate) mod drop_tx_locs_by_spends;
@ -35,6 +35,68 @@ pub(crate) mod drop_tx_locs_by_spends;
#[cfg(feature = "indexer")]
pub(crate) mod track_tx_locs_by_spends;
/// Defines method signature for running disk format upgrades.
pub trait DiskFormatUpgrade {
/// Returns the version at which this upgrade is applied.
fn version(&self) -> Version;
/// Returns the description of this upgrade.
fn description(&self) -> &'static str;
/// Runs disk format upgrade.
fn run(
&self,
initial_tip_height: Height,
db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<(), CancelFormatChange>;
/// Check that state has been upgraded to this format correctly.
///
/// The outer `Result` indicates whether the validation was cancelled (due to e.g. node shutdown).
/// The inner `Result` indicates whether the validation itself failed or not.
fn validate(
&self,
_db: &ZebraDb,
_cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<Result<(), String>, CancelFormatChange> {
Ok(Ok(()))
}
/// Prepare for disk format upgrade.
fn prepare(
&self,
_initial_tip_height: Height,
_upgrade_db: &ZebraDb,
_cancel_receiver: &Receiver<CancelFormatChange>,
_older_disk_version: &Version,
) -> Result<(), CancelFormatChange> {
Ok(())
}
/// Returns true if the [`DiskFormatUpgrade`] needs to run a migration on existing data in the db.
fn needs_migration(&self) -> bool {
true
}
}
fn format_upgrades(
min_version: Option<Version>,
) -> impl Iterator<Item = Box<dyn DiskFormatUpgrade>> {
let min_version = move || min_version.clone().unwrap_or(Version::new(0, 0, 0));
// Note: Disk format upgrades must be run in order of database version.
([
Box::new(prune_trees::PruneTrees),
Box::new(add_subtrees::AddSubtrees),
Box::new(tree_keys_and_caches_upgrade::FixTreeKeyTypeAndCacheGenesisRoots),
// Value balance upgrade
Box::new(no_migration::NoMigration::new(26, 0, 0)),
] as [Box<dyn DiskFormatUpgrade>; 4])
.into_iter()
.filter(move |upgrade| upgrade.version() > min_version())
}
/// The kind of database format change or validity check we're performing.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DbFormatChange {
@ -474,155 +536,31 @@ impl DbFormatChange {
return Ok(());
};
// Note commitment tree de-duplication database upgrade task.
// Apply or validate format upgrades
for upgrade in format_upgrades(Some(older_disk_version.clone())) {
if upgrade.needs_migration() {
let timer = CodeTimer::start();
let version_for_pruning_trees =
Version::parse("25.1.1").expect("Hardcoded version string should be valid.");
upgrade.prepare(initial_tip_height, db, cancel_receiver, older_disk_version)?;
upgrade.run(initial_tip_height, db, cancel_receiver)?;
// Check if we need to prune the note commitment trees in the database.
if older_disk_version < &version_for_pruning_trees {
let timer = CodeTimer::start();
// Before marking the state as upgraded, check that the upgrade completed successfully.
upgrade
.validate(db, cancel_receiver)?
.expect("db should be valid after upgrade");
// 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;
timer.finish(module_path!(), line!(), upgrade.description());
}
// 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;
}
// Before marking the state as upgraded, check that the upgrade completed successfully.
Self::check_for_duplicate_trees(db, cancel_receiver)?
.expect("database format is valid after upgrade");
// Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
// database is marked, so the upgrade MUST be complete at this point.
Self::mark_as_upgraded_to(db, &version_for_pruning_trees);
timer.finish(module_path!(), line!(), "deduplicate trees upgrade");
info!(
newer_running_version = ?upgrade.version(),
"Zebra automatically upgraded the database format"
);
Self::mark_as_upgraded_to(db, &upgrade.version());
}
// Note commitment subtree creation database upgrade task.
let latest_version_for_adding_subtrees = latest_version_for_adding_subtrees();
let first_version_for_adding_subtrees =
Version::parse("25.2.0").expect("Hardcoded version string should be valid.");
// Check if we need to add or fix note commitment subtrees in the database.
if older_disk_version < &latest_version_for_adding_subtrees {
let timer = CodeTimer::start();
if older_disk_version >= &first_version_for_adding_subtrees {
// Clear previous upgrade data, because it was incorrect.
add_subtrees::reset(initial_tip_height, db, cancel_receiver)?;
}
add_subtrees::run(initial_tip_height, db, cancel_receiver)?;
// Before marking the state as upgraded, check that the upgrade completed successfully.
add_subtrees::subtree_format_validity_checks_detailed(db, cancel_receiver)?
.expect("database format is valid after upgrade");
// Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
// database is marked, so the upgrade MUST be complete at this point.
Self::mark_as_upgraded_to(db, &latest_version_for_adding_subtrees);
timer.finish(module_path!(), line!(), "add subtrees upgrade");
}
// Sprout & history tree key formats, and cached genesis tree roots database upgrades.
let version_for_tree_keys_and_caches =
Version::parse("25.3.0").expect("Hardcoded version string should be valid.");
// Check if we need to do the upgrade.
if older_disk_version < &version_for_tree_keys_and_caches {
let timer = CodeTimer::start();
// It shouldn't matter what order these are run in.
cache_genesis_roots::run(initial_tip_height, db, cancel_receiver)?;
fix_tree_key_type::run(initial_tip_height, db, cancel_receiver)?;
// Before marking the state as upgraded, check that the upgrade completed successfully.
cache_genesis_roots::detailed_check(db, cancel_receiver)?
.expect("database format is valid after upgrade");
fix_tree_key_type::detailed_check(db, cancel_receiver)?
.expect("database format is valid after upgrade");
// Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
// database is marked, so the upgrade MUST be complete at this point.
Self::mark_as_upgraded_to(db, &version_for_tree_keys_and_caches);
timer.finish(module_path!(), line!(), "tree keys and caches upgrade");
}
let version_for_upgrading_value_balance_format =
Version::parse("26.0.0").expect("hard-coded version string should be valid");
// Check if we need to do the upgrade.
if older_disk_version < &version_for_upgrading_value_balance_format {
Self::mark_as_upgraded_to(db, &version_for_upgrading_value_balance_format)
}
// # New Upgrades Usually Go Here
//
// New code goes above this comment!
//
// Run the latest format upgrade code after the other upgrades are complete,
// then mark the format as upgraded. The code should check `cancel_receiver`
// every time it runs its inner update loop.
info!(
%newer_running_version,
"Zebra automatically upgraded the database format to:"
);
Ok(())
}
@ -669,13 +607,9 @@ impl DbFormatChange {
// Do the quick checks first, so we don't have to do this in every detailed check.
results.push(Self::format_validity_checks_quick(db));
results.push(Self::check_for_duplicate_trees(db, cancel_receiver)?);
results.push(add_subtrees::subtree_format_validity_checks_detailed(
db,
cancel_receiver,
)?);
results.push(cache_genesis_roots::detailed_check(db, cancel_receiver)?);
results.push(fix_tree_key_type::detailed_check(db, cancel_receiver)?);
for upgrade in format_upgrades(None) {
results.push(upgrade.validate(db, cancel_receiver)?);
}
// The work is done in the functions we just called.
timer.finish(module_path!(), line!(), "format_validity_checks_detailed()");
@ -689,66 +623,6 @@ impl DbFormatChange {
Ok(Ok(()))
}
/// Check that note commitment trees were correctly de-duplicated.
//
// TODO: move this method into an deduplication upgrade module file,
// along with the upgrade code above.
#[allow(clippy::unwrap_in_result)]
fn check_for_duplicate_trees(
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)
}
/// Mark a newly created database with the current format version.
///
/// This should be called when a newly created database is opened.
@ -950,3 +824,12 @@ impl Drop for DbFormatChangeThreadHandle {
}
}
}
#[test]
fn format_upgrades_are_in_version_order() {
let mut last_version = Version::new(0, 0, 0);
for upgrade in format_upgrades(None) {
assert!(upgrade.version() > last_version);
last_version = upgrade.version();
}
}

233
zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
View File

@ -5,6 +5,7 @@ use std::sync::Arc;
use crossbeam_channel::{Receiver, TryRecvError};
use hex_literal::hex;
use itertools::Itertools;
use semver::Version;
use tracing::instrument;
use zebra_chain::{
@ -17,90 +18,144 @@ use zebra_chain::{
};
use crate::service::finalized_state::{
disk_format::upgrade::CancelFormatChange, DiskWriteBatch, ZebraDb,
disk_format::upgrade::{CancelFormatChange, DiskFormatUpgrade},
DiskWriteBatch, ZebraDb,
};
/// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
///
/// Trees are added to the database in reverse height order, so that wallets can sync correctly
/// while the upgrade is running.
///
/// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
#[allow(clippy::unwrap_in_result)]
#[instrument(skip(upgrade_db, cancel_receiver))]
pub fn run(
initial_tip_height: Height,
upgrade_db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<(), CancelFormatChange> {
// # Consensus
//
// Zebra stores exactly one note commitment tree for every block with sapling notes.
// (It also stores the empty note commitment tree for the genesis block, but we skip that.)
//
// The consensus rules limit blocks to less than 2^16 sapling and 2^16 orchard outputs. So a
// block can't complete multiple level 16 subtrees (or complete an entire subtree by itself).
// Currently, with 2MB blocks and v4/v5 sapling and orchard output sizes, the subtree index can
// increase by at most 1 every ~20 blocks.
//
// # Compatibility
//
// Because wallets search backwards from the chain tip, subtrees need to be added to the
// database in reverse height order. (Tip first, genesis last.)
//
// Otherwise, wallets that sync during the upgrade will be missing some notes.
/// Implements [`DiskFormatUpgrade`] for populating Sapling and Orchard note commitment subtrees.
pub struct AddSubtrees;
// Generate a list of sapling subtree inputs: previous and current trees, and their end heights.
let subtrees = upgrade_db
.sapling_tree_by_reversed_height_range(..=initial_tip_height)
// We need both the tree and its previous tree for each shielded block.
.tuple_windows()
// Because the iterator is reversed, the larger tree is first.
.map(|((end_height, tree), (prev_end_height, prev_tree))| {
(prev_end_height, prev_tree, end_height, tree)
})
// Find new subtrees.
.filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
tree.contains_new_subtree(prev_tree)
});
for (prev_end_height, prev_tree, end_height, tree) in subtrees {
// Return early if the upgrade is cancelled.
if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
return Err(CancelFormatChange);
}
let subtree =
calculate_sapling_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
write_sapling_subtree(upgrade_db, subtree);
impl DiskFormatUpgrade for AddSubtrees {
fn version(&self) -> Version {
Version::new(25, 2, 2)
}
// Generate a list of orchard subtree inputs: previous and current trees, and their end heights.
let subtrees = upgrade_db
.orchard_tree_by_reversed_height_range(..=initial_tip_height)
// We need both the tree and its previous tree for each shielded block.
.tuple_windows()
// Because the iterator is reversed, the larger tree is first.
.map(|((end_height, tree), (prev_end_height, prev_tree))| {
(prev_end_height, prev_tree, end_height, tree)
})
// Find new subtrees.
.filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
tree.contains_new_subtree(prev_tree)
});
for (prev_end_height, prev_tree, end_height, tree) in subtrees {
// Return early if the upgrade is cancelled.
if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
return Err(CancelFormatChange);
}
let subtree =
calculate_orchard_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
write_orchard_subtree(upgrade_db, subtree);
fn description(&self) -> &'static str {
"add subtrees upgrade"
}
Ok(())
fn prepare(
&self,
initial_tip_height: Height,
upgrade_db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
older_disk_version: &Version,
) -> Result<(), CancelFormatChange> {
let first_version_for_adding_subtrees = Version::new(25, 2, 0);
if older_disk_version >= &first_version_for_adding_subtrees {
// Clear previous upgrade data, because it was incorrect.
reset(initial_tip_height, upgrade_db, cancel_receiver)?;
}
Ok(())
}
/// Runs disk format upgrade for adding Sapling and Orchard note commitment subtrees to database.
///
/// Trees are added to the database in reverse height order, so that wallets can sync correctly
/// while the upgrade is running.
///
/// Returns `Ok` if the upgrade completed, and `Err` if it was cancelled.
fn run(
&self,
initial_tip_height: Height,
upgrade_db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<(), CancelFormatChange> {
// # Consensus
//
// Zebra stores exactly one note commitment tree for every block with sapling notes.
// (It also stores the empty note commitment tree for the genesis block, but we skip that.)
//
// The consensus rules limit blocks to less than 2^16 sapling and 2^16 orchard outputs. So a
// block can't complete multiple level 16 subtrees (or complete an entire subtree by itself).
// Currently, with 2MB blocks and v4/v5 sapling and orchard output sizes, the subtree index can
// increase by at most 1 every ~20 blocks.
//
// # Compatibility
//
// Because wallets search backwards from the chain tip, subtrees need to be added to the
// database in reverse height order. (Tip first, genesis last.)
//
// Otherwise, wallets that sync during the upgrade will be missing some notes.
// Generate a list of sapling subtree inputs: previous and current trees, and their end heights.
let subtrees = upgrade_db
.sapling_tree_by_reversed_height_range(..=initial_tip_height)
// We need both the tree and its previous tree for each shielded block.
.tuple_windows()
// Because the iterator is reversed, the larger tree is first.
.map(|((end_height, tree), (prev_end_height, prev_tree))| {
(prev_end_height, prev_tree, end_height, tree)
})
// Find new subtrees.
.filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
tree.contains_new_subtree(prev_tree)
});
for (prev_end_height, prev_tree, end_height, tree) in subtrees {
// Return early if the upgrade is cancelled.
if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
return Err(CancelFormatChange);
}
let subtree =
calculate_sapling_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
write_sapling_subtree(upgrade_db, subtree);
}
// Generate a list of orchard subtree inputs: previous and current trees, and their end heights.
let subtrees = upgrade_db
.orchard_tree_by_reversed_height_range(..=initial_tip_height)
// We need both the tree and its previous tree for each shielded block.
.tuple_windows()
// Because the iterator is reversed, the larger tree is first.
.map(|((end_height, tree), (prev_end_height, prev_tree))| {
(prev_end_height, prev_tree, end_height, tree)
})
// Find new subtrees.
.filter(|(_prev_end_height, prev_tree, _end_height, tree)| {
tree.contains_new_subtree(prev_tree)
});
for (prev_end_height, prev_tree, end_height, tree) in subtrees {
// Return early if the upgrade is cancelled.
if !matches!(cancel_receiver.try_recv(), Err(TryRecvError::Empty)) {
return Err(CancelFormatChange);
}
let subtree =
calculate_orchard_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
write_orchard_subtree(upgrade_db, subtree);
}
Ok(())
}
#[allow(clippy::unwrap_in_result)]
fn validate(
&self,
db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<Result<(), String>, CancelFormatChange> {
// This is redundant in some code paths, but not in others. But it's quick anyway.
let quick_result = subtree_format_calculation_pre_checks(db);
// Check the entire format before returning any errors.
let sapling_result = check_sapling_subtrees(db, cancel_receiver)?;
let orchard_result = check_orchard_subtrees(db, cancel_receiver)?;
if quick_result.is_err() || sapling_result.is_err() || orchard_result.is_err() {
let err = Err(format!(
"missing or invalid subtree(s): \
quick: {quick_result:?}, sapling: {sapling_result:?}, orchard: {orchard_result:?}"
));
warn!(?err);
return Ok(err);
}
Ok(Ok(()))
}
}
/// 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> {
Ok(())
}
/// Check that note commitment subtrees were correctly added.
pub fn subtree_format_validity_checks_detailed(
db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<Result<(), String>, CancelFormatChange> {
// This is redundant in some code paths, but not in others. But it's quick anyway.
let quick_result = subtree_format_calculation_pre_checks(db);
// Check the entire format before returning any errors.
let sapling_result = check_sapling_subtrees(db, cancel_receiver)?;
let orchard_result = check_orchard_subtrees(db, cancel_receiver)?;
if quick_result.is_err() || sapling_result.is_err() || orchard_result.is_err() {
let err = Err(format!(
"missing or invalid subtree(s): \
quick: {quick_result:?}, sapling: {sapling_result:?}, orchard: {orchard_result:?}"
));
warn!(?err);
return Ok(err);
}
Ok(Ok(()))
}
/// Check that Sapling note commitment subtrees were correctly added.
///
/// Returns an error if a note commitment subtree is missing or incorrect.

49
zebra-state/src/service/finalized_state/disk_format/upgrade/no_migration.rs Normal file
View File

@ -0,0 +1,49 @@
//! An implementation of [`DiskFormatUpgrade`] for marking the database as upgraded to a new format version.
use crossbeam_channel::Receiver;
use semver::Version;
use zebra_chain::block::Height;
use crate::service::finalized_state::ZebraDb;
use super::{CancelFormatChange, DiskFormatUpgrade};
/// Implements [`DiskFormatUpgrade`] for in-place upgrades that do not involve any migration
/// of existing data into the new format.
pub struct NoMigration {
version: Version,
}
impl NoMigration {
/// Creates a new instance of the [`NoMigration`] upgrade.
pub fn new(major: u64, minor: u64, patch: u64) -> Self {
Self {
version: Version::new(major, minor, patch),
}
}
}
impl DiskFormatUpgrade for NoMigration {
fn version(&self) -> Version {
self.version.clone()
}
fn description(&self) -> &'static str {
"no migration"
}
#[allow(clippy::unwrap_in_result)]
fn run(
&self,
_initial_tip_height: Height,
_db: &ZebraDb,
_cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<(), CancelFormatChange> {
Ok(())
}
fn needs_migration(&self) -> bool {
false
}
}

145
zebra-state/src/service/finalized_state/disk_format/upgrade/prune_trees.rs Normal file
View File

@ -0,0 +1,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)
}
}

57
zebra-state/src/service/finalized_state/disk_format/upgrade/tree_keys_and_caches_upgrade.rs Normal file
View File

@ -0,0 +1,57 @@
//! Applies the [`fix_tree_key_type`] and [`cache_genesis_roots`] upgrades to the database.
use crossbeam_channel::Receiver;
use semver::Version;
use zebra_chain::block::Height;
use crate::service::finalized_state::ZebraDb;
use super::{cache_genesis_roots, fix_tree_key_type, CancelFormatChange, DiskFormatUpgrade};
/// Implements [`DiskFormatUpgrade`] for updating the sprout and history tree key type from
/// `Height` to the empty key `()` and the genesis note commitment trees to cache their roots
pub struct FixTreeKeyTypeAndCacheGenesisRoots;
impl DiskFormatUpgrade for FixTreeKeyTypeAndCacheGenesisRoots {
fn version(&self) -> Version {
Version::new(25, 3, 0)
}
fn description(&self) -> &'static str {
"tree keys and caches upgrade"
}
#[allow(clippy::unwrap_in_result)]
fn run(
&self,
initial_tip_height: Height,
db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<(), CancelFormatChange> {
// It shouldn't matter what order these are run in.
cache_genesis_roots::run(initial_tip_height, db, cancel_receiver)?;
fix_tree_key_type::run(initial_tip_height, db, cancel_receiver)?;
Ok(())
}
#[allow(clippy::unwrap_in_result)]
fn validate(
&self,
db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<Result<(), String>, CancelFormatChange> {
let results = [
cache_genesis_roots::detailed_check(db, cancel_receiver)?,
fix_tree_key_type::detailed_check(db, cancel_receiver)?,
];
let result = if results.iter().any(Result::is_err) {
Err(format!("{results:?}"))
} else {
Ok(())
};
Ok(result)
}
}