feat: add a sanity check for sync state of rocksdb vs sqlite, minor partitioning changes

CHANGELOG.md

@@ -27,6 +27,7 @@
 - The network monitor now marks the chain as unhealthy if it fails to create new blocks ([#1512](https://github.com/0xMiden/miden-node/pull/1512)).
 - Block producer now detects if it is desync'd from the store's chain tip and aborts ([#1520](https://github.com/0xMiden/miden-node/pull/1520)).
 - Pin tool versions in CI ([#1523](https://github.com/0xMiden/miden-node/pull/1523)).
+- Add check to ensure tree store state is in sync with database storage ([#1532](https://github.com/0xMiden/miden-node/issues/1534)).
 
 ### Changes
 

crates/block-producer/src/server/tests.rs

@@ -94,6 +94,7 @@ async fn block_producer_startup_is_robust_to_network_failures() {
                 ntx_builder_listener,
                 block_producer_listener,
                 data_directory: dir,
+                grpc_timeout: std::time::Duration::from_secs(30),
             }
             .serve()
             .await
@@ -159,6 +160,7 @@ async fn restart_store(
             ntx_builder_listener,
             block_producer_listener,
             data_directory: dir,
+            grpc_timeout: std::time::Duration::from_secs(30),
         }
         .serve()
         .await

crates/store/src/db/mod.rs

@@ -3,8 +3,7 @@ use std::ops::RangeInclusive;
 use std::path::PathBuf;
 
 use anyhow::Context;
-use diesel::prelude::QueryableByName;
-use diesel::{Connection, RunQueryDsl, SqliteConnection};
+use diesel::{Connection, QueryableByName, RunQueryDsl, SqliteConnection};
 use miden_node_proto::domain::account::{AccountInfo, AccountSummary, NetworkAccountPrefix};
 use miden_node_proto::generated as proto;
 use miden_node_utils::tracing::OpenTelemetrySpanExt;

crates/store/src/errors.rs

@@ -206,6 +206,21 @@ pub enum StateInitializationError {
     DatabaseLoadError(#[from] DatabaseSetupError),
     #[error("inner forest error")]
     InnerForestError(#[from] InnerForestError),
+    #[error(
+        "{tree_name} SMT root ({tree_root:?}) does not match expected root from block {block_num} \
+         ({block_root:?}). Delete the tree storage directories and restart the node to rebuild \
+         from the database."
+    )]
+    TreeStorageDiverged {
+        tree_name: &'static str,
+        block_num: BlockNumber,
+        tree_root: Word,
+        block_root: Word,
+    },
+    #[error("public account {0} is missing details in database")]
+    PublicAccountMissingDetails(AccountId),
+    #[error("failed to convert account to delta: {0}")]
+    AccountToDeltaConversionFailed(String),
 }
 
 #[derive(Debug, Error)]

crates/store/src/server/mod.rs

@@ -88,7 +88,8 @@ impl Store {
         let ntx_builder_address = self.ntx_builder_listener.local_addr()?;
         let block_producer_address = self.block_producer_listener.local_addr()?;
         info!(target: COMPONENT, rpc_endpoint=?rpc_address, ntx_builder_endpoint=?ntx_builder_address,
-            block_producer_endpoint=?block_producer_address, ?self.data_directory, ?self.grpc_timeout, "Loading database");
+            block_producer_endpoint=?block_producer_address, ?self.data_directory, ?self.grpc_timeout,
+            "Loading database");
 
         let state =
             Arc::new(State::load(&self.data_directory).await.context("failed to load state")?);

crates/store/src/state/loader.rs

@@ -0,0 +1,295 @@
+//! Tree loading logic for the store state.
+//!
+//! This module handles loading and initializing the Merkle trees (account tree, nullifier tree,
+//! and SMT forest) from storage backends. It supports different loading modes:
+//!
+//! - **Memory mode** (`rocksdb` feature disabled): Trees are rebuilt from the database on each
+//!   startup.
+//! - **Persistent mode** (`rocksdb` feature enabled): Trees are loaded from persistent storage if
+//!   data exists, otherwise rebuilt from the database and persisted.
+
+use std::path::Path;
+
+use miden_protocol::Word;
+use miden_protocol::block::account_tree::account_id_to_smt_key;
+use miden_protocol::block::nullifier_tree::NullifierTree;
+use miden_protocol::block::{BlockHeader, BlockNumber, Blockchain};
+#[cfg(not(feature = "rocksdb"))]
+use miden_protocol::crypto::merkle::smt::MemoryStorage;
+use miden_protocol::crypto::merkle::smt::{LargeSmt, LargeSmtError, SmtStorage};
+#[cfg(feature = "rocksdb")]
+use tracing::info;
+use tracing::instrument;
+#[cfg(feature = "rocksdb")]
+use {
+    miden_crypto::merkle::smt::RocksDbStorage,
+    miden_protocol::crypto::merkle::smt::RocksDbConfig,
+};
+
+use crate::COMPONENT;
+use crate::db::Db;
+use crate::errors::{DatabaseError, StateInitializationError};
+use crate::inner_forest::InnerForest;
+
+// CONSTANTS
+// ================================================================================================
+
+/// Directory name for the account tree storage within the data directory.
+pub const ACCOUNT_TREE_STORAGE_DIR: &str = "accounttree";
+
+/// Directory name for the nullifier tree storage within the data directory.
+pub const NULLIFIER_TREE_STORAGE_DIR: &str = "nullifiertree";
+
+// STORAGE TYPE ALIAS
+// ================================================================================================
+
+/// The storage backend for trees.
+#[cfg(feature = "rocksdb")]
+pub type TreeStorage = RocksDbStorage;
+#[cfg(not(feature = "rocksdb"))]
+pub type TreeStorage = MemoryStorage;
+
+// ERROR CONVERSION
+// ================================================================================================
+
+/// Converts a `LargeSmtError` into a `StateInitializationError`.
+pub fn account_tree_large_smt_error_to_init_error(e: LargeSmtError) -> StateInitializationError {
+    use miden_node_utils::ErrorReport;
+    match e {
+        LargeSmtError::Merkle(merkle_error) => {
+            StateInitializationError::DatabaseError(DatabaseError::MerkleError(merkle_error))
+        },
+        LargeSmtError::Storage(err) => {
+            StateInitializationError::AccountTreeIoError(err.as_report())
+        },
+    }
+}
+
+// STORAGE LOADER TRAIT
+// ================================================================================================
+
+/// Trait for loading trees from storage.
+///
+/// For `MemoryStorage`, the tree is rebuilt from database entries on each startup.
+/// For `RocksDbStorage`, the tree is loaded directly from disk (much faster for large trees).
+///
+/// Missing or corrupted storage is handled by the `verify_tree_consistency` check after loading,
+/// which detects divergence between persistent storage and the database. If divergence is detected,
+/// the user should manually delete the tree storage directories and restart the node.
+pub trait StorageLoader: SmtStorage + Sized {
+    /// Creates a storage backend for the given domain.
+    fn create(data_dir: &Path, domain: &'static str) -> Result<Self, StateInitializationError>;
+
+    /// Loads an account tree, either from persistent storage or by rebuilding from DB.
+    fn load_account_tree(
+        self,
+        db: &mut Db,
+    ) -> impl std::future::Future<Output = Result<LargeSmt<Self>, StateInitializationError>> + Send;
+
+    /// Loads a nullifier tree, either from persistent storage or by rebuilding from DB.
+    fn load_nullifier_tree(
+        self,
+        db: &mut Db,
+    ) -> impl std::future::Future<
+        Output = Result<NullifierTree<LargeSmt<Self>>, StateInitializationError>,
+    > + Send;
+}
+
+// MEMORY STORAGE IMPLEMENTATION
+// ================================================================================================
+
+#[cfg(not(feature = "rocksdb"))]
+impl StorageLoader for MemoryStorage {
+    fn create(_data_dir: &Path, _domain: &'static str) -> Result<Self, StateInitializationError> {
+        Ok(MemoryStorage::default())
+    }
+
+    async fn load_account_tree(
+        self,
+        db: &mut Db,
+    ) -> Result<LargeSmt<Self>, StateInitializationError> {
+        let account_data = db.select_all_account_commitments().await?;
+        let smt_entries = account_data
+            .into_iter()
+            .map(|(id, commitment)| (account_id_to_smt_key(id), commitment));
+        LargeSmt::with_entries(self, smt_entries)
+            .map_err(account_tree_large_smt_error_to_init_error)
+    }
+
+    async fn load_nullifier_tree(
+        self,
+        db: &mut Db,
+    ) -> Result<NullifierTree<LargeSmt<Self>>, StateInitializationError> {
+        let nullifiers = db.select_all_nullifiers().await?;
+        let entries = nullifiers.into_iter().map(|info| (info.nullifier, info.block_num));
+        NullifierTree::with_storage_from_entries(self, entries)
+            .map_err(StateInitializationError::FailedToCreateNullifierTree)
+    }
+}
+
+// ROCKSDB STORAGE IMPLEMENTATION
+// ================================================================================================
+
+#[cfg(feature = "rocksdb")]
+impl StorageLoader for RocksDbStorage {
+    fn create(data_dir: &Path, domain: &'static str) -> Result<Self, StateInitializationError> {
+        let storage_path = data_dir.join(domain);
+
+        fs_err::create_dir_all(&storage_path)
+            .map_err(|e| StateInitializationError::AccountTreeIoError(e.to_string()))?;
+        RocksDbStorage::open(RocksDbConfig::new(storage_path))
+            .map_err(|e| StateInitializationError::AccountTreeIoError(e.to_string()))
+    }
+
+    async fn load_account_tree(
+        self,
+        db: &mut Db,
+    ) -> Result<LargeSmt<Self>, StateInitializationError> {
+        // If RocksDB storage has data, load from it directly
+        let has_data = self
+            .has_leaves()
+            .map_err(|e| StateInitializationError::AccountTreeIoError(e.to_string()))?;
+        if has_data {
+            return load_smt(self);
+        }
+
+        info!(target: COMPONENT, "RocksDB account tree storage is empty, populating from SQLite");
+        let account_data = db.select_all_account_commitments().await?;
+        let smt_entries = account_data
+            .into_iter()
+            .map(|(id, commitment)| (account_id_to_smt_key(id), commitment));
+        LargeSmt::with_entries(self, smt_entries)
+            .map_err(account_tree_large_smt_error_to_init_error)
+    }
+
+    async fn load_nullifier_tree(
+        self,
+        db: &mut Db,
+    ) -> Result<NullifierTree<LargeSmt<Self>>, StateInitializationError> {
+        // If RocksDB storage has data, load from it directly
+        let has_data = self
+            .has_leaves()
+            .map_err(|e| StateInitializationError::NullifierTreeIoError(e.to_string()))?;
+        if has_data {
+            let smt = load_smt(self)?;
+            return Ok(NullifierTree::new_unchecked(smt));
+        }
+
+        info!(target: COMPONENT, "RocksDB nullifier tree storage is empty, populating from SQLite");
+        let nullifiers = db.select_all_nullifiers().await?;
+        let entries = nullifiers.into_iter().map(|info| (info.nullifier, info.block_num));
+        NullifierTree::with_storage_from_entries(self, entries)
+            .map_err(StateInitializationError::FailedToCreateNullifierTree)
+    }
+}
+
+// HELPER FUNCTIONS
+// ================================================================================================
+
+/// Loads an SMT from persistent storage.
+#[cfg(feature = "rocksdb")]
+pub fn load_smt<S: SmtStorage>(storage: S) -> Result<LargeSmt<S>, StateInitializationError> {
+    LargeSmt::new(storage).map_err(account_tree_large_smt_error_to_init_error)
+}
+
+// TREE LOADING FUNCTIONS
+// ================================================================================================
+
+/// Loads the blockchain MMR from all block headers in the database.
+#[instrument(target = COMPONENT, skip_all)]
+pub async fn load_mmr(db: &mut Db) -> Result<Blockchain, StateInitializationError> {
+    let block_commitments: Vec<miden_protocol::Word> = db
+        .select_all_block_headers()
+        .await?
+        .iter()
+        .map(BlockHeader::commitment)
+        .collect();
+
+    // SAFETY: We assume the loaded MMR is valid and does not have more than u32::MAX
+    // entries.
+    let chain_mmr = Blockchain::from_mmr_unchecked(block_commitments.into());
+
+    Ok(chain_mmr)
+}
+
+/// Loads SMT forest with storage map and vault Merkle paths for all public accounts.
+#[instrument(target = COMPONENT, skip_all, fields(block_num = %block_num))]
+pub async fn load_smt_forest(
+    db: &mut Db,
+    block_num: BlockNumber,
+) -> Result<InnerForest, StateInitializationError> {
+    use miden_protocol::account::delta::AccountDelta;
+
+    let public_account_ids = db.select_all_public_account_ids().await?;
+
+    // Acquire write lock once for the entire initialization
+    let mut forest = InnerForest::new();
+
+    // Process each account
+    for account_id in public_account_ids {
+        // Get the full account from the database
+        let account_info = db.select_account(account_id).await?;
+        let account = account_info.details.expect("public accounts always have details in DB");
+
+        // Convert the full account to a full-state delta
+        let delta =
+            AccountDelta::try_from(account).expect("accounts from DB should not have seeds");
+
+        // Use the unified update method (will recognize it's a full-state delta)
+        forest.update_account(block_num, &delta)?;
+    }
+
+    Ok(forest)
+}
+
+// CONSISTENCY VERIFICATION
+// ================================================================================================
+
+/// Verifies that tree roots match the expected roots from the latest block header.
+///
+/// This check ensures the database and tree storage (memory or persistent) haven't diverged due to
+/// corruption or incomplete shutdown. When trees are rebuilt from the database, they will naturally
+/// match; when loaded from persistent storage, this catches any inconsistencies.
+///
+/// # Arguments
+/// * `account_tree_root` - Root of the loaded account tree
+/// * `nullifier_tree_root` - Root of the loaded nullifier tree
+/// * `db` - Database connection to fetch the latest block header
+///
+/// # Errors
+/// Returns `StateInitializationError::TreeStorageDiverged` if any root doesn't match.
+#[instrument(target = COMPONENT, skip_all)]
+pub async fn verify_tree_consistency(
+    account_tree_root: Word,
+    nullifier_tree_root: Word,
+    db: &mut Db,
+) -> Result<(), StateInitializationError> {
+    // Fetch the latest block header to get the expected roots
+    let latest_header = db.select_block_header_by_block_num(None).await?;
+
+    let (block_num, expected_account_root, expected_nullifier_root) = latest_header
+        .map(|header| (header.block_num(), header.account_root(), header.nullifier_root()))
+        .unwrap_or_default();
+
+    // Verify account tree root
+    if account_tree_root != expected_account_root {
+        return Err(StateInitializationError::TreeStorageDiverged {
+            tree_name: "Account",
+            block_num,
+            tree_root: account_tree_root,
+            block_root: expected_account_root,
+        });
+    }
+
+    // Verify nullifier tree root
+    if nullifier_tree_root != expected_nullifier_root {
+        return Err(StateInitializationError::TreeStorageDiverged {
+            tree_name: "Nullifier",
+            block_num,
+            tree_root: nullifier_tree_root,
+            block_root: expected_nullifier_root,
+        });
+    }
+
+    Ok(())
+}