unwind.rs

//! Command that unwinds the OP proofs storage to a specific block number.

use clap::Parser;
use reth_cli::chainspec::ChainSpecParser;
use reth_cli_commands::common::{AccessRights, CliNodeTypes, Environment, EnvironmentArgs};
use reth_node_core::version::version_metadata;
use reth_optimism_chainspec::OpChainSpec;
use reth_optimism_primitives::OpPrimitives;
use reth_optimism_node::args::ProofsStorageVersion;
use reth_optimism_trie::{
    OpProofsProviderRO, OpProofsProviderRw, OpProofsStore,
    db::{MdbxProofsStorage, MdbxProofsStorageV2},
};
use reth_provider::{BlockReader, TransactionVariant};
use std::{path::PathBuf, sync::Arc};
use tracing::{info, warn};

/// Unwinds the proofs storage to a specific block number.
///
/// This command removes all proof history and state updates after the target block number.
#[derive(Debug, Parser)]
pub struct UnwindCommand<C: ChainSpecParser> {
    #[command(flatten)]
    env: EnvironmentArgs<C>,

    /// The path to the storage DB for proofs history.
    #[arg(
        long = "proofs-history.storage-path",
        value_name = "PROOFS_HISTORY_STORAGE_PATH",
        required = true
    )]
    pub storage_path: PathBuf,

    /// The target block number to unwind to.
    ///
    /// All history *after* this block will be removed.
    #[arg(long, value_name = "TARGET_BLOCK")]
    pub target: u64,

    /// Storage schema version. Must match the version used when starting the node.
    #[arg(
        long = "proofs-history.storage-version",
        value_name = "PROOFS_HISTORY_STORAGE_VERSION",
        default_value = "v1"
    )]
    pub storage_version: ProofsStorageVersion,
}

impl<C: ChainSpecParser> UnwindCommand<C> {
    /// Validates that the target block number is within a valid range for unwinding.
    fn validate_unwind_range<Store: OpProofsStore>(&self, storage: Store) -> eyre::Result<bool> {
        let provider_ro = storage.provider_ro()?;
        let (Some((earliest, _)), Some((latest, _))) =
            (provider_ro.get_earliest_block_number()?, provider_ro.get_latest_block_number()?)
        else {
            warn!(target: "reth::cli", "No blocks found in proofs storage. Nothing to unwind.");
            return Ok(false);
        };

        if self.target <= earliest {
            warn!(target: "reth::cli", unwind_target = ?self.target, ?earliest, "Target block is less than the earliest block in proofs storage. Nothing to unwind.");
            return Ok(false);
        }

        if self.target > latest {
            warn!(target: "reth::cli", unwind_target = ?self.target, ?latest, "Target block is not less than the latest block in proofs storage. Nothing to unwind.");
            return Ok(false);
        }

        Ok(true)
    }
}

impl<C: ChainSpecParser<ChainSpec = OpChainSpec>> UnwindCommand<C> {
    /// Execute [`UnwindCommand`].
    pub async fn execute<N: CliNodeTypes<ChainSpec = C::ChainSpec, Primitives = OpPrimitives>>(
        self,
        runtime: reth_tasks::Runtime,
    ) -> eyre::Result<()> {
        info!(target: "reth::cli", "reth {} starting", version_metadata().short_version);
        info!(target: "reth::cli", "Unwinding OP proofs storage at: {:?}", self.storage_path);

        // Initialize the environment with read-only access
        let Environment { provider_factory, .. } = self.env.init::<N>(AccessRights::RO, runtime)?;

        macro_rules! run_unwind {
            ($storage:expr) => {{
                let storage = $storage;

                // Validate that the target block is within a valid range for unwinding
                if !self.validate_unwind_range(storage.clone())? {
                    return Ok(());
                }

                // Get the target block from the main database
                let block = provider_factory
                    .recovered_block(self.target.into(), TransactionVariant::NoHash)?
                    .ok_or_else(|| {
                        eyre::eyre!("Target block {} not found in the main database", self.target)
                    })?;

                info!(target: "reth::cli", block_number = block.number, block_hash = %block.hash(), "Unwinding to target block");
                let provider_rw = storage.provider_rw()?;
                provider_rw.unwind_history(block.block_with_parent())?;
                provider_rw.commit()?;
            }};
        }

        match self.storage_version {
            ProofsStorageVersion::V1 => {
                let storage: Arc<MdbxProofsStorage> = Arc::new(
                    MdbxProofsStorage::new(&self.storage_path)
                        .map_err(|e| eyre::eyre!("Failed to create MdbxProofsStorage: {e}"))?,
                );
                run_unwind!(storage);
            }
            ProofsStorageVersion::V2 => {
                let storage: Arc<MdbxProofsStorageV2> = Arc::new(
                    MdbxProofsStorageV2::new(&self.storage_path)
                        .map_err(|e| eyre::eyre!("Failed to create MdbxProofsStorageV2: {e}"))?,
                );
                run_unwind!(storage);
            }
        }

        Ok(())
    }
}

impl<C: ChainSpecParser> UnwindCommand<C> {
    /// Returns the underlying chain being used to run this command
    pub const fn chain_spec(&self) -> Option<&Arc<C::ChainSpec>> {
        Some(&self.env.chain)
    }
}