Skip to content

Commit

Permalink
refactors add subtrees format upgrade to use new trait
Browse files Browse the repository at this point in the history
  • Loading branch information
@arya2
arya2 committed Feb 18, 2025
1 parent 2e3e177 commit ac1eb1b
Show file tree
Hide file tree
Showing 5 changed files with 148 additions and 144 deletions.
7 changes: 0 additions & 7 deletions zebra-state/src/constants.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
3 changes: 0 additions & 3 deletions zebra-state/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
52 changes: 16 additions & 36 deletions zebra-state/src/service/finalized_state/disk_format/upgrade.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ use zebra_chain::{

use DbFormatChange::*;

use crate::{constants::latest_version_for_adding_subtrees, service::finalized_state::ZebraDb};
use crate::service::finalized_state::ZebraDb;

pub(crate) mod add_subtrees;
pub(crate) mod cache_genesis_roots;
Expand All @@ -47,7 +47,7 @@ pub trait DiskFormatUpgrade {
initial_tip_height: Height,
db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
);
) -> Result<(), CancelFormatChange>;

/// Check that state has been upgraded to this format correctly.
///
Expand All @@ -61,13 +61,24 @@ pub trait DiskFormatUpgrade {
) -> 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(())
}
}

fn format_upgrades() -> Vec<Box<dyn DiskFormatUpgrade>> {
vec![
Box::new(prune_trees::PruneTrees),
Box::new(add_subtrees::AddSubtrees),
// TODO:
// Box::new(add_subtrees::AddSubtrees),
// Box::new(cache_genesis_roots::CacheGenesisRoots),
// Box::new(fix_tree_key_type::FixTreeKeyType),
]
Expand Down Expand Up @@ -523,7 +534,8 @@ impl DbFormatChange {

let timer = CodeTimer::start();

upgrade.run(initial_tip_height, db, cancel_receiver);
upgrade.prepare(initial_tip_height, db, cancel_receiver, older_disk_version)?;
upgrade.run(initial_tip_height, db, cancel_receiver)?;

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

// 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 =
Expand Down Expand Up @@ -666,10 +650,6 @@ impl DbFormatChange {
results.push(upgrade.validate(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)?);

Expand Down
221 changes: 126 additions & 95 deletions zebra-state/src/service/finalized_state/disk_format/upgrade/add_subtrees.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand All @@ -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)
});
impl DiskFormatUpgrade for AddSubtrees {
fn version(&self) -> Version {
Version::new(25, 2, 2)
}

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);
fn description(&self) -> &'static str {
"add subtrees upgrade"
}

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)?;
}

let subtree =
calculate_sapling_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
write_sapling_subtree(upgrade_db, subtree);
Ok(())
}

// 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)
});
/// 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);
}

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);
}

let subtree =
calculate_orchard_subtree(upgrade_db, prev_end_height, prev_tree, end_height, tree);
write_orchard_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(())
}

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.
Expand Down Expand Up @@ -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.
Expand Down
Loading

0 comments on commit ac1eb1b

Please sign in to comment.