zakura-state 11.0.0

//! Non-finalized chain state management as defined by [RFC0005]
//!
//! [RFC0005]: https://zebra.zfnd.org/dev/rfcs/0005-state-updates.html

use std::{
    collections::{BTreeSet, HashMap},
    mem,
    num::NonZeroU64,
    path::{Path, PathBuf},
    sync::Arc,
    time::Instant,
};

use indexmap::IndexMap;
use tokio::sync::watch;
use zakura_chain::{
    block::{self, Hash, Height},
    parameters::Network,
    sprout::{self},
    transparent,
};
use zakura_node_services::task::spawn_erased;

use crate::{
    constants::{MAX_INVALIDATED_BLOCKS, MAX_NON_FINALIZED_CHAIN_FORKS},
    error::ReconsiderError,
    request::{ContextuallyVerifiedBlock, FinalizableBlock},
    service::{check, finalized_state::ZakuraDb, InvalidateError, PreciousError},
    SemanticallyVerifiedBlock, ValidateContextError, WatchReceiver,
};

mod address_transfers;
mod backup;
mod chain;
mod created_utxos;
mod shared_map;

#[cfg(test)]
pub(crate) use backup::MIN_DURATION_BETWEEN_BACKUP_UPDATES;

#[cfg(test)]
mod tests;

pub(crate) use address_transfers::AddressTransfers;
pub(crate) use backup::write_semantically_verified_backup_block;
#[cfg(feature = "internal-bench")]
pub use chain::{benchmark_height_map, HeightMapBenchmarkSample};
pub(crate) use chain::{Chain, SpendingTransactionId};
pub(crate) use created_utxos::CreatedUtxos;
pub(crate) use shared_map::SharedMap;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ContextualMetrics {
    Disabled,
    AllBlocks,
    Mined,
}

impl ContextualMetrics {
    pub(crate) fn for_commit(is_mined: bool) -> Self {
        if is_mined {
            Self::Mined
        } else {
            Self::AllBlocks
        }
    }

    pub(crate) fn record_duration(
        self,
        metric_name: &'static str,
        mined_metric_name: &'static str,
        duration: std::time::Duration,
    ) {
        if self == Self::Disabled {
            return;
        }

        let duration = duration.as_secs_f64();
        metrics::histogram!(metric_name).record(duration);
        if self == Self::Mined {
            metrics::histogram!(mined_metric_name).record(duration);
        }
    }
}

/// The state of the chains in memory, including queued blocks.
///
/// Clones of the non-finalized state contain independent copies of the chains.
/// This is different from `FinalizedState::clone()`,
/// which returns a shared reference to the database.
///
/// Most chain data is clone-on-write using [`Arc`].
pub struct NonFinalizedState {
    // Chain Data
    //
    /// Verified, non-finalized chains, in ascending work order.
    ///
    /// The best chain is [`NonFinalizedState::best_chain()`], or `chain_iter().next()`.
    /// Using `chain_set.last()` or `chain_set.iter().next_back()` is deprecated,
    /// callers should migrate to `chain_iter().next()`.
    chain_set: BTreeSet<Arc<Chain>>,

    /// Blocks that have been invalidated in, and removed from, the non finalized
    /// state.
    invalidated_blocks: IndexMap<Height, Arc<Vec<ContextuallyVerifiedBlock>>>,

    /// The sequence of the latest `preciousblock` call, so later calls win equal-work ties.
    ///
    /// Operator preferences are local and are not restored from the backup after a restart.
    precious_sequence: u64,

    /// The latest `preciousblock` sequence of each preferred block, so chains reverted to that
    /// block get its preference back, and chains with the same tip compare equal.
    ///
    /// Every operation that can remove blocks drops, once it finishes, the preferences of blocks
    /// that are in no chain and not in `invalidated_blocks`. So the map stays bounded by the
    /// non-finalized state, and an evicted or discarded block downloaded again has no preference.
    precious_blocks: HashMap<block::Hash, NonZeroU64>,

    // Configuration
    //
    /// The configured Zcash network.
    pub network: Network,

    // Diagnostics
    //
    /// Configures the non-finalized state to count metrics.
    ///
    /// Used for skipping metrics when testing block proposals
    /// with a commit to a cloned non-finalized state.
    //
    // TODO: make this field private and set it via an argument to NonFinalizedState::new()
    should_count_metrics: bool,
}

impl std::fmt::Debug for NonFinalizedState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut f = f.debug_struct("NonFinalizedState");

        f.field("chain_set", &self.chain_set)
            .field("network", &self.network);

        f.field("should_count_metrics", &self.should_count_metrics);

        f.finish()
    }
}

impl Clone for NonFinalizedState {
    fn clone(&self) -> Self {
        Self {
            chain_set: self.chain_set.clone(),
            network: self.network.clone(),
            invalidated_blocks: self.invalidated_blocks.clone(),
            precious_sequence: self.precious_sequence,
            precious_blocks: self.precious_blocks.clone(),
            should_count_metrics: self.should_count_metrics,
        }
    }
}

impl NonFinalizedState {
    /// Returns a new non-finalized state for `network`.
    pub fn new(network: &Network) -> NonFinalizedState {
        NonFinalizedState {
            chain_set: Default::default(),
            network: network.clone(),
            invalidated_blocks: Default::default(),
            precious_sequence: 0,
            precious_blocks: HashMap::new(),
            should_count_metrics: true,
        }
    }

    /// Writes the current non-finalized state to the backup directory at `backup_dir_path`.
    ///
    /// Reads the existing backup directory contents, writes any blocks that are in the
    /// non-finalized state but missing from the backup, and deletes any backup files that
    /// are no longer present in the non-finalized state.
    ///
    /// This method performs blocking I/O and should only be called from a blocking context.
    pub(crate) fn write_to_backup(&self, backup_dir_path: &Path) {
        let backup_blocks: HashMap<block::Hash, PathBuf> =
            backup::list_backup_dir_entries(backup_dir_path).collect();
        backup::update_non_finalized_state_backup(backup_dir_path, self, backup_blocks);
    }

    /// Accepts an optional path to the non-finalized state backup directory and a handle to the database.
    ///
    /// If a backup directory path is provided:
    /// - Creates a new backup directory at the provided path if none exists,
    /// - Restores non-finalized blocks from the backup directory, if any, and
    /// - Unless `skip_backup_task` is true, spawns a task that updates the non-finalized
    ///   backup cache with the latest non-finalized state sent to the returned watch channel.
    ///
    /// Returns the non-finalized state with a watch channel sender and receiver.
    pub async fn with_backup(
        self,
        backup_dir_path: Option<PathBuf>,
        finalized_state: &ZakuraDb,
        should_restore_backup: bool,
        skip_backup_task: bool,
    ) -> (
        Self,
        watch::Sender<NonFinalizedState>,
        WatchReceiver<NonFinalizedState>,
    ) {
        let with_watch_channel = |non_finalized_state: NonFinalizedState| {
            let (sender, receiver) = watch::channel(non_finalized_state.clone());
            (non_finalized_state, sender, WatchReceiver::new(receiver))
        };

        let Some(backup_dir_path) = backup_dir_path else {
            return with_watch_channel(self);
        };

        if skip_backup_task {
            tracing::info!(
                ?backup_dir_path,
                "restoring non-finalized blocks from backup (sync write mode, backup task skipped)"
            );
        } else {
            tracing::info!(
                ?backup_dir_path,
                "restoring non-finalized blocks from backup and spawning backup task"
            );
        }

        let non_finalized_state = {
            let backup_dir_path = backup_dir_path.clone();
            let finalized_state = finalized_state.clone();
            tokio::task::spawn_blocking(move || {
                // Create a new backup directory if none exists
                std::fs::create_dir_all(&backup_dir_path)
                    .expect("failed to create non-finalized state backup directory");

                if should_restore_backup {
                    backup::restore_backup(self, &backup_dir_path, &finalized_state)
                } else {
                    self
                }
            })
            .await
            .expect("failed to join blocking task")
        };

        let (non_finalized_state, sender, receiver) = with_watch_channel(non_finalized_state);

        if !skip_backup_task {
            spawn_erased(backup::run_backup_task(receiver.clone(), backup_dir_path));
        }

        if !non_finalized_state.is_chain_set_empty() {
            let num_blocks_restored = non_finalized_state
                .best_chain()
                .expect("must have best chain if chain set is not empty")
                .len();

            tracing::info!(
                ?num_blocks_restored,
                "restored blocks from non-finalized backup cache"
            );
        }

        (non_finalized_state, sender, receiver)
    }

    /// Is the internal state of `self` the same as `other`?
    ///
    /// [`Chain`] has a custom [`Eq`] implementation based on proof of work,
    /// which is used to select the best chain. So we can't derive [`Eq`] for [`NonFinalizedState`].
    ///
    /// Unlike the custom trait impl, this method returns `true` if the entire internal state
    /// of two non-finalized states is equal.
    ///
    /// If the internal states are different, it returns `false`,
    /// even if the chains and blocks are equal.
    #[cfg(any(test, feature = "proptest-impl"))]
    #[allow(dead_code)]
    pub fn eq_internal_state(&self, other: &NonFinalizedState) -> bool {
        // this method must be updated every time a consensus-critical field is added to NonFinalizedState
        // (diagnostic fields can be ignored)

        self.chain_set.len() == other.chain_set.len()
            && self
                .chain_set
                .iter()
                .zip(other.chain_set.iter())
                .all(|(self_chain, other_chain)| self_chain.eq_internal_state(other_chain))
            && self.network == other.network
    }

    /// Returns an iterator over the non-finalized chains, with the best chain first.
    //
    // TODO: replace chain_set.iter().rev() with this method
    pub fn chain_iter(&self) -> impl Iterator<Item = &Arc<Chain>> {
        self.chain_set.iter().rev()
    }

    /// Insert `chain` into `self.chain_set`, apply `chain_filter` to the chains,
    /// then limit the number of tracked chains.
    fn insert_with<F>(&mut self, chain: Arc<Chain>, chain_filter: F)
    where
        F: FnOnce(&mut BTreeSet<Arc<Chain>>),
    {
        let inserted_tip = chain.non_finalized_tip_hash();
        self.chain_set.insert(self.with_tip_preference(chain));

        chain_filter(&mut self.chain_set);

        while self.chain_set.len() > MAX_NON_FINALIZED_CHAIN_FORKS {
            // Keep the accepted parent available for a child that could make it best.
            // Mining still uses the best chain, and the total fork limit is unchanged.
            let evicted = self
                .chain_set
                .iter()
                .find(|chain| chain.non_finalized_tip_hash() != inserted_tip)
                .expect("the fork limit leaves room for both the best and inserted chains")
                .clone();
            self.chain_set.remove(&evicted);
        }
    }

    /// Returns `chain` with the `preciousblock` preference of its current tip.
    ///
    /// Chain::cmp reads the preference, so call this only while the chain is outside the set.
    fn with_tip_preference(&self, mut chain: Arc<Chain>) -> Arc<Chain> {
        let tip = chain.non_finalized_tip_hash();
        let precious = self
            .precious_blocks
            .get(&tip)
            .map(|&sequence| (tip, sequence));
        if chain.precious() != precious {
            Arc::make_mut(&mut chain).set_precious(precious);
        }
        chain
    }

    /// Forget preferences for blocks that are no longer in a chain or invalidated.
    ///
    /// Call this only once an operation has finished changing the chains and
    /// `invalidated_blocks`, because blocks can be briefly in neither while it runs.
    fn retain_precious_blocks(&mut self) {
        let stale: Vec<block::Hash> = self
            .precious_blocks
            .keys()
            .copied()
            .filter(|hash| !self.any_chain_contains(hash) && !self.is_invalidated(hash))
            .collect();
        for hash in stale {
            self.precious_blocks.remove(&hash);
        }
    }

    /// Returns true if `hash` is in `invalidated_blocks`.
    fn is_invalidated(&self, hash: &block::Hash) -> bool {
        self.invalidated_blocks
            .values()
            .any(|blocks| blocks.iter().any(|block| block.hash == *hash))
    }

    /// Insert `chain` into `self.chain_set`, then limit the number of tracked chains.
    fn insert(&mut self, chain: Arc<Chain>) {
        self.insert_with(chain, |_ignored_chain| { /* no filter */ })
    }

    /// Blocks removed by a fork eviction, excluding prefixes shared with retained chains.
    /// Only compare states at the same finalized height.
    pub(crate) fn evicted_blocks(&self, after: &Self) -> Vec<block::Hash> {
        let mut evicted = std::collections::HashSet::new();
        for chain in self.chain_iter() {
            for block in chain.blocks.values().rev() {
                if after.any_chain_contains(&block.hash) || !evicted.insert(block.hash) {
                    break;
                }
            }
        }
        let mut evicted: Vec<_> = evicted.into_iter().collect();
        evicted.sort_unstable_by_key(|hash| hash.0);
        evicted
    }

    /// Returns the latest `preciousblock` sequence, and the sequence of each preferred block.
    #[cfg(test)]
    pub(crate) fn precious_preferences(&self) -> (u64, &HashMap<block::Hash, NonZeroU64>) {
        (self.precious_sequence, &self.precious_blocks)
    }

    #[cfg(test)]
    #[allow(dead_code)]
    pub(crate) fn insert_test_chain(&mut self, chain: Arc<Chain>) {
        self.insert(chain);
    }

    /// Finalize the lowest height block in the non-finalized portion of the best
    /// chain and update all side-chains to match.
    pub fn finalize(&mut self) -> FinalizableBlock {
        // Chain::cmp uses work, operator preference, receipt order, and the tip hash.
        // None of these fields has interior mutability.
        // (And when the tip block is dropped for a chain, the chain is also dropped.)
        #[allow(clippy::mutable_key_type)]
        let chains = mem::take(&mut self.chain_set);
        let mut chains = chains.into_iter();

        // extract best chain
        let mut best_chain = chains.next_back().expect("there's at least one chain");

        // clone if required
        let mut_best_chain = Arc::make_mut(&mut best_chain);

        // extract the rest into side_chains so they can be mutated
        let side_chains = chains;

        // Pop the lowest height block from the best chain to be finalized, and
        // also obtain its associated treestate.
        let (best_chain_root, root_treestate) = mut_best_chain.pop_root();

        // add best_chain back to `self.chain_set`
        if !best_chain.is_empty() {
            self.insert(best_chain);
        }

        // for each remaining chain in side_chains
        for mut side_chain in side_chains.rev() {
            if side_chain.non_finalized_root_hash() != best_chain_root.hash {
                // If we popped the root, the chain would be empty or orphaned,
                // so just drop it now.
                drop(side_chain);

                continue;
            }

            // otherwise, the popped root block is the same as the finalizing block

            // clone if required
            let mut_side_chain = Arc::make_mut(&mut side_chain);

            // remove the first block from `chain`
            let (side_chain_root, _treestate) = mut_side_chain.pop_root();
            assert_eq!(side_chain_root.hash, best_chain_root.hash);

            // add the chain back to `self.chain_set`
            if !side_chain.is_empty() {
                self.insert(side_chain);
            }
        }

        // Remove all invalidated_blocks at or below the finalized height
        self.invalidated_blocks
            .retain(|height, _blocks| *height >= best_chain_root.height);
        self.retain_precious_blocks();

        self.update_metrics_for_chains();

        // Add the treestate to the finalized block.
        FinalizableBlock::new(Arc::unwrap_or_clone(best_chain_root), root_treestate)
    }

    /// Commit block to the non-finalized state, on top of:
    /// - an existing chain's tip, or
    /// - a newly forked chain.
    pub fn commit_block(
        &mut self,
        prepared: SemanticallyVerifiedBlock,
        finalized_state: &ZakuraDb,
    ) -> Result<(), ValidateContextError> {
        self.commit_block_with_metrics(prepared, finalized_state, ContextualMetrics::Disabled)
    }

    #[tracing::instrument(
        name = "commit_block",
        level = "debug",
        skip(self, finalized_state, prepared)
    )]
    pub(crate) fn commit_block_with_metrics(
        &mut self,
        prepared: SemanticallyVerifiedBlock,
        finalized_state: &ZakuraDb,
        contextual_metrics: ContextualMetrics,
    ) -> Result<(), ValidateContextError> {
        let parent_hash = prepared.block.header.previous_block_hash;
        let (height, hash) = (prepared.height, prepared.hash);

        let parent_chain_start = Instant::now();
        let parent_chain = self.parent_chain(parent_hash);
        contextual_metrics.record_duration(
            "state.contextual.parent_chain.duration_seconds",
            "state.contextual.mined.parent_chain.duration_seconds",
            parent_chain_start.elapsed(),
        );
        let parent_chain = parent_chain?;

        // If the block is invalid, return the error,
        // and drop the cloned parent Arc, or newly created chain fork.
        let modified_chain =
            self.validate_and_commit(parent_chain, prepared, finalized_state, contextual_metrics)?;

        // If the block is valid:
        // - add the new chain fork or updated chain to the set of recent chains
        // - remove the parent chain, if it was in the chain set
        //   (if it was a newly created fork, it won't be in the chain set)
        self.insert_with(modified_chain, |chain_set| {
            chain_set.retain(|chain| chain.non_finalized_tip_hash() != parent_hash)
        });
        // The fork limit can evict a chain.
        self.retain_precious_blocks();

        self.update_metrics_for_committed_block(height, hash);

        Ok(())
    }

    /// Invalidate block with hash `block_hash` and all descendants from the non-finalized state. Insert
    /// the new chain into the chain_set and discard the previous.
    #[allow(clippy::unwrap_in_result)]
    pub fn invalidate_block(&mut self, block_hash: Hash) -> Result<block::Hash, InvalidateError> {
        let Some(chain) = self.find_chain(|chain| chain.contains_block_hash(block_hash)) else {
            return Err(InvalidateError::BlockNotFound(block_hash));
        };

        let invalidated_blocks = if chain.non_finalized_root_hash() == block_hash {
            self.chain_set
                .retain(|chain| !chain.contains_block_hash(block_hash));
            chain
                .blocks
                .values()
                .map(|block| block.as_ref().clone())
                .collect()
        } else {
            let (new_chain, invalidated_blocks) = chain
                .invalidate_block(block_hash)
                .expect("already checked that chain contains hash");

            // Add the new chain fork or updated chain to the set of recent chains, and
            // remove the chain containing the hash of the block from chain set
            self.insert_with(Arc::new(new_chain), |chain_set| {
                chain_set.retain(|c| !c.contains_block_hash(block_hash))
            });

            invalidated_blocks
        };

        // TODO: Allow for invalidating multiple block hashes at a given height (#9552).
        self.invalidated_blocks.insert(
            invalidated_blocks
                .first()
                .expect("should not be empty")
                .clone()
                .height,
            Arc::new(invalidated_blocks),
        );

        while self.invalidated_blocks.len() > MAX_INVALIDATED_BLOCKS {
            self.invalidated_blocks.shift_remove_index(0);
        }
        // The new record keeps the removed blocks' preferences, but a replaced or dropped
        // record, or a fork-limit eviction, can discard blocks.
        self.retain_precious_blocks();

        self.update_metrics_for_chains();

        Ok(block_hash)
    }

    /// Prefers the chain tip with hash `block_hash` over every other tip with the same work, as if
    /// it had been received first.
    ///
    /// A later call overrides an earlier one. Greater work still wins, so a block with less work
    /// than the best tip, including a block below another chain tip, an invalidated block, or a
    /// finalized block, leaves the state unchanged.
    ///
    /// The preference lasts while the block is in a chain or an invalidation record, so a block
    /// that is evicted or discarded and then downloaded again is not preferred.
    ///
    /// Returns [`PreciousError::BlockNotFound`] if the block is not known to either state.
    /// Invalidated blocks are known, like Bitcoin Core's block index.
    pub fn precious_block(
        &mut self,
        block_hash: block::Hash,
        finalized_state: &ZakuraDb,
    ) -> Result<(), PreciousError> {
        let Some(chain) = self
            .chain_set
            .iter()
            .find(|chain| chain.non_finalized_tip_hash() == block_hash)
            .cloned()
        else {
            if self.any_chain_contains(&block_hash)
                || self.is_invalidated(&block_hash)
                || finalized_state.height(block_hash).is_some()
            {
                return Ok(());
            }
            return Err(PreciousError::BlockNotFound(block_hash));
        };

        let best_work = self
            .best_chain()
            .expect("the chain set contains the chain found above")
            .partial_cumulative_work;
        if chain.partial_cumulative_work < best_work {
            return Ok(());
        }

        self.precious_sequence = self.precious_sequence.saturating_add(1);
        let sequence = NonZeroU64::new(self.precious_sequence)
            .expect("a saturating increment from zero is at least one");
        self.precious_blocks.insert(block_hash, sequence);
        self.retain_precious_blocks();

        // Chain::cmp reads the preference, so change it only while the chain is outside the set.
        // Chains with the same tip compare equal, so the set holds one chain for this tip.
        self.chain_set.remove(&chain);
        let chain = self.with_tip_preference(chain);
        self.chain_set.insert(chain);

        Ok(())
    }

    /// Reconsiders a previously invalidated block and its descendants into the non-finalized state
    /// based on a block_hash. Reconsidered blocks are inserted into the previous chain and re-inserted
    /// into the chain_set.
    #[allow(clippy::unwrap_in_result)]
    pub fn reconsider_block(
        &mut self,
        block_hash: block::Hash,
        finalized_state: &ZakuraDb,
    ) -> Result<Vec<block::Hash>, ReconsiderError> {
        // Locate the record but keep it live until replay succeeds, so a
        // recoverable error can't lose it; it is `shift_remove`d atomically with
        // the insert below.
        let (height, invalidated_blocks) = self
            .invalidated_blocks
            .iter()
            .find_map(|(height, blocks)| {
                if blocks.first()?.hash == block_hash {
                    Some((*height, (**blocks).clone()))
                } else {
                    None
                }
            })
            .ok_or(ReconsiderError::MissingInvalidatedBlock(block_hash))?;

        let invalidated_block_hashes = invalidated_blocks
            .iter()
            .map(|block| block.hash)
            .collect::<Vec<_>>();

        // Find and fork the parent chain of the invalidated_root. Update the parent chain
        // with the invalidated_descendants
        let invalidated_root = invalidated_blocks
            .first()
            .ok_or(ReconsiderError::InvalidatedBlocksEmpty)?;

        let root_parent_hash = invalidated_root.block.header.previous_block_hash;

        // If the parent is the tip of the finalized_state we create a new chain and insert it
        // into the non finalized state
        let chain_result = if root_parent_hash == finalized_state.finalized_tip_hash() {
            let chain = Chain::new(
                &self.network,
                finalized_state
                    .finalized_tip_height()
                    .ok_or(ReconsiderError::ParentChainNotFound(block_hash))?,
                finalized_state.sprout_tree_for_tip()?,
                finalized_state.sapling_tree_for_tip(),
                finalized_state.orchard_tree_for_tip(),
                finalized_state.ironwood_tree_for_tip(),
                finalized_state.history_tree(),
                finalized_state.finalized_value_pool(),
            );
            Arc::new(chain)
        } else {
            // The parent is not the finalized_tip and still exist in the NonFinalizedState
            // or else we return an error due to the parent not existing in the NonFinalizedState
            self.parent_chain(root_parent_hash)
                .map_err(|_| ReconsiderError::ParentChainNotFound(block_hash))?
        };

        let mut modified_chain = Arc::unwrap_or_clone(chain_result);
        for block in invalidated_blocks {
            modified_chain = modified_chain
                .push(block)
                .map_err(ReconsiderError::ReplayFailed)?;
        }

        let (tip_height, tip_hash) = modified_chain.non_finalized_tip();

        // All fallible steps have succeeded; remove the invalidation record
        // atomically with installing the restored chain so a failed attempt
        // does not destroy the record.
        self.invalidated_blocks.shift_remove(&height);

        // Only track invalidated_blocks that are not yet finalized. Once blocks are finalized (below the best_chain_root_height)
        // we can discard the block.
        if let Some(best_chain_root_height) = finalized_state.finalized_tip_height() {
            self.invalidated_blocks
                .retain(|height, _blocks| *height >= best_chain_root_height);
        }

        self.insert_with(Arc::new(modified_chain), |chain_set| {
            chain_set.retain(|chain| chain.non_finalized_tip_hash() != root_parent_hash)
        });
        // The fork limit can evict a chain, and finalized records were dropped above.
        self.retain_precious_blocks();

        self.update_metrics_for_committed_block(tip_height, tip_hash);

        Ok(invalidated_block_hashes)
    }

    /// Commit block to the non-finalized state as a new chain where its parent
    /// is the finalized tip.
    #[allow(clippy::unwrap_in_result)]
    pub fn commit_new_chain(
        &mut self,
        prepared: SemanticallyVerifiedBlock,
        finalized_state: &ZakuraDb,
    ) -> Result<(), ValidateContextError> {
        self.commit_new_chain_with_metrics(prepared, finalized_state, ContextualMetrics::Disabled)
    }

    #[tracing::instrument(
        name = "commit_new_chain",
        level = "debug",
        skip(self, finalized_state, prepared)
    )]
    #[allow(clippy::unwrap_in_result)]
    pub(crate) fn commit_new_chain_with_metrics(
        &mut self,
        prepared: SemanticallyVerifiedBlock,
        finalized_state: &ZakuraDb,
        contextual_metrics: ContextualMetrics,
    ) -> Result<(), ValidateContextError> {
        let chain_new_start = Instant::now();
        let chain: Result<Chain, ValidateContextError> = (|| {
            let finalized_tip_height = finalized_state.finalized_tip_height();

            // TODO: fix tests that don't initialize the finalized state
            #[cfg(not(test))]
            let finalized_tip_height =
                finalized_tip_height.expect("finalized state contains blocks");
            #[cfg(test)]
            let finalized_tip_height =
                finalized_tip_height.unwrap_or(zakura_chain::block::Height(0));

            Ok(Chain::new(
                &self.network,
                finalized_tip_height,
                finalized_state.sprout_tree_for_tip()?,
                finalized_state.sapling_tree_for_tip(),
                finalized_state.orchard_tree_for_tip(),
                finalized_state.ironwood_tree_for_tip(),
                finalized_state.history_tree(),
                finalized_state.finalized_value_pool(),
            ))
        })();
        contextual_metrics.record_duration(
            "state.contextual.chain_new.duration_seconds",
            "state.contextual.mined.chain_new.duration_seconds",
            chain_new_start.elapsed(),
        );
        let chain = chain?;

        let (height, hash) = (prepared.height, prepared.hash);

        // If the block is invalid, return the error, and drop the newly created chain fork
        let chain = self.validate_and_commit(
            Arc::new(chain),
            prepared,
            finalized_state,
            contextual_metrics,
        )?;

        // If the block is valid, add the new chain fork to the set of recent chains.
        self.insert(chain);
        // The fork limit can evict a chain.
        self.retain_precious_blocks();
        self.update_metrics_for_committed_block(height, hash);

        Ok(())
    }

    /// Contextually validate `prepared` using `finalized_state`.
    /// If validation succeeds, push `prepared` onto `new_chain`.
    ///
    /// `new_chain` should start as a clone of the parent chain fork,
    /// or the finalized tip.
    #[tracing::instrument(level = "debug", skip(self, finalized_state, new_chain))]
    fn validate_and_commit(
        &self,
        new_chain: Arc<Chain>,
        mut prepared: SemanticallyVerifiedBlock,
        finalized_state: &ZakuraDb,
        contextual_metrics: ContextualMetrics,
    ) -> Result<Arc<Chain>, ValidateContextError> {
        // Re-delivery cannot change the ordering key of a retained block.
        if let Some(chain) = self.find_chain(|chain| chain.contains_block_hash(prepared.hash)) {
            prepared.receipt_order = chain
                .block(prepared.hash.into())
                .expect("the selected chain contains this block")
                .receipt_order;
        }
        if self
            .invalidated_blocks
            .values()
            .any(|blocks| blocks.iter().any(|block| block.hash == prepared.hash))
        {
            return Err(ValidateContextError::BlockPreviouslyInvalidated {
                block_hash: prepared.hash,
            });
        }

        let transparent_spend_start = Instant::now();
        let spent_utxos = check::utxo::transparent_spend(
            &prepared,
            &new_chain.created_utxos,
            &new_chain.spent_utxos,
            finalized_state,
        );
        contextual_metrics.record_duration(
            "state.contextual.transparent_spend.duration_seconds",
            "state.contextual.mined.transparent_spend.duration_seconds",
            transparent_spend_start.elapsed(),
        );
        let spent_utxos = spent_utxos?;

        // Reads from disk
        let shielded_anchor_start = Instant::now();
        let shielded_anchors =
            check::anchors::block_sapling_orchard_ironwood_anchors_refer_to_final_treestates(
                finalized_state,
                &new_chain,
                &prepared,
            );
        contextual_metrics.record_duration(
            "state.contextual.shielded_anchors.duration_seconds",
            "state.contextual.mined.shielded_anchors.duration_seconds",
            shielded_anchor_start.elapsed(),
        );
        shielded_anchors?;

        // Reads from disk
        let sprout_anchor_fetch_start = Instant::now();
        let sprout_final_treestates = check::anchors::block_fetch_sprout_final_treestates(
            finalized_state,
            &new_chain,
            &prepared,
        );
        contextual_metrics.record_duration(
            "state.contextual.sprout_anchor_fetch.duration_seconds",
            "state.contextual.mined.sprout_anchor_fetch.duration_seconds",
            sprout_anchor_fetch_start.elapsed(),
        );

        // Quick check that doesn't read from disk
        let contextual_block_start = Instant::now();
        let height = prepared.height;
        let block_hash = prepared.hash;
        let transaction_count = prepared.block.transactions.len();
        let spent_utxo_count = spent_utxos.len();
        let contextual = ContextuallyVerifiedBlock::with_block_and_spent_utxos(
            &self.network,
            prepared,
            spent_utxos,
        )
        .map_err(|value_balance_error| {
            ValidateContextError::CalculateBlockChainValueChange {
                value_balance_error,
                height,
                block_hash,
                transaction_count,
                spent_utxo_count,
            }
        });
        contextual_metrics.record_duration(
            "state.contextual.block_construction.duration_seconds",
            "state.contextual.mined.block_construction.duration_seconds",
            contextual_block_start.elapsed(),
        );
        let contextual = contextual?;

        let parallel_update_start = Instant::now();
        let result = Self::validate_and_update_parallel(
            new_chain,
            contextual,
            sprout_final_treestates,
            contextual_metrics,
        );
        contextual_metrics.record_duration(
            "state.contextual.parallel_update.duration_seconds",
            "state.contextual.mined.parallel_update.duration_seconds",
            parallel_update_start.elapsed(),
        );
        result
    }

    /// Validate `contextual` and update `new_chain`, doing CPU-intensive work in parallel batches.
    #[allow(clippy::unwrap_in_result)]
    #[tracing::instrument(
        skip(new_chain, contextual, sprout_final_treestates),
        fields(height = ?contextual.height, hash = %contextual.hash)
    )]
    fn validate_and_update_parallel(
        new_chain: Arc<Chain>,
        contextual: ContextuallyVerifiedBlock,
        sprout_final_treestates: HashMap<sprout::tree::Root, Arc<sprout::tree::NoteCommitmentTree>>,
        contextual_metrics: ContextualMetrics,
    ) -> Result<Arc<Chain>, ValidateContextError> {
        let mut block_commitment_result = None;
        let mut sprout_anchor_result = None;
        let mut chain_push_result = None;

        // Clone function arguments for different threads
        let block = contextual.block.clone();
        let network = new_chain.network();
        let history_tree = new_chain.history_block_commitment_tree();

        let block2 = contextual.block.clone();
        let height = contextual.height;
        let transaction_hashes = contextual.transaction_hashes.clone();

        rayon::in_place_scope_fifo(|scope| {
            scope.spawn_fifo(|_scope| {
                let start = Instant::now();
                let result = check::block_commitment_is_valid_for_chain_history(
                    block,
                    &network,
                    &history_tree,
                    None,
                );
                block_commitment_result = Some((result, start.elapsed()));
            });

            scope.spawn_fifo(|_scope| {
                let start = Instant::now();
                let result = check::anchors::block_sprout_anchors_refer_to_treestates(
                    sprout_final_treestates,
                    block2,
                    transaction_hashes,
                    height,
                );
                sprout_anchor_result = Some((result, start.elapsed()));
            });

            // We're pretty sure the new block is valid,
            // so clone the inner chain if needed, then add the new block.
            //
            // Pushing a block onto a Chain can launch additional parallel batches.
            // TODO: should we pass _scope into Chain::push()?
            scope.spawn_fifo(|_scope| {
                let chain_clone_start = Instant::now();
                let new_chain = Arc::unwrap_or_clone(new_chain);
                let chain_clone_duration = chain_clone_start.elapsed();

                let chain_push_start = Instant::now();
                let result = new_chain.push(contextual).map(Arc::new);
                chain_push_result =
                    Some((result, chain_clone_duration, chain_push_start.elapsed()));
            });
        });

        // Don't return the updated Chain unless all the parallel results were Ok
        let (block_commitment_result, block_commitment_duration) =
            block_commitment_result.expect("scope has finished");
        let (sprout_anchor_result, sprout_anchor_duration) =
            sprout_anchor_result.expect("scope has finished");
        let (chain_push_result, chain_clone_duration, chain_push_duration) =
            chain_push_result.expect("scope has finished");

        // These task durations overlap. Only `parallel_update` measures their
        // combined critical-path wall time.
        contextual_metrics.record_duration(
            "state.contextual.parallel_task.block_commitment.duration_seconds",
            "state.contextual.mined.parallel_task.block_commitment.duration_seconds",
            block_commitment_duration,
        );
        contextual_metrics.record_duration(
            "state.contextual.parallel_task.sprout_anchor_check.duration_seconds",
            "state.contextual.mined.parallel_task.sprout_anchor_check.duration_seconds",
            sprout_anchor_duration,
        );
        contextual_metrics.record_duration(
            "state.contextual.parallel_task.chain_clone.duration_seconds",
            "state.contextual.mined.parallel_task.chain_clone.duration_seconds",
            chain_clone_duration,
        );
        contextual_metrics.record_duration(
            "state.contextual.parallel_task.chain_push.duration_seconds",
            "state.contextual.mined.parallel_task.chain_push.duration_seconds",
            chain_push_duration,
        );

        block_commitment_result?;
        sprout_anchor_result?;

        chain_push_result
    }

    /// Returns the length of the non-finalized portion of the current best chain
    /// or `None` if the best chain has no blocks.
    pub fn best_chain_len(&self) -> Option<u32> {
        // This `as` can't overflow because the number of blocks in the chain is limited to i32::MAX,
        // and the non-finalized chain is further limited by the rollback window
        // (`MAX_BLOCK_REORG_HEIGHT`, currently 1000 blocks).
        Some(self.best_chain()?.blocks.len() as u32)
    }

    /// Returns the root height of the non-finalized state, if the non-finalized state is not empty.
    pub fn root_height(&self) -> Option<block::Height> {
        self.best_chain()
            .map(|chain| chain.non_finalized_root_height())
    }

    /// Returns `true` if `hash` is contained in the non-finalized portion of any
    /// known chain.
    #[allow(dead_code)]
    pub fn any_chain_contains(&self, hash: &block::Hash) -> bool {
        self.chain_set
            .iter()
            .rev()
            .any(|chain| chain.height_by_hash.contains_key(hash))
    }

    /// Returns the first chain satisfying the given predicate.
    ///
    /// If multiple chains satisfy the predicate, returns the chain with the highest difficulty.
    /// (Using operator preference, receipt order, then tip hash to break ties.)
    pub fn find_chain<P>(&self, mut predicate: P) -> Option<Arc<Chain>>
    where
        P: FnMut(&Chain) -> bool,
    {
        // Reverse the iteration order, to find highest difficulty chains first.
        self.chain_set
            .iter()
            .rev()
            .find(|chain| predicate(chain))
            .cloned()
    }

    /// Returns the [`transparent::Utxo`] pointed to by the given
    /// [`transparent::OutPoint`] if it is present in any chain.
    ///
    /// UTXOs are returned regardless of whether they have been spent.
    pub fn any_utxo(&self, outpoint: &transparent::OutPoint) -> Option<transparent::Utxo> {
        self.chain_set
            .iter()
            .rev()
            .find_map(|chain| chain.created_utxo(outpoint))
    }

    /// Returns the tip of the best chain.
    #[allow(dead_code)]
    pub fn best_tip(&self) -> Option<(block::Height, block::Hash)> {
        let best_chain = self.best_chain()?;
        let height = best_chain.non_finalized_tip_height();
        let hash = best_chain.non_finalized_tip_hash();

        Some((height, hash))
    }

    /// Returns the block at the tip of the best chain.
    #[allow(dead_code)]
    pub fn best_tip_block(&self) -> Option<&ContextuallyVerifiedBlock> {
        let best_chain = self.best_chain()?;

        best_chain.tip_block()
    }

    /// Returns `true` if the best chain contains `sprout_nullifier`.
    #[cfg(test)]
    pub fn best_contains_sprout_nullifier(&self, sprout_nullifier: &sprout::Nullifier) -> bool {
        self.best_chain()
            .map(|best_chain| best_chain.sprout_nullifiers.contains_key(sprout_nullifier))
            .unwrap_or(false)
    }

    /// Returns `true` if the best chain contains `sapling_nullifier`.
    #[cfg(test)]
    pub fn best_contains_sapling_nullifier(
        &self,
        sapling_nullifier: &zakura_chain::sapling::Nullifier,
    ) -> bool {
        self.best_chain()
            .map(|best_chain| {
                best_chain
                    .sapling_nullifiers
                    .contains_key(sapling_nullifier)
            })
            .unwrap_or(false)
    }

    /// Returns `true` if the best chain contains `orchard_nullifier`.
    #[cfg(test)]
    pub fn best_contains_orchard_nullifier(
        &self,
        orchard_nullifier: &zakura_chain::orchard::Nullifier,
    ) -> bool {
        self.best_chain()
            .map(|best_chain| {
                best_chain
                    .orchard_nullifiers
                    .contains_key(orchard_nullifier)
            })
            .unwrap_or(false)
    }

    /// Return the non-finalized portion of the current best chain.
    pub fn best_chain(&self) -> Option<&Arc<Chain>> {
        self.chain_iter().next()
    }

    /// Return the number of chains.
    pub fn chain_count(&self) -> usize {
        self.chain_set.len()
    }

    /// Returns true if this [`NonFinalizedState`] contains no chains.
    pub fn is_chain_set_empty(&self) -> bool {
        self.chain_count() == 0
    }

    /// Return the invalidated blocks.
    pub fn invalidated_blocks(&self) -> IndexMap<Height, Arc<Vec<ContextuallyVerifiedBlock>>> {
        self.invalidated_blocks.clone()
    }

    /// Return the chain whose tip block hash is `parent_hash`.
    ///
    /// The chain can be an existing chain in the non-finalized state, or a freshly
    /// created fork.
    fn parent_chain(&self, parent_hash: block::Hash) -> Result<Arc<Chain>, ValidateContextError> {
        match self.find_chain(|chain| chain.non_finalized_tip_hash() == parent_hash) {
            // Clone the existing Arc<Chain> in the non-finalized state
            Some(chain) => Ok(chain.clone()),
            // Create a new fork
            None => {
                // Check the lowest difficulty chains first,
                // because the fork could be closer to their tip.
                let fork_chain = self
                    .chain_set
                    .iter()
                    .rev()
                    .find_map(|chain| chain.fork(parent_hash))
                    .ok_or(ValidateContextError::NotReadyToBeCommitted)?;

                Ok(Arc::new(fork_chain))
            }
        }
    }

    /// Should this `NonFinalizedState` instance track metrics?
    fn should_count_metrics(&self) -> bool {
        self.should_count_metrics
    }

    /// Update the metrics after `block` is committed
    fn update_metrics_for_committed_block(&self, height: block::Height, hash: block::Hash) {
        if !self.should_count_metrics() {
            return;
        }

        metrics::counter!("state.memory.committed.block.count").increment(1);
        metrics::gauge!("state.memory.committed.block.height").set(height.0 as f64);

        if self
            .best_chain()
            .expect("metrics are only updated after initialization")
            .non_finalized_tip_hash()
            == hash
        {
            metrics::counter!("state.memory.best.committed.block.count").increment(1);
            metrics::gauge!("state.memory.best.committed.block.height").set(height.0 as f64);
        }

        self.update_metrics_for_chains();
    }

    /// Update the metrics after `self.chain_set` is modified
    fn update_metrics_for_chains(&self) {
        if !self.should_count_metrics() {
            return;
        }

        metrics::gauge!("state.memory.chain.count").set(self.chain_set.len() as f64);
        metrics::gauge!("state.memory.best.chain.length",)
            .set(self.best_chain_len().unwrap_or_default() as f64);
    }

    /// Stop tracking metrics for this non-finalized state and all its chains.
    pub fn disable_metrics(&mut self) {
        self.should_count_metrics = false;
    }
}

impl Drop for NonFinalizedState {
    fn drop(&mut self) {
        self.disable_metrics();
    }
}