reth_cli_commands/db/checksum/

mod.rs

use crate::{
    common::CliNodeTypes,
    db::get::{maybe_json_value_parser, table_key},
};
use alloy_primitives::map::foldhash::fast::FixedState;
use clap::Parser;
use itertools::Itertools;
use reth_chainspec::EthereumHardforks;
use reth_db::{static_file::iter_static_files, DatabaseEnv};
use reth_db_api::{
    cursor::DbCursorRO,
    table::{Decode, Table},
    transaction::DbTx,
    DatabaseError, RawKey, RawTable, TableViewer, Tables,
};
use reth_db_common::DbTool;
use reth_node_builder::{NodeTypesWithDB, NodeTypesWithDBAdapter};
use reth_provider::{providers::ProviderNodeTypes, DBProvider, StaticFileProviderFactory};
use reth_static_file_types::{ChangesetOffset, StaticFileSegment};
use std::{
    hash::{BuildHasher, Hasher},
    ops::Bound,
    time::{Duration, Instant},
};
use tracing::{info, warn};

mod rocksdb;

/// Interval for logging progress during checksum computation.
const PROGRESS_LOG_INTERVAL: usize = 100_000;

#[derive(Parser, Debug)]
/// The arguments for the `reth db checksum` command
pub struct Command {
    #[command(subcommand)]
    subcommand: Subcommand,
}

#[derive(clap::Subcommand, Debug)]
enum Subcommand {
    /// Calculates the checksum of a database table
    Mdbx {
        /// The table name
        table: Tables,

        /// The start of the range to checksum.
        #[arg(long, value_parser = maybe_json_value_parser)]
        start_key: Option<String>,

        /// The end of the range to checksum.
        #[arg(long, value_parser = maybe_json_value_parser)]
        end_key: Option<String>,

        /// The maximum number of records that are queried and used to compute the
        /// checksum.
        #[arg(long)]
        limit: Option<usize>,
    },
    /// Calculates the checksum of a static file segment
    StaticFile {
        /// The static file segment
        #[arg(value_enum)]
        segment: StaticFileSegment,

        /// The block number to start from (inclusive).
        #[arg(long)]
        start_block: Option<u64>,

        /// The block number to end at (inclusive).
        #[arg(long)]
        end_block: Option<u64>,

        /// The maximum number of rows to checksum.
        #[arg(long)]
        limit: Option<usize>,
    },
    /// Calculates the checksum of a RocksDB table
    Rocksdb {
        /// The RocksDB table
        #[arg(value_enum)]
        table: rocksdb::RocksDbTable,

        /// The maximum number of records to checksum.
        #[arg(long)]
        limit: Option<usize>,
    },
}

impl Command {
    /// Execute `db checksum` command
    pub fn execute<N: CliNodeTypes<ChainSpec: EthereumHardforks>>(
        self,
        tool: &DbTool<NodeTypesWithDBAdapter<N, DatabaseEnv>>,
    ) -> eyre::Result<()> {
        warn!("This command should be run without the node running!");

        match self.subcommand {
            Subcommand::Mdbx { table, start_key, end_key, limit } => {
                table.view(&ChecksumViewer { tool, start_key, end_key, limit })?;
            }
            Subcommand::StaticFile { segment, start_block, end_block, limit } => {
                checksum_static_file(tool, segment, start_block, end_block, limit)?;
            }
            Subcommand::Rocksdb { table, limit } => {
                rocksdb::checksum_rocksdb(tool, table, limit)?;
            }
        }

        Ok(())
    }
}

/// Creates a new hasher with the standard seed used for checksum computation.
fn checksum_hasher() -> impl Hasher {
    FixedState::with_seed(u64::from_be_bytes(*b"RETHRETH")).build_hasher()
}

fn checksum_static_file<N: CliNodeTypes<ChainSpec: EthereumHardforks>>(
    tool: &DbTool<NodeTypesWithDBAdapter<N, DatabaseEnv>>,
    segment: StaticFileSegment,
    start_block: Option<u64>,
    end_block: Option<u64>,
    limit: Option<usize>,
) -> eyre::Result<()> {
    let static_file_provider = tool.provider_factory.static_file_provider();
    if let Err(err) = static_file_provider.check_consistency(&tool.provider_factory.provider()?) {
        warn!("Error checking consistency of static files: {err}");
    }

    let static_files = iter_static_files(static_file_provider.directory())?;

    let ranges = static_files
        .get(segment)
        .ok_or_else(|| eyre::eyre!("No static files found for segment: {}", segment))?;

    let start_time = Instant::now();
    let limit = limit.unwrap_or(usize::MAX);
    let mut checksummer = Checksummer::new(checksum_hasher(), limit);

    let start_block = start_block.unwrap_or(0);
    let end_block = end_block.unwrap_or(u64::MAX);
    let is_change_based = segment.is_change_based();

    info!(
        "Computing checksum for {} static files, start_block={}, end_block={}, limit={:?}",
        segment,
        start_block,
        end_block,
        if limit == usize::MAX { None } else { Some(limit) }
    );

    let mut reached_limit = false;
    for (block_range, _header) in ranges.iter().sorted_by_key(|(range, _)| range.start()) {
        if block_range.end() < start_block || block_range.start() > end_block {
            continue;
        }

        let fixed_block_range = static_file_provider.find_fixed_range(segment, block_range.start());
        let jar_provider = static_file_provider
            .get_segment_provider_for_range(segment, || Some(fixed_block_range), None)?
            .ok_or_else(|| {
                eyre::eyre!(
                    "Failed to get segment provider for segment {} at range {}",
                    segment,
                    block_range
                )
            })?;

        let mut cursor = jar_provider.cursor()?;

        if is_change_based {
            let offsets = jar_provider.read_changeset_offsets()?.ok_or_else(|| {
                eyre::eyre!(
                    "Missing changeset offsets sidecar for segment {} at range {}",
                    segment,
                    block_range
                )
            })?;
            let input = ChangeBasedChecksumInput {
                segment,
                block_range_start: block_range.start(),
                start_block,
                end_block,
                offsets: &offsets,
            };

            reached_limit = checksum_change_based_segment(&mut checksummer, input, &mut cursor)?;
        } else {
            while let Some(row) = cursor.next_row()? {
                if checksummer.write_row(&row) {
                    reached_limit = true;
                    break;
                }
            }
        }

        // Explicitly drop provider before removing from cache to avoid deadlock
        drop(jar_provider);
        static_file_provider.remove_cached_provider(segment, fixed_block_range.end());

        if reached_limit {
            break;
        }
    }

    let (checksum, total) = checksummer.finish();
    let elapsed = start_time.elapsed();

    info!(
        "Checksum for static file segment `{}`: {:#x} ({} entries, elapsed: {:?})",
        segment, checksum, total, elapsed
    );

    Ok(())
}

pub(crate) struct ChecksumViewer<'a, N: NodeTypesWithDB> {
    tool: &'a DbTool<N>,
    start_key: Option<String>,
    end_key: Option<String>,
    limit: Option<usize>,
}

impl<N: NodeTypesWithDB> ChecksumViewer<'_, N> {
    pub(crate) const fn new(tool: &'_ DbTool<N>) -> ChecksumViewer<'_, N> {
        ChecksumViewer { tool, start_key: None, end_key: None, limit: None }
    }
}

impl<N: ProviderNodeTypes> TableViewer<(u64, Duration)> for ChecksumViewer<'_, N> {
    type Error = eyre::Report;

    fn view<T: Table>(&self) -> Result<(u64, Duration), Self::Error> {
        let provider =
            self.tool.provider_factory.provider()?.disable_long_read_transaction_safety();
        let tx = provider.tx_ref();
        log_checksum_start(&self.start_key, &self.end_key, self.limit);

        let mut cursor = tx.cursor_read::<RawTable<T>>()?;
        let parse_key = |key: &str| table_key::<T>(key).map(RawKey::new);
        let start = self.start_key.as_deref().map(parse_key).transpose()?;
        let end = self.end_key.as_deref().map(parse_key).transpose()?;
        let bound = |key: Option<RawKey<T::Key>>| key.map_or(Bound::Unbounded, Bound::Included);
        let mut entries = cursor
            .walk_range((bound(start), bound(end)))?
            .map(|entry| entry.map(|(k, v)| (k.into_key(), v.into_value())));

        checksum_entries(
            &mut entries,
            self.limit.unwrap_or(usize::MAX),
            &|key| Ok(serde_json::to_string(&T::Key::decode(key)?).unwrap_or_default()),
            T::NAME,
        )
    }
}

fn log_checksum_start(start: &Option<String>, end: &Option<String>, limit: Option<usize>) {
    info!("Start computing checksum, start={:?}, end={:?}, limit={:?}", start, end, limit);
}

/// Hashes the raw key-value entries of a table and logs the result.
///
/// `key_json` serializes a raw key to JSON, as it can only be decoded with the table's key type.
fn checksum_entries(
    entries: &mut dyn Iterator<Item = Result<(Vec<u8>, Vec<u8>), DatabaseError>>,
    limit: usize,
    key_json: &dyn Fn(&[u8]) -> eyre::Result<String>,
    table: &str,
) -> eyre::Result<(u64, Duration)> {
    let start_time = Instant::now();
    let mut hasher = checksum_hasher();
    let mut total = 0;

    let mut enumerate_start_key = None;
    let mut enumerate_end_key = None;
    for (index, entry) in entries.enumerate() {
        let (k, v) = entry?;

        if index.is_multiple_of(PROGRESS_LOG_INTERVAL) {
            info!("Hashed {index} entries.");
        }

        hasher.write(&k);
        hasher.write(&v);

        if enumerate_start_key.is_none() {
            enumerate_start_key = Some(k.clone());
        }
        enumerate_end_key = Some(k);

        total = index + 1;
        if total >= limit {
            break;
        }
    }

    info!("Hashed {total} entries.");
    if let (Some(s), Some(e)) = (enumerate_start_key, enumerate_end_key) {
        info!("start-key: {}", key_json(&s)?);
        info!("end-key: {}", key_json(&e)?);
    }

    let checksum = hasher.finish();
    let elapsed = start_time.elapsed();

    info!("Checksum for table `{}`: {:#x} (elapsed: {:?})", table, checksum, elapsed);

    Ok((checksum, elapsed))
}

/// Accumulates a checksum over key-value entries, tracking count and limit.
struct Checksummer<H> {
    hasher: H,
    total: usize,
    limit: usize,
}

impl<H: Hasher> Checksummer<H> {
    fn new(hasher: H, limit: usize) -> Self {
        Self { hasher, total: 0, limit }
    }

    /// Hash a row's columns (non-changeset segments). Returns `true` if the limit is reached.
    fn write_row(&mut self, row: &[&[u8]]) -> bool {
        for col in row {
            self.hasher.write(col);
        }
        self.advance()
    }

    /// Hash a key + value as two separate writes, matching MDBX raw entry semantics.
    /// Write boundaries matter: foldhash rotates its accumulator by `len` on each `write`.
    fn write_entry(&mut self, key: &[u8], value: &[u8]) -> bool {
        self.hasher.write(key);
        self.hasher.write(value);
        self.advance()
    }

    fn advance(&mut self) -> bool {
        self.total += 1;
        if self.total.is_multiple_of(PROGRESS_LOG_INTERVAL) {
            info!("Hashed {} entries.", self.total);
        }
        self.total >= self.limit
    }

    fn finish(self) -> (u64, usize) {
        (self.hasher.finish(), self.total)
    }
}

/// Reconstruct MDBX `StorageChangeSets` key/value boundaries from a static-file row.
///
/// MDBX layout:
/// - key: `BlockNumberAddress` => `[8B block_number][20B address]`
/// - value: `StorageEntry` => `[32B storage_key][compact U256 value]`
///
/// Static-file row layout for `StorageBeforeTx`:
/// - `[20B address][32B storage_key][compact U256 value]`
fn split_storage_changeset_row(block_number: u64, row: &[u8]) -> eyre::Result<([u8; 28], &[u8])> {
    if row.len() < 20 {
        return Err(eyre::eyre!(
            "Storage changeset row too short: expected at least 20 bytes, got {}",
            row.len()
        ));
    }

    let mut key_buf = [0u8; 28];
    key_buf[..8].copy_from_slice(&block_number.to_be_bytes());
    key_buf[8..].copy_from_slice(&row[..20]);
    Ok((key_buf, &row[20..]))
}

struct ChangeBasedChecksumInput<'a> {
    segment: StaticFileSegment,
    block_range_start: u64,
    start_block: u64,
    end_block: u64,
    offsets: &'a [ChangesetOffset],
}

fn checksum_change_based_segment<H: Hasher>(
    checksummer: &mut Checksummer<H>,
    input: ChangeBasedChecksumInput<'_>,
    cursor: &mut reth_db::static_file::StaticFileCursor<'_>,
) -> eyre::Result<bool> {
    let ChangeBasedChecksumInput { segment, block_range_start, start_block, end_block, offsets } =
        input;
    let is_storage = segment.is_storage_change_sets();
    let mut reached_limit = false;

    for (offset_index, offset) in offsets.iter().enumerate() {
        let block_number = block_range_start + offset_index as u64;
        let include = block_number >= start_block && block_number <= end_block;

        for _ in 0..offset.num_changes() {
            let row = cursor.next_row()?.ok_or_else(|| {
                eyre::eyre!(
                    "Unexpected EOF while checksumming {} static file at range starting {}",
                    segment,
                    block_range_start
                )
            })?;

            if !include {
                continue;
            }

            // Reconstruct MDBX key/value write boundaries. foldhash rotates
            // its accumulator by `len` on each write(), so boundaries must
            // match exactly.
            let done = if is_storage {
                // StorageChangeSets: MDBX key = BlockNumberAddress (28B),
                // value = compact StorageEntry. Column 0 is compact
                // StorageBeforeTx = [20B address][32B key][compact U256].
                let col = row[0];
                let (key, value) = split_storage_changeset_row(block_number, col)?;
                checksummer.write_entry(&key, value)
            } else {
                // AccountChangeSets: MDBX key = BlockNumber (8B),
                // value = compact AccountBeforeTx (= column 0).
                checksummer.write_entry(&block_number.to_be_bytes(), row[0])
            };

            if done {
                reached_limit = true;
                break;
            }
        }

        if reached_limit {
            break;
        }
    }

    if !reached_limit && cursor.next_row()?.is_some() {
        return Err(eyre::eyre!(
            "Changeset offsets do not cover all rows for {} at range starting {}",
            segment,
            block_range_start
        ));
    }

    Ok(reached_limit)
}