Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
274066b
Adds a new trait for disk format upgrades, implements in on a new str…
arya2 Feb 18, 2025
9433654
refactors add subtrees format upgrade to use new trait
arya2 Feb 18, 2025
59c0b24
refactors fix tree keys, cache genesis roots, and value balance upgra…
arya2 Feb 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 0 additions & 7 deletions zebra-state/src/constants.rs
View file
Open in desktop
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
Open in desktop
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
280 changes: 86 additions & 194 deletions zebra-state/src/service/finalized_state/disk_format/upgrade.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -20,21 +20,78 @@ 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;

#[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(
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

run_migration() may be a better name for this method.

&self,
initial_tip_height: Height,
db: &ZebraDb,
cancel_receiver: &Receiver<CancelFormatChange>,
) -> Result<(), CancelFormatChange>;

/// Check that state has been upgraded to this format correctly.
///
/// # Panics
///
/// If the state has not been upgraded to this format correctly.
Comment on lines +55 to +58
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
///
/// # Panics
///
/// If the state has not been upgraded to this format correctly.

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() -> Vec<Box<dyn DiskFormatUpgrade>> {
// Note: Disk format upgrades must be run in order.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would be good to check if they upgrades are really in order and panic if they're not

Suggested change
// Note: Disk format upgrades must be run in order.
// Note: Disk format upgrades must be run in order of database version.

vec![
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)),
]
}

/// The kind of database format change or validity check we're performing.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DbFormatChange {
Expand Down Expand Up @@ -474,132 +531,39 @@ impl DbFormatChange {
return Ok(());
};

// Note commitment tree de-duplication database upgrade task.

let version_for_pruning_trees =
Version::parse("25.1.1").expect("Hardcoded version string should be valid.");

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

// 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;
// Apply or validate format upgrades
for upgrade in format_upgrades() {
if older_disk_version >= &upgrade.version() {
upgrade
.validate(db, cancel_receiver)?
.expect("failed to validate db format");
continue;
Comment on lines +536 to +540
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a bit confused about this. The original code does not seem to do validate if the disk version >= upgrade version. Is this required or just good to have? What is the performance impact?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change was wrong, I thought some of the upgrades were being validated here even when they weren't being applied, but none are.

What is the performance impact?

It shouldn't have a significant performance impact, it won't block anything. These will still always run at startup anyway but should run after any required format upgrades have been applied.

}

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

// 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)?;
if !upgrade.needs_migration() {
Self::mark_as_upgraded_to(db, &upgrade.version());
continue;
}

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)?;
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.
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");
upgrade
.validate(db, cancel_receiver)?
.expect("db should be 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);
info!(
newer_running_version = ?upgrade.version(),
"Zebra automatically upgraded the database format"
);
Self::mark_as_upgraded_to(db, &upgrade.version());

timer.finish(module_path!(), line!(), "tree keys and caches upgrade");
timer.finish(module_path!(), line!(), upgrade.description());
}

let version_for_upgrading_value_balance_format =
Expand All @@ -610,14 +574,6 @@ impl DbFormatChange {
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:"
Expand Down Expand Up @@ -669,13 +625,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() {
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()");
Expand All @@ -689,66 +641,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.
Expand Down
Loading
Loading