From d1a6006fda1d1a76b7e60037a60c9ee9afc37c64 Mon Sep 17 00:00:00 2001 From: Bernhard Schuster Date: Mon, 19 Jan 2026 14:10:00 +0100 Subject: [PATCH 1/5] initial --- CHANGELOG.md | 1 + crates/store/src/db/mod.rs | 194 ++++++++-- .../store/src/db/models/queries/accounts.rs | 189 ++++++++++ crates/store/src/db/tests.rs | 337 +++++++++++++++++- crates/store/src/errors.rs | 2 + crates/store/src/inner_forest/mod.rs | 141 ++++++++ crates/store/src/inner_forest/tests.rs | 152 ++++++++ crates/store/src/state.rs | 15 + 8 files changed, 1009 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fee80111c..1da045315 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Enhancements +- Added periodic cleanup of old account data from the database and in-memory forest ([#1304](https://github.com/0xMiden/miden-node/issues/1304)). - Added support for timeouts in the WASM remote prover clients ([#1383](https://github.com/0xMiden/miden-node/pull/1383)). - Added block validation endpoint to validator and integrated with block producer ([#1382](https://github.com/0xMiden/miden-node/pull/1381)). - Added support for caching mempool statistics in the block producer server ([#1388](https://github.com/0xMiden/miden-node/pull/1388)). diff --git a/crates/store/src/db/mod.rs b/crates/store/src/db/mod.rs index a2dacd235..5e68c7e3a 100644 --- a/crates/store/src/db/mod.rs +++ b/crates/store/src/db/mod.rs @@ -48,8 +48,10 @@ pub(crate) mod schema; pub type Result = std::result::Result; +#[derive(Clone)] pub struct Db { pool: deadpool_diesel::Pool>, + notify_cleanup_task: tokio::sync::mpsc::Sender, } /// Describes the value of an asset for an account ID at `block_num` specifically. @@ -305,6 +307,8 @@ impl Db { } /// Open a connection to the DB and apply any pending migrations. + /// + /// This also spawns a background task that handles periodic cleanup of old account data. #[instrument(target = COMPONENT, skip_all)] pub async fn load(database_filepath: PathBuf) -> Result { let manager = ConnectionManager::new(database_filepath.to_str().unwrap()); @@ -316,8 +320,19 @@ impl Db { "Connected to the database" ); - let me = Db { pool }; + // Create channel for cleanup notifications + // Buffer size of 2 is sufficient since cleanup coalesces multiple notifications + let (notify_cleanup_task, rx) = tokio::sync::mpsc::channel(2); + + let me = Db { pool, notify_cleanup_task }; + let me2 = me.clone(); + + // Spawn background cleanup task + let _cleanup_task_handle = + tokio::spawn(async move { Self::periodic_cleanup_task(me2, rx).await }); + me.query("migrations", apply_migrations).await?; + Ok(me) } @@ -566,30 +581,167 @@ impl Db { block: ProvenBlock, notes: Vec<(NoteRecord, Option)>, ) -> Result<()> { - self.transact("apply block", move |conn| -> Result<()> { - // TODO: This span is logged in a root span, we should connect it to the parent one. - let _span = info_span!(target: COMPONENT, "write_block_to_db").entered(); + let block_num = block.header().block_num(); - models::queries::apply_block( - conn, - block.header(), - ¬es, - block.body().created_nullifiers(), - block.body().updated_accounts(), - block.body().transactions(), - )?; - - // XXX FIXME TODO free floating mutex MUST NOT exist - // it doesn't bind it properly to the data locked! - if allow_acquire.send(()).is_err() { - tracing::warn!(target: COMPONENT, "failed to send notification for successful block application, potential deadlock"); + let result = self + .transact("apply block", move |conn| -> Result<()> { + // TODO: This span is logged in a root span, we should connect it to the parent one. + let _span = info_span!(target: COMPONENT, "write_block_to_db").entered(); + + models::queries::apply_block( + conn, + block.header(), + ¬es, + block.body().created_nullifiers(), + block.body().updated_accounts(), + block.body().transactions(), + )?; + + // XXX FIXME TODO free floating mutex MUST NOT exist + // it doesn't bind it properly to the data locked! + if allow_acquire.send(()).is_err() { + tracing::warn!(target: COMPONENT, "failed to send notification for successful block application, potential deadlock"); + } + + acquire_done.blocking_recv()?; + + Ok(()) + }) + .await; + + // Notify the cleanup task of the latest applied block + // Ignore errors since cleanup is non-critical and shouldn't block block application + let _res = self.notify_cleanup_task.try_send(block_num); + + result + } + + /// Background task that handles periodic cleanup of old account data. + /// + /// This task runs indefinitely, receiving block numbers from the `apply_block` method + /// and triggering cleanup whenever new blocks are available. The cleanup process: + /// + /// 1. Batches incoming notifications using `recv_many` to avoid excessive cleanup operations + /// 2. Only processes the most recent block number from the batch (coalescing multiple updates) + /// 3. Runs cleanup with a 30-second timeout to prevent blocking + /// 4. Logs success or failure but continues running regardless of cleanup outcome + /// + /// # Batching Strategy + /// + /// The batching approach ensures that if multiple blocks are applied quickly (e.g., during + /// initial sync), only the latest block number triggers cleanup. This prevents redundant + /// cleanup operations while ensuring cleanup runs on the most recent state. + /// + /// # Error Handling + /// + /// This task never exits on cleanup errors. Cleanup failures are logged but the task + /// continues to process future blocks. This ensures that temporary issues (like database + /// locks or high load) don't permanently disable the cleanup mechanism. + /// + /// The task only exits if the channel is closed (i.e., all `Db` instances are dropped), + /// which typically happens during application shutdown. + async fn periodic_cleanup_task(db: Self, mut notify: tokio::sync::mpsc::Receiver) { + let mut buf = Vec::with_capacity(128); + + loop { + // Receive many notifications at once to batch them + // If the channel is closed (returns 0), exit the task + let received = notify.recv_many(&mut buf, 128).await; + if received == 0 { + tracing::info!(target: COMPONENT, "Cleanup task shutting down: channel closed"); + break; } - acquire_done.blocking_recv()?; + // Only process the most recent block number from the batch + // This coalesces multiple cleanup requests during fast block processing + if let Some(block_num) = buf.pop() { + match db.run_periodic_cleanup(block_num).await { + Ok((vault_deleted, storage_deleted)) => { + tracing::info!( + target: COMPONENT, + block_num = block_num.as_u32(), + vault_assets_deleted = vault_deleted, + storage_map_values_deleted = storage_deleted, + "Periodic cleanup completed successfully" + ); + }, + Err(e) => { + tracing::warn!( + target: COMPONENT, + block_num = block_num.as_u32(), + error = %e, + "Periodic cleanup failed, will retry on next block" + ); + }, + } + } - Ok(()) - }) - .await + // Clear the buffer for the next batch + buf.clear(); + } + } + + /// Runs periodic cleanup of old account data with a timeout. + /// + /// This function cleans up old vault asset and storage map value entries for all accounts, + /// keeping only the latest entry and up to MAX_HISTORICAL_ENTRIES_PER_ACCOUNT historical + /// entries per key. + /// + /// The cleanup operation has a 30-second timeout to prevent it from blocking for too long. + /// If the timeout is reached, the cleanup is aborted and returns an error. + /// + /// # Parameters + /// * `block_num` - The block number at which cleanup was triggered (used for logging) + /// + /// # Returns + /// A tuple of (vault_assets_deleted, storage_map_values_deleted) on success, or an error + /// if the operation fails or times out. + #[instrument(level = "debug", target = COMPONENT, skip(self), fields(block_num = %block_num.as_u32()))] + async fn run_periodic_cleanup(&self, block_num: BlockNumber) -> Result<(usize, usize)> { + use std::time::Duration; + + let cleanup_timeout = Duration::from_secs(30); + let start = std::time::Instant::now(); + + let cleanup_task = self.transact("periodic cleanup", models::queries::cleanup_all_accounts); + + // Run cleanup with timeout + let result = tokio::time::timeout(cleanup_timeout, cleanup_task).await; + + let duration = start.elapsed(); + + match result { + Ok(Ok((vault_deleted, storage_deleted))) => { + tracing::info!( + target: COMPONENT, + block_num = block_num.as_u32(), + vault_assets_deleted = vault_deleted, + storage_map_values_deleted = storage_deleted, + duration_ms = duration.as_millis(), + "Cleanup completed within timeout" + ); + Ok((vault_deleted, storage_deleted)) + }, + Ok(Err(e)) => { + tracing::error!( + target: COMPONENT, + block_num = block_num.as_u32(), + duration_ms = duration.as_millis(), + error = %e, + "Cleanup failed" + ); + Err(e) + }, + Err(_timeout_err) => { + tracing::warn!( + target: COMPONENT, + block_num = block_num.as_u32(), + timeout_ms = cleanup_timeout.as_millis(), + "Cleanup timed out - operation was aborted" + ); + Err(DatabaseError::QueryTimeout("periodic cleanup".to_string())) + }, + } } /// Selects storage map values for syncing storage maps for a specific account ID. diff --git a/crates/store/src/db/models/queries/accounts.rs b/crates/store/src/db/models/queries/accounts.rs index f517360cd..7a0fc1ff9 100644 --- a/crates/store/src/db/models/queries/accounts.rs +++ b/crates/store/src/db/models/queries/accounts.rs @@ -1170,3 +1170,192 @@ pub(crate) struct AccountStorageMapRowInsert { pub(crate) value: Vec, pub(crate) is_latest: bool, } + +// CLEANUP FUNCTIONS +// ================================================================================================ + +/// Maximum number of historical entries to retain per key for vault assets and storage map +/// values. This ensures we don't accumulate unbounded historical data in the database. +/// The latest entry (marked with is_latest=true) is always retained regardless of this limit. +pub const MAX_HISTORICAL_ENTRIES_PER_KEY: usize = 50; + +/// Clean up old vault asset entries for a specific account, keeping only the latest entry +/// (is_latest=true) and up to MAX_HISTORICAL_ENTRIES_PER_KEY older entries per +/// (account_id, vault_key) combination. +/// +/// # Parameters +/// * `conn`: Database connection +/// * `account_id`: The account to clean up +/// +/// # Returns +/// The number of rows deleted +/// +/// # Notes +/// This function ensures we don't accumulate unbounded historical data while preserving: +/// - The latest state for each vault key (is_latest=true) +/// - Up to MAX_HISTORICAL_ENTRIES_PER_KEY historical entries per vault key +#[cfg(test)] +pub(crate) fn cleanup_old_account_vault_assets( + conn: &mut SqliteConnection, + account_id: AccountId, +) -> Result { + let account_id_bytes = account_id.to_bytes(); + let limit = i64::try_from(MAX_HISTORICAL_ENTRIES_PER_KEY) + .expect("MAX_HISTORICAL_ENTRIES_PER_KEY should fit in i64"); + + // This query: + // 1. For each vault_key, ranks only non-latest entries by block_num descending (newest first) + // 2. Keeps entries where rank <= MAX_HISTORICAL_ENTRIES_PER_KEY + // 3. Deletes everything else using the primary key (account_id, block_num, vault_key) + // Note: The latest entry (is_latest=true) is never ranked and never deleted + + diesel::sql_query( + r#" + DELETE FROM account_vault_assets + WHERE (account_id, block_num, vault_key) IN ( + SELECT account_id, block_num, vault_key FROM ( + SELECT + account_id, + block_num, + vault_key, + ROW_NUMBER() OVER ( + PARTITION BY vault_key + ORDER BY block_num DESC + ) as row_num + FROM account_vault_assets + WHERE account_id = ?1 AND is_latest = 0 + ) + WHERE row_num > ?2 + ) + "#, + ) + .bind::(&account_id_bytes) + .bind::(limit) + .execute(conn) + .map_err(DatabaseError::Diesel) +} + +/// Clean up old storage map value entries for a specific account, keeping only the latest entry +/// (is_latest=true) and up to MAX_HISTORICAL_ENTRIES_PER_KEY older entries per +/// (account_id, slot_name, key) combination. +/// +/// # Parameters +/// * `conn`: Database connection +/// * `account_id`: The account to clean up +/// +/// # Returns +/// The number of rows deleted +/// +/// # Notes +/// This function ensures we don't accumulate unbounded historical data while preserving: +/// - The latest state for each (slot_name, key) combination (is_latest=true) +/// - Up to MAX_HISTORICAL_ENTRIES_PER_KEY historical entries per (slot_name, key) +#[cfg(test)] +pub(crate) fn cleanup_old_account_storage_map_values( + conn: &mut SqliteConnection, + account_id: AccountId, +) -> Result { + let account_id_bytes = account_id.to_bytes(); + let limit = i64::try_from(MAX_HISTORICAL_ENTRIES_PER_KEY) + .expect("MAX_HISTORICAL_ENTRIES_PER_KEY should fit in i64"); + + // This query: + // 1. For each (slot_name, key) combination, ranks only non-latest entries by block_num + // descending (newest first) + // 2. Keeps entries where rank <= MAX_HISTORICAL_ENTRIES_PER_KEY + // 3. Deletes everything else using the primary key (account_id, block_num, slot_name, key) + // Note: The latest entry (is_latest=true) is never ranked and never deleted + + diesel::sql_query( + r#" + DELETE FROM account_storage_map_values + WHERE (account_id, block_num, slot_name, key) IN ( + SELECT account_id, block_num, slot_name, key FROM ( + SELECT + account_id, + block_num, + slot_name, + key, + ROW_NUMBER() OVER ( + PARTITION BY slot_name, key + ORDER BY block_num DESC + ) as row_num + FROM account_storage_map_values + WHERE account_id = ?1 AND is_latest = 0 + ) + WHERE row_num > ?2 + ) + "#, + ) + .bind::(&account_id_bytes) + .bind::(limit) + .execute(conn) + .map_err(DatabaseError::Diesel) +} + +/// Clean up old entries for all accounts in the database using efficient batched SQL. +/// +/// Uses window functions to delete old historical entries across ALL accounts in a single +/// SQL statement per table, avoiding the need to iterate through accounts. +/// +/// # Returns +/// A tuple of (vault_assets_deleted, storage_map_values_deleted) +pub fn cleanup_all_accounts(conn: &mut SqliteConnection) -> Result<(usize, usize), DatabaseError> { + let limit = i64::try_from(MAX_HISTORICAL_ENTRIES_PER_KEY) + .expect("MAX_HISTORICAL_ENTRIES_PER_KEY should fit in i64"); + + // Delete old vault assets across ALL accounts in a single query. + // Window function partitions by (account_id, vault_key) and ranks by block_num DESC. + let vault_deleted = diesel::sql_query( + r#" + DELETE FROM account_vault_assets + WHERE (account_id, block_num, vault_key) IN ( + SELECT account_id, block_num, vault_key FROM ( + SELECT + account_id, + block_num, + vault_key, + ROW_NUMBER() OVER ( + PARTITION BY account_id, vault_key + ORDER BY block_num DESC + ) as row_num + FROM account_vault_assets + WHERE is_latest = 0 + ) + WHERE row_num > ?1 + ) + "#, + ) + .bind::(limit) + .execute(conn) + .map_err(DatabaseError::Diesel)?; + + // Delete old storage map values across ALL accounts in a single query. + // Window function partitions by (account_id, slot_name, key) and ranks by block_num DESC. + let storage_deleted = diesel::sql_query( + r#" + DELETE FROM account_storage_map_values + WHERE (account_id, block_num, slot_name, key) IN ( + SELECT account_id, block_num, slot_name, key FROM ( + SELECT + account_id, + block_num, + slot_name, + key, + ROW_NUMBER() OVER ( + PARTITION BY account_id, slot_name, key + ORDER BY block_num DESC + ) as row_num + FROM account_storage_map_values + WHERE is_latest = 0 + ) + WHERE row_num > ?1 + ) + "#, + ) + .bind::(limit) + .execute(conn) + .map_err(DatabaseError::Diesel)?; + + Ok((vault_deleted, storage_deleted)) +} diff --git a/crates/store/src/db/tests.rs b/crates/store/src/db/tests.rs index 3988e160d..9b5f22d7a 100644 --- a/crates/store/src/db/tests.rs +++ b/crates/store/src/db/tests.rs @@ -72,10 +72,10 @@ use pretty_assertions::assert_eq; use rand::Rng; use super::{AccountInfo, NoteRecord, NullifierInfo}; -use crate::db::TransactionSummary; use crate::db::migrations::apply_migrations; use crate::db::models::queries::{StorageMapValue, insert_account_storage_map_value}; use crate::db::models::{Page, queries, utils}; +use crate::db::{TransactionSummary, schema}; use crate::errors::DatabaseError; fn create_db() -> SqliteConnection { @@ -2348,3 +2348,338 @@ fn db_roundtrip_account_storage_with_maps() { "Full account commitment must match after DB roundtrip" ); } + +// CLEANUP TESTS +// ================================================================================================ + +#[test] +#[miden_node_test_macro::enable_logging] +fn test_cleanup_old_account_vault_assets() { + use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl}; + + use crate::db::models::queries::{ + MAX_HISTORICAL_ENTRIES_PER_KEY, + cleanup_old_account_vault_assets, + }; + + let mut conn = create_db(); + let account_id = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap(); + + // Create blocks + for i in 1..=60 { + create_block(&mut conn, BlockNumber::from(i)); + } + + // Create account + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + BlockNumber::from(1), + ) + .unwrap(); + + let vault_key = AssetVaultKey::new_unchecked(num_to_word(100)); + let asset = Asset::Fungible(FungibleAsset::new(account_id, 1000).unwrap()); + + // Insert 60 vault asset entries for the same vault_key + for block_num in 1..=60 { + queries::insert_account_vault_asset( + &mut conn, + account_id, + BlockNumber::from(block_num), + vault_key, + Some(asset), + ) + .unwrap(); + } + + // Verify we have 60 entries using Diesel API + use schema::account_vault_assets::dsl; + let count = dsl::account_vault_assets + .filter(dsl::account_id.eq(account_id.to_bytes())) + .count() + .get_result::(&mut conn) + .unwrap(); + assert_eq!(count, 60, "Should have 60 entries before cleanup"); + + // Run cleanup + let deleted = cleanup_old_account_vault_assets(&mut conn, account_id).unwrap(); + + // We should have deleted entries beyond MAX_HISTORICAL_ENTRIES_PER_KEY + // The latest entry (is_latest=true) is always kept + // Plus up to MAX_HISTORICAL_ENTRIES_PER_KEY non-latest entries + let expected_remaining = MAX_HISTORICAL_ENTRIES_PER_KEY + 1; // +1 for the latest + let expected_deleted = 60 - expected_remaining; + + assert_eq!( + deleted, expected_deleted, + "Should have deleted {} old entries", + expected_deleted + ); + + // Verify remaining count using Diesel API + let remaining = dsl::account_vault_assets + .filter(dsl::account_id.eq(account_id.to_bytes())) + .count() + .get_result::(&mut conn) + .unwrap(); + assert_eq!( + remaining as usize, expected_remaining, + "Should have {} entries remaining", + expected_remaining + ); + + // Verify the latest entry is still marked as latest using Diesel API + let latest_count = dsl::account_vault_assets + .filter(dsl::account_id.eq(account_id.to_bytes())) + .filter(dsl::is_latest.eq(true)) + .count() + .get_result::(&mut conn) + .unwrap(); + assert_eq!(latest_count, 1, "Should have exactly one latest entry"); + + // Verify the latest entry is from block 60 using Diesel API + let latest_block = dsl::account_vault_assets + .filter(dsl::account_id.eq(account_id.to_bytes())) + .filter(dsl::is_latest.eq(true)) + .select(dsl::block_num) + .first::(&mut conn) + .unwrap(); + assert_eq!(latest_block, 60, "Latest entry should be from block 60"); +} + +#[test] +#[miden_node_test_macro::enable_logging] +fn test_cleanup_old_account_storage_map_values() { + use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl}; + + use crate::db::models::queries::{ + MAX_HISTORICAL_ENTRIES_PER_KEY, + cleanup_old_account_storage_map_values, + }; + + let mut conn = create_db(); + let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE).unwrap(); + + // Create blocks + for i in 1..=70 { + create_block(&mut conn, BlockNumber::from(i)); + } + + let slot_name = StorageSlotName::mock(5); + let key = num_to_word(123); + let value_base = num_to_word(456); + + // Insert 70 storage map value entries for the same (slot_name, key) combination + for block_num in 1..=70 { + queries::insert_account_storage_map_value( + &mut conn, + account_id, + BlockNumber::from(block_num), + slot_name.clone(), + key, + value_base, + ) + .unwrap(); + } + + // Verify we have 70 entries using Diesel API + use schema::account_storage_map_values::dsl; + let count = dsl::account_storage_map_values + .filter(dsl::account_id.eq(account_id.to_bytes())) + .count() + .get_result::(&mut conn) + .unwrap(); + assert_eq!(count, 70, "Should have 70 entries before cleanup"); + + // Run cleanup + let deleted = cleanup_old_account_storage_map_values(&mut conn, account_id).unwrap(); + + // We should have deleted entries beyond MAX_HISTORICAL_ENTRIES_PER_KEY + let expected_remaining = MAX_HISTORICAL_ENTRIES_PER_KEY + 1; // +1 for the latest + let expected_deleted = 70 - expected_remaining; + + assert_eq!( + deleted, expected_deleted, + "Should have deleted {} old entries", + expected_deleted + ); + + // Verify remaining count using Diesel API + let remaining = dsl::account_storage_map_values + .filter(dsl::account_id.eq(account_id.to_bytes())) + .count() + .get_result::(&mut conn) + .unwrap(); + assert_eq!( + remaining as usize, expected_remaining, + "Should have {} entries remaining", + expected_remaining + ); + + // Verify the latest entry is still marked as latest using Diesel API + let latest_count = dsl::account_storage_map_values + .filter(dsl::account_id.eq(account_id.to_bytes())) + .filter(dsl::is_latest.eq(true)) + .count() + .get_result::(&mut conn) + .unwrap(); + assert_eq!(latest_count, 1, "Should have exactly one latest entry"); +} + +#[test] +#[miden_node_test_macro::enable_logging] +fn test_cleanup_preserves_latest_state() { + use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl}; + + use crate::db::models::queries::cleanup_old_account_vault_assets; + + let mut conn = create_db(); + let account_id = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap(); + + // Create blocks + for i in 1..=10 { + create_block(&mut conn, BlockNumber::from(i)); + } + + // Create account + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + BlockNumber::from(1), + ) + .unwrap(); + + // Test with multiple vault keys to ensure per-key cleanup + let vault_key_1 = AssetVaultKey::new_unchecked(num_to_word(100)); + let vault_key_2 = AssetVaultKey::new_unchecked(num_to_word(200)); + let asset = Asset::Fungible(FungibleAsset::new(account_id, 1000).unwrap()); + + // Insert entries for both keys + for block_num in 1..=10 { + queries::insert_account_vault_asset( + &mut conn, + account_id, + BlockNumber::from(block_num), + vault_key_1, + Some(asset), + ) + .unwrap(); + + queries::insert_account_vault_asset( + &mut conn, + account_id, + BlockNumber::from(block_num), + vault_key_2, + Some(asset), + ) + .unwrap(); + } + + // Run cleanup (should not delete anything since we're under the limit) + let deleted = cleanup_old_account_vault_assets(&mut conn, account_id).unwrap(); + assert_eq!(deleted, 0, "Should not delete anything when under limit"); + + // Verify both latest entries exist using Diesel API + use schema::account_vault_assets::dsl; + let latest_entries: Vec<(Vec, i64)> = dsl::account_vault_assets + .filter(dsl::account_id.eq(account_id.to_bytes())) + .filter(dsl::is_latest.eq(true)) + .select((dsl::vault_key, dsl::block_num)) + .order(dsl::vault_key.asc()) + .load(&mut conn) + .unwrap(); + + assert_eq!(latest_entries.len(), 2, "Should have two latest entries"); + assert_eq!(latest_entries[0].1, 10, "Latest for key 1 should be block 10"); + assert_eq!(latest_entries[1].1, 10, "Latest for key 2 should be block 10"); +} + +#[test] +#[miden_node_test_macro::enable_logging] +fn test_cleanup_all_accounts() { + use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl}; + + use crate::db::models::queries::cleanup_all_accounts; + + let mut conn = create_db(); + + // Create two accounts + let account1 = AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE).unwrap(); + let account2 = AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE_2).unwrap(); + let faucet_id = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap(); + + // Create blocks + for i in 1..=60 { + create_block(&mut conn, BlockNumber::from(i)); + } + + // Create accounts + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account1, 0)], + BlockNumber::from(1), + ) + .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account2, 0)], + BlockNumber::from(1), + ) + .unwrap(); + + // Insert many entries for both accounts + let vault_key = AssetVaultKey::new_unchecked(num_to_word(100)); + let asset = Asset::Fungible(FungibleAsset::new(faucet_id, 1000).unwrap()); + + for block_num in 1..=60 { + queries::insert_account_vault_asset( + &mut conn, + account1, + BlockNumber::from(block_num), + vault_key, + Some(asset), + ) + .unwrap(); + + let asset2 = Asset::Fungible(FungibleAsset::new(faucet_id, 2000).unwrap()); + queries::insert_account_vault_asset( + &mut conn, + account2, + BlockNumber::from(block_num), + vault_key, + Some(asset2), + ) + .unwrap(); + } + + // Run cleanup for all accounts + let (vault_deleted, _) = cleanup_all_accounts(&mut conn).unwrap(); + + // We should have deleted entries from both accounts + assert!(vault_deleted > 0, "Should have deleted some entries"); + + // Each account should have MAX_HISTORICAL_ENTRIES_PER_KEY + 1 entries remaining using + // Diesel API + use schema::account_vault_assets::dsl; + let count1 = dsl::account_vault_assets + .filter(dsl::account_id.eq(account1.to_bytes())) + .count() + .get_result::(&mut conn) + .unwrap(); + + let count2 = dsl::account_vault_assets + .filter(dsl::account_id.eq(account2.to_bytes())) + .count() + .get_result::(&mut conn) + .unwrap(); + + assert_eq!( + count1 as usize, + crate::db::models::queries::MAX_HISTORICAL_ENTRIES_PER_KEY + 1 + ); + assert_eq!( + count2 as usize, + crate::db::models::queries::MAX_HISTORICAL_ENTRIES_PER_KEY + 1 + ); +} diff --git a/crates/store/src/errors.rs b/crates/store/src/errors.rs index 7ebed5a74..a0426b031 100644 --- a/crates/store/src/errors.rs +++ b/crates/store/src/errors.rs @@ -134,6 +134,8 @@ pub enum DatabaseError { SqlValueConversion(#[from] DatabaseTypeConversionError), #[error("Not implemented: {0}")] NotImplemented(String), + #[error("query timeout: {0}")] + QueryTimeout(String), #[error("storage root not found for account {account_id}, slot {slot_name}, block {block_num}")] StorageRootNotFound { account_id: AccountId, diff --git a/crates/store/src/inner_forest/mod.rs b/crates/store/src/inner_forest/mod.rs index a5d47ac51..8c144b270 100644 --- a/crates/store/src/inner_forest/mod.rs +++ b/crates/store/src/inner_forest/mod.rs @@ -13,6 +13,13 @@ use thiserror::Error; #[cfg(test)] mod tests; +// CONSTANTS +// ================================================================================================ + +/// Maximum number of historical block entries to retain per account key in the in-memory forest. +/// This matches the database retention policy in [`crate::db::models::queries::accounts`]. +pub const MAX_HISTORICAL_BLOCKS_PER_KEY: usize = 50; + // ERRORS // ================================================================================================ @@ -442,4 +449,138 @@ impl InnerForest { ); } } + + // PRUNING + // -------------------------------------------------------------------------------------------- + + /// Prunes old entries from the in-memory forest data structures. + /// + /// Removes historical entries exceeding the retention limit while preserving the most recent + /// entries for each key: + /// - For vault_roots: keeps at most MAX_HISTORICAL_BLOCKS_PER_KEY entries per account_id + /// - For storage_map_roots: keeps at most MAX_HISTORICAL_BLOCKS_PER_KEY entries per + /// (account_id, slot_name) + /// - For storage_entries: keeps at most MAX_HISTORICAL_BLOCKS_PER_KEY entries per (account_id, + /// slot_name) + /// + /// The SmtForest itself is not pruned directly as it uses structural sharing and old roots + /// are naturally garbage-collected when they become unreachable. + /// + /// # Returns + /// + /// A tuple containing the number of entries removed from: + /// (vault_roots_removed, storage_map_roots_removed, storage_entries_removed) + pub(crate) fn prune(&mut self) -> (usize, usize, usize) { + let vault_removed = Self::prune_map_by_account(&mut self.vault_roots); + let storage_roots_removed = Self::prune_map_by_account_slot(&mut self.storage_map_roots); + let storage_entries_removed = Self::prune_map_by_account_slot(&mut self.storage_entries); + + if vault_removed > 0 || storage_roots_removed > 0 || storage_entries_removed > 0 { + tracing::info!( + target: crate::COMPONENT, + vault_roots_removed = vault_removed, + storage_map_roots_removed = storage_roots_removed, + storage_entries_removed = storage_entries_removed, + "Forest pruning completed" + ); + } + + // Prune the underlying SmtForest by telling it which roots are still referenced. + // Collect all roots that we're keeping from both vault_roots and storage_map_roots. + let roots_to_keep: Vec = self + .vault_roots + .values() + .chain(self.storage_map_roots.values()) + .copied() + .collect(); + self.forest.pop_smts(roots_to_keep); + + (vault_removed, storage_roots_removed, storage_entries_removed) + } + + /// Prunes a map keyed by (AccountId, BlockNumber), keeping at most + /// MAX_HISTORICAL_BLOCKS_PER_KEY entries per account. + fn prune_map_by_account(map: &mut BTreeMap<(AccountId, BlockNumber), V>) -> usize { + // Group block numbers by account_id + let mut entries_by_account: BTreeMap> = BTreeMap::new(); + for (account_id, block_num) in map.keys() { + entries_by_account.entry(*account_id).or_default().push(*block_num); + } + + let mut removed = 0; + for (account_id, mut block_nums) in entries_by_account { + let len = block_nums.len(); + if len <= MAX_HISTORICAL_BLOCKS_PER_KEY { + continue; + } + + // Find the cutoff block number using O(n) selection algorithm. + // We want to keep the MAX_HISTORICAL_BLOCKS_PER_KEY largest block numbers. + // select_nth_unstable_by partitions so that element at index is in sorted position. + // After this, elements at indices >= cutoff_idx are the largest. + let cutoff_idx = len - MAX_HISTORICAL_BLOCKS_PER_KEY; + block_nums.select_nth_unstable(cutoff_idx); + let cutoff_block = block_nums[cutoff_idx]; + + // Remove entries with block_num < cutoff_block + // Since BTreeMap is sorted, we can use range to find entries to remove. + let keys_to_remove: Vec<_> = map + .range((account_id, BlockNumber::GENESIS)..(account_id, cutoff_block)) + .map(|(k, _)| *k) + .collect(); + + for key in keys_to_remove { + map.remove(&key); + removed += 1; + } + } + + removed + } + + /// Prunes a map keyed by (AccountId, StorageSlotName, BlockNumber), keeping at most + /// MAX_HISTORICAL_BLOCKS_PER_KEY entries per (account_id, slot_name). + fn prune_map_by_account_slot( + map: &mut BTreeMap<(AccountId, StorageSlotName, BlockNumber), V>, + ) -> usize { + // Group block numbers by (account_id, slot_name) + // Clone slot_name to avoid holding references to map during mutation + let mut entries_by_key: BTreeMap<(AccountId, StorageSlotName), Vec> = + BTreeMap::new(); + for (account_id, slot_name, block_num) in map.keys() { + entries_by_key + .entry((*account_id, slot_name.clone())) + .or_default() + .push(*block_num); + } + + let mut removed = 0; + for ((account_id, slot_name), mut block_nums) in entries_by_key { + let len = block_nums.len(); + if len <= MAX_HISTORICAL_BLOCKS_PER_KEY { + continue; + } + + // Find the cutoff block number using O(n) selection algorithm + let cutoff_idx = len - MAX_HISTORICAL_BLOCKS_PER_KEY; + block_nums.select_nth_unstable(cutoff_idx); + let cutoff_block = block_nums[cutoff_idx]; + + // Remove entries with block_num < cutoff_block + let keys_to_remove: Vec<_> = map + .range( + (account_id, slot_name.clone(), BlockNumber::GENESIS) + ..(account_id, slot_name.clone(), cutoff_block), + ) + .map(|(k, _)| k.clone()) + .collect(); + + for key in keys_to_remove { + map.remove(&key); + removed += 1; + } + } + + removed + } } diff --git a/crates/store/src/inner_forest/tests.rs b/crates/store/src/inner_forest/tests.rs index da311db46..f55bb3cd9 100644 --- a/crates/store/src/inner_forest/tests.rs +++ b/crates/store/src/inner_forest/tests.rs @@ -461,3 +461,155 @@ fn test_open_storage_map_returns_limit_exceeded_for_too_many_keys() { let details = result.expect("Should return Some").expect("Should not error"); assert_matches!(details.entries, StorageMapEntries::LimitExceeded); } + +// PRUNING TESTS +// ================================================================================================ + +#[test] +fn test_prune_vault_roots_removes_old_entries() { + let mut forest = InnerForest::new(); + let account_id = dummy_account(); + let faucet_id = dummy_faucet(); + + // Add entries for MAX_HISTORICAL_BLOCKS_PER_KEY + 10 blocks + let num_blocks = MAX_HISTORICAL_BLOCKS_PER_KEY + 10; + for i in 1..=num_blocks { + let block_num = BlockNumber::from(i as u32); + let amount = (i * 100) as u64; + let mut vault_delta = AccountVaultDelta::default(); + vault_delta.add_asset(dummy_fungible_asset(faucet_id, amount)).unwrap(); + let delta = dummy_partial_delta(account_id, vault_delta, AccountStorageDelta::default()); + forest.update_account(block_num, &delta).unwrap(); + } + + // Verify we have all entries before pruning + assert_eq!(forest.vault_roots.len(), num_blocks); + + // Prune + let (vault_removed, ..) = forest.prune(); + + // Should have removed 10 entries + assert_eq!(vault_removed, 10); + + // Should have exactly MAX_HISTORICAL_BLOCKS_PER_KEY entries remaining + assert_eq!(forest.vault_roots.len(), MAX_HISTORICAL_BLOCKS_PER_KEY); + + // The remaining entries should be the most recent ones + let remaining_blocks: Vec<_> = forest.vault_roots.keys().map(|(_, b)| b.as_u32()).collect(); + let oldest_remaining = *remaining_blocks.iter().min().unwrap(); + // Oldest remaining should be block 11 (since we kept the most recent 50 out of 60) + assert_eq!(oldest_remaining, 11); +} + +#[test] +fn test_prune_storage_map_roots_removes_old_entries() { + use std::collections::BTreeMap; + + use miden_protocol::account::delta::{StorageMapDelta, StorageSlotDelta}; + + let mut forest = InnerForest::new(); + let account_id = dummy_account(); + let slot_name = StorageSlotName::mock(3); + + // Add entries for MAX_HISTORICAL_BLOCKS_PER_KEY + 5 blocks + let num_blocks = MAX_HISTORICAL_BLOCKS_PER_KEY + 5; + for i in 1..=num_blocks { + let block_num = BlockNumber::from(i as u32); + let key = Word::from([i as u32, 0, 0, 0]); + let value = Word::from([0, 0, 0, i as u32]); + + let mut map_delta = StorageMapDelta::default(); + map_delta.insert(key, value); + let raw = BTreeMap::from_iter([(slot_name.clone(), StorageSlotDelta::Map(map_delta))]); + let storage_delta = AccountStorageDelta::from_raw(raw); + let delta = dummy_partial_delta(account_id, AccountVaultDelta::default(), storage_delta); + forest.update_account(block_num, &delta).unwrap(); + } + + // Verify we have all entries before pruning + assert_eq!(forest.storage_map_roots.len(), num_blocks); + + // Prune + let (_, storage_roots_removed, storage_entries_removed) = forest.prune(); + + // Should have removed 5 entries from each structure + assert_eq!(storage_roots_removed, 5); + assert_eq!(storage_entries_removed, 5); + + // Should have exactly MAX_HISTORICAL_BLOCKS_PER_KEY entries remaining + assert_eq!(forest.storage_map_roots.len(), MAX_HISTORICAL_BLOCKS_PER_KEY); + assert_eq!(forest.storage_entries.len(), MAX_HISTORICAL_BLOCKS_PER_KEY); +} + +#[test] +fn test_prune_does_nothing_when_under_limit() { + let mut forest = InnerForest::new(); + let account_id = dummy_account(); + let faucet_id = dummy_faucet(); + + // Add fewer entries than the limit + let num_blocks = 10; + for i in 1..=num_blocks { + let block_num = BlockNumber::from(i as u32); + let mut vault_delta = AccountVaultDelta::default(); + vault_delta + .add_asset(dummy_fungible_asset(faucet_id, (i * 100) as u64)) + .unwrap(); + let delta = dummy_partial_delta(account_id, vault_delta, AccountStorageDelta::default()); + forest.update_account(block_num, &delta).unwrap(); + } + + // Prune + let (vault_removed, storage_roots_removed, storage_entries_removed) = forest.prune(); + + // Nothing should be removed + assert_eq!(vault_removed, 0); + assert_eq!(storage_roots_removed, 0); + assert_eq!(storage_entries_removed, 0); + + // All entries should remain + assert_eq!(forest.vault_roots.len(), num_blocks); +} + +#[test] +fn test_prune_handles_multiple_accounts() { + let mut forest = InnerForest::new(); + let account1 = dummy_account(); + let account2 = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap(); + let faucet_id = dummy_faucet(); + + // Add entries for multiple accounts, each exceeding the limit + let num_blocks = MAX_HISTORICAL_BLOCKS_PER_KEY + 5; + + for i in 1..=num_blocks { + let block_num = BlockNumber::from(i as u32); + let amount = (i * 100) as u64; + + // Account 1 + let mut vault_delta1 = AccountVaultDelta::default(); + vault_delta1.add_asset(dummy_fungible_asset(faucet_id, amount)).unwrap(); + let delta1 = dummy_partial_delta(account1, vault_delta1, AccountStorageDelta::default()); + forest.update_account(block_num, &delta1).unwrap(); + + // Account 2 (using a different faucet for variety) + let mut vault_delta2 = AccountVaultDelta::default(); + vault_delta2.add_asset(dummy_fungible_asset(account2, amount * 2)).unwrap(); + let delta2 = dummy_partial_delta(account2, vault_delta2, AccountStorageDelta::default()); + forest.update_account(block_num, &delta2).unwrap(); + } + + // Verify we have entries for both accounts before pruning + assert_eq!(forest.vault_roots.len(), num_blocks * 2); + + // Prune + let (vault_removed, ..) = forest.prune(); + + // Should have removed 5 entries from each account = 10 total + assert_eq!(vault_removed, 10); + + // Each account should have exactly MAX_HISTORICAL_BLOCKS_PER_KEY entries + let account1_entries = forest.vault_roots.keys().filter(|(id, _)| *id == account1).count(); + let account2_entries = forest.vault_roots.keys().filter(|(id, _)| *id == account2).count(); + assert_eq!(account1_entries, MAX_HISTORICAL_BLOCKS_PER_KEY); + assert_eq!(account2_entries, MAX_HISTORICAL_BLOCKS_PER_KEY); +} diff --git a/crates/store/src/state.rs b/crates/store/src/state.rs index 2ff1f887b..e31ee3c3b 100644 --- a/crates/store/src/state.rs +++ b/crates/store/src/state.rs @@ -569,6 +569,21 @@ impl State { self.forest.write().await.apply_block_updates(block_num, account_deltas)?; + // Prune old entries from the in-memory forest + // This runs after each block to prevent unbounded memory growth + let (vault_pruned, storage_roots_pruned, storage_entries_pruned) = + self.forest.write().await.prune(); + if vault_pruned > 0 || storage_roots_pruned > 0 || storage_entries_pruned > 0 { + tracing::debug!( + target: COMPONENT, + block_num = block_num.as_u32(), + vault_pruned, + storage_roots_pruned, + storage_entries_pruned, + "Forest pruning applied" + ); + } + info!(%block_commitment, block_num = block_num.as_u32(), COMPONENT, "apply_block successful"); Ok(()) From 1d45f47b2be490e7e4b2eee9cfe4358609446359 Mon Sep 17 00:00:00 2001 From: Bernhard Schuster Date: Mon, 19 Jan 2026 14:54:05 +0100 Subject: [PATCH 2/5] use a single write lock --- crates/store/src/state.rs | 34 ++++++++++++++++++---------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/crates/store/src/state.rs b/crates/store/src/state.rs index e31ee3c3b..1172d152f 100644 --- a/crates/store/src/state.rs +++ b/crates/store/src/state.rs @@ -565,23 +565,25 @@ impl State { .apply_mutations(account_tree_update) .expect("Unreachable: old account tree root must be checked before this step"); inner.blockchain.push(block_commitment); - } - self.forest.write().await.apply_block_updates(block_num, account_deltas)?; - - // Prune old entries from the in-memory forest - // This runs after each block to prevent unbounded memory growth - let (vault_pruned, storage_roots_pruned, storage_entries_pruned) = - self.forest.write().await.prune(); - if vault_pruned > 0 || storage_roots_pruned > 0 || storage_entries_pruned > 0 { - tracing::debug!( - target: COMPONENT, - block_num = block_num.as_u32(), - vault_pruned, - storage_roots_pruned, - storage_entries_pruned, - "Forest pruning applied" - ); + // Update forest while still holding the inner lock to ensure consistency. + // Readers acquiring the inner read lock will see both account tree and forest + // updated atomically. + let mut forest = self.forest.write().await; + forest.apply_block_updates(block_num, account_deltas)?; + + // Prune old entries from the in-memory forest while holding both locks + let (vault_pruned, storage_roots_pruned, storage_entries_pruned) = forest.prune(); + if vault_pruned > 0 || storage_roots_pruned > 0 || storage_entries_pruned > 0 { + tracing::debug!( + target: COMPONENT, + block_num = block_num.as_u32(), + vault_pruned, + storage_roots_pruned, + storage_entries_pruned, + "Forest pruning applied" + ); + } } info!(%block_commitment, block_num = block_num.as_u32(), COMPONENT, "apply_block successful"); From 925ca237f0966e07da26b63be1fbb86cd8bcdc5a Mon Sep 17 00:00:00 2001 From: Bernhard Schuster Date: Mon, 19 Jan 2026 20:19:29 +0100 Subject: [PATCH 3/5] simplify --- crates/store/src/db/mod.rs | 4 +- .../store/src/db/models/queries/accounts.rs | 157 ++++-------------- crates/store/src/db/tests.rs | 138 +++++++-------- crates/store/src/inner_forest/mod.rs | 127 +++++--------- crates/store/src/inner_forest/tests.rs | 119 +++++++------ crates/store/src/state.rs | 3 +- 6 files changed, 219 insertions(+), 329 deletions(-) diff --git a/crates/store/src/db/mod.rs b/crates/store/src/db/mod.rs index 5e68c7e3a..ff9300b11 100644 --- a/crates/store/src/db/mod.rs +++ b/crates/store/src/db/mod.rs @@ -703,7 +703,9 @@ impl Db { let cleanup_timeout = Duration::from_secs(30); let start = std::time::Instant::now(); - let cleanup_task = self.transact("periodic cleanup", models::queries::cleanup_all_accounts); + let cleanup_task = self.transact("periodic cleanup", move |conn| { + models::queries::cleanup_all_accounts(conn, block_num) + }); // Run cleanup with timeout let result = tokio::time::timeout(cleanup_timeout, cleanup_task).await; diff --git a/crates/store/src/db/models/queries/accounts.rs b/crates/store/src/db/models/queries/accounts.rs index 7a0fc1ff9..ff0722e42 100644 --- a/crates/store/src/db/models/queries/accounts.rs +++ b/crates/store/src/db/models/queries/accounts.rs @@ -1174,186 +1174,91 @@ pub(crate) struct AccountStorageMapRowInsert { // CLEANUP FUNCTIONS // ================================================================================================ -/// Maximum number of historical entries to retain per key for vault assets and storage map -/// values. This ensures we don't accumulate unbounded historical data in the database. -/// The latest entry (marked with is_latest=true) is always retained regardless of this limit. -pub const MAX_HISTORICAL_ENTRIES_PER_KEY: usize = 50; - -/// Clean up old vault asset entries for a specific account, keeping only the latest entry -/// (is_latest=true) and up to MAX_HISTORICAL_ENTRIES_PER_KEY older entries per -/// (account_id, vault_key) combination. -/// -/// # Parameters -/// * `conn`: Database connection -/// * `account_id`: The account to clean up -/// -/// # Returns -/// The number of rows deleted +/// Number of historical blocks to retain for vault assets and storage map values. +/// Entries older than `chain_tip - HISTORICAL_BLOCK_RETENTION` will be deleted, +/// except for entries marked with `is_latest=true` which are always retained. +pub const HISTORICAL_BLOCK_RETENTION: u32 = 50; + +/// Clean up old vault asset entries for a specific account, deleting entries older than +/// the retention window. /// -/// # Notes -/// This function ensures we don't accumulate unbounded historical data while preserving: -/// - The latest state for each vault key (is_latest=true) -/// - Up to MAX_HISTORICAL_ENTRIES_PER_KEY historical entries per vault key +/// Keeps entries where `block_num >= cutoff_block` OR `is_latest = true`. #[cfg(test)] pub(crate) fn cleanup_old_account_vault_assets( conn: &mut SqliteConnection, account_id: AccountId, + chain_tip: BlockNumber, ) -> Result { let account_id_bytes = account_id.to_bytes(); - let limit = i64::try_from(MAX_HISTORICAL_ENTRIES_PER_KEY) - .expect("MAX_HISTORICAL_ENTRIES_PER_KEY should fit in i64"); - - // This query: - // 1. For each vault_key, ranks only non-latest entries by block_num descending (newest first) - // 2. Keeps entries where rank <= MAX_HISTORICAL_ENTRIES_PER_KEY - // 3. Deletes everything else using the primary key (account_id, block_num, vault_key) - // Note: The latest entry (is_latest=true) is never ranked and never deleted + let cutoff_block = chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION) as i64; diesel::sql_query( r#" DELETE FROM account_vault_assets - WHERE (account_id, block_num, vault_key) IN ( - SELECT account_id, block_num, vault_key FROM ( - SELECT - account_id, - block_num, - vault_key, - ROW_NUMBER() OVER ( - PARTITION BY vault_key - ORDER BY block_num DESC - ) as row_num - FROM account_vault_assets - WHERE account_id = ?1 AND is_latest = 0 - ) - WHERE row_num > ?2 - ) + WHERE account_id = ?1 AND block_num < ?2 AND is_latest = 0 "#, ) .bind::(&account_id_bytes) - .bind::(limit) + .bind::(cutoff_block) .execute(conn) .map_err(DatabaseError::Diesel) } -/// Clean up old storage map value entries for a specific account, keeping only the latest entry -/// (is_latest=true) and up to MAX_HISTORICAL_ENTRIES_PER_KEY older entries per -/// (account_id, slot_name, key) combination. -/// -/// # Parameters -/// * `conn`: Database connection -/// * `account_id`: The account to clean up +/// Clean up old storage map value entries for a specific account, deleting entries older than +/// the retention window. /// -/// # Returns -/// The number of rows deleted -/// -/// # Notes -/// This function ensures we don't accumulate unbounded historical data while preserving: -/// - The latest state for each (slot_name, key) combination (is_latest=true) -/// - Up to MAX_HISTORICAL_ENTRIES_PER_KEY historical entries per (slot_name, key) +/// Keeps entries where `block_num >= cutoff_block` OR `is_latest = true`. #[cfg(test)] pub(crate) fn cleanup_old_account_storage_map_values( conn: &mut SqliteConnection, account_id: AccountId, + chain_tip: BlockNumber, ) -> Result { let account_id_bytes = account_id.to_bytes(); - let limit = i64::try_from(MAX_HISTORICAL_ENTRIES_PER_KEY) - .expect("MAX_HISTORICAL_ENTRIES_PER_KEY should fit in i64"); - - // This query: - // 1. For each (slot_name, key) combination, ranks only non-latest entries by block_num - // descending (newest first) - // 2. Keeps entries where rank <= MAX_HISTORICAL_ENTRIES_PER_KEY - // 3. Deletes everything else using the primary key (account_id, block_num, slot_name, key) - // Note: The latest entry (is_latest=true) is never ranked and never deleted + let cutoff_block = chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION) as i64; diesel::sql_query( r#" DELETE FROM account_storage_map_values - WHERE (account_id, block_num, slot_name, key) IN ( - SELECT account_id, block_num, slot_name, key FROM ( - SELECT - account_id, - block_num, - slot_name, - key, - ROW_NUMBER() OVER ( - PARTITION BY slot_name, key - ORDER BY block_num DESC - ) as row_num - FROM account_storage_map_values - WHERE account_id = ?1 AND is_latest = 0 - ) - WHERE row_num > ?2 - ) + WHERE account_id = ?1 AND block_num < ?2 AND is_latest = 0 "#, ) .bind::(&account_id_bytes) - .bind::(limit) + .bind::(cutoff_block) .execute(conn) .map_err(DatabaseError::Diesel) } -/// Clean up old entries for all accounts in the database using efficient batched SQL. +/// Clean up old entries for all accounts, deleting entries older than the retention window. /// -/// Uses window functions to delete old historical entries across ALL accounts in a single -/// SQL statement per table, avoiding the need to iterate through accounts. +/// Deletes rows where `block_num < chain_tip - HISTORICAL_BLOCK_RETENTION` and `is_latest = false`. +/// This is a simple and efficient approach that doesn't require window functions. /// /// # Returns /// A tuple of (vault_assets_deleted, storage_map_values_deleted) -pub fn cleanup_all_accounts(conn: &mut SqliteConnection) -> Result<(usize, usize), DatabaseError> { - let limit = i64::try_from(MAX_HISTORICAL_ENTRIES_PER_KEY) - .expect("MAX_HISTORICAL_ENTRIES_PER_KEY should fit in i64"); +pub fn cleanup_all_accounts( + conn: &mut SqliteConnection, + chain_tip: BlockNumber, +) -> Result<(usize, usize), DatabaseError> { + let cutoff_block = chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION) as i64; - // Delete old vault assets across ALL accounts in a single query. - // Window function partitions by (account_id, vault_key) and ranks by block_num DESC. let vault_deleted = diesel::sql_query( r#" DELETE FROM account_vault_assets - WHERE (account_id, block_num, vault_key) IN ( - SELECT account_id, block_num, vault_key FROM ( - SELECT - account_id, - block_num, - vault_key, - ROW_NUMBER() OVER ( - PARTITION BY account_id, vault_key - ORDER BY block_num DESC - ) as row_num - FROM account_vault_assets - WHERE is_latest = 0 - ) - WHERE row_num > ?1 - ) + WHERE block_num < ?1 AND is_latest = 0 "#, ) - .bind::(limit) + .bind::(cutoff_block) .execute(conn) .map_err(DatabaseError::Diesel)?; - // Delete old storage map values across ALL accounts in a single query. - // Window function partitions by (account_id, slot_name, key) and ranks by block_num DESC. let storage_deleted = diesel::sql_query( r#" DELETE FROM account_storage_map_values - WHERE (account_id, block_num, slot_name, key) IN ( - SELECT account_id, block_num, slot_name, key FROM ( - SELECT - account_id, - block_num, - slot_name, - key, - ROW_NUMBER() OVER ( - PARTITION BY account_id, slot_name, key - ORDER BY block_num DESC - ) as row_num - FROM account_storage_map_values - WHERE is_latest = 0 - ) - WHERE row_num > ?1 - ) + WHERE block_num < ?1 AND is_latest = 0 "#, ) - .bind::(limit) + .bind::(cutoff_block) .execute(conn) .map_err(DatabaseError::Diesel)?; diff --git a/crates/store/src/db/tests.rs b/crates/store/src/db/tests.rs index 9b5f22d7a..217203af2 100644 --- a/crates/store/src/db/tests.rs +++ b/crates/store/src/db/tests.rs @@ -2352,21 +2352,24 @@ fn db_roundtrip_account_storage_with_maps() { // CLEANUP TESTS // ================================================================================================ +/// Chain length used in cleanup tests - must be > HISTORICAL_BLOCK_RETENTION for meaningful tests. +const TEST_CHAIN_LENGTH: u32 = 100; + #[test] #[miden_node_test_macro::enable_logging] fn test_cleanup_old_account_vault_assets() { use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl}; use crate::db::models::queries::{ - MAX_HISTORICAL_ENTRIES_PER_KEY, + HISTORICAL_BLOCK_RETENTION, cleanup_old_account_vault_assets, }; let mut conn = create_db(); let account_id = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap(); - // Create blocks - for i in 1..=60 { + // Create blocks 1-TEST_CHAIN_LENGTH + for i in 1..=TEST_CHAIN_LENGTH { create_block(&mut conn, BlockNumber::from(i)); } @@ -2381,8 +2384,8 @@ fn test_cleanup_old_account_vault_assets() { let vault_key = AssetVaultKey::new_unchecked(num_to_word(100)); let asset = Asset::Fungible(FungibleAsset::new(account_id, 1000).unwrap()); - // Insert 60 vault asset entries for the same vault_key - for block_num in 1..=60 { + // Insert vault asset entries for blocks 1-TEST_CHAIN_LENGTH + for block_num in 1..=TEST_CHAIN_LENGTH { queries::insert_account_vault_asset( &mut conn, account_id, @@ -2393,43 +2396,45 @@ fn test_cleanup_old_account_vault_assets() { .unwrap(); } - // Verify we have 60 entries using Diesel API + // Verify we have TEST_CHAIN_LENGTH entries use schema::account_vault_assets::dsl; let count = dsl::account_vault_assets .filter(dsl::account_id.eq(account_id.to_bytes())) .count() .get_result::(&mut conn) .unwrap(); - assert_eq!(count, 60, "Should have 60 entries before cleanup"); - - // Run cleanup - let deleted = cleanup_old_account_vault_assets(&mut conn, account_id).unwrap(); + assert_eq!(count, TEST_CHAIN_LENGTH as i64, "Should have TEST_CHAIN_LENGTH entries before cleanup"); - // We should have deleted entries beyond MAX_HISTORICAL_ENTRIES_PER_KEY - // The latest entry (is_latest=true) is always kept - // Plus up to MAX_HISTORICAL_ENTRIES_PER_KEY non-latest entries - let expected_remaining = MAX_HISTORICAL_ENTRIES_PER_KEY + 1; // +1 for the latest - let expected_deleted = 60 - expected_remaining; + // Run cleanup with chain_tip at block TEST_CHAIN_LENGTH + // cutoff = TEST_CHAIN_LENGTH - 50, so blocks < cutoff should be deleted + let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); + let deleted = cleanup_old_account_vault_assets(&mut conn, account_id, chain_tip).unwrap(); + // Blocks 1 to cutoff-1 are older than cutoff, but block TEST_CHAIN_LENGTH is is_latest=true + // So we should delete (cutoff-1) entries (blocks 1 to cutoff-1, excluding is_latest) + let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; + let expected_deleted = cutoff - 1; // blocks 1 to cutoff-1 assert_eq!( - deleted, expected_deleted, - "Should have deleted {} old entries", - expected_deleted + deleted, expected_deleted as usize, + "Should have deleted entries older than block {}", + cutoff ); - // Verify remaining count using Diesel API + // Verify remaining count let remaining = dsl::account_vault_assets .filter(dsl::account_id.eq(account_id.to_bytes())) .count() .get_result::(&mut conn) .unwrap(); + // Remaining: blocks cutoff to TEST_CHAIN_LENGTH + let expected_remaining = TEST_CHAIN_LENGTH - expected_deleted; assert_eq!( - remaining as usize, expected_remaining, + remaining as u32, expected_remaining, "Should have {} entries remaining", expected_remaining ); - // Verify the latest entry is still marked as latest using Diesel API + // Verify the latest entry is still marked as latest let latest_count = dsl::account_vault_assets .filter(dsl::account_id.eq(account_id.to_bytes())) .filter(dsl::is_latest.eq(true)) @@ -2438,14 +2443,14 @@ fn test_cleanup_old_account_vault_assets() { .unwrap(); assert_eq!(latest_count, 1, "Should have exactly one latest entry"); - // Verify the latest entry is from block 60 using Diesel API + // Verify the latest entry is from block TEST_CHAIN_LENGTH let latest_block = dsl::account_vault_assets .filter(dsl::account_id.eq(account_id.to_bytes())) .filter(dsl::is_latest.eq(true)) .select(dsl::block_num) .first::(&mut conn) .unwrap(); - assert_eq!(latest_block, 60, "Latest entry should be from block 60"); + assert_eq!(latest_block, TEST_CHAIN_LENGTH as i64, "Latest entry should be from block TEST_CHAIN_LENGTH"); } #[test] @@ -2454,15 +2459,15 @@ fn test_cleanup_old_account_storage_map_values() { use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl}; use crate::db::models::queries::{ - MAX_HISTORICAL_ENTRIES_PER_KEY, + HISTORICAL_BLOCK_RETENTION, cleanup_old_account_storage_map_values, }; let mut conn = create_db(); let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE).unwrap(); - // Create blocks - for i in 1..=70 { + // Create blocks 1-TEST_CHAIN_LENGTH + for i in 1..=TEST_CHAIN_LENGTH { create_block(&mut conn, BlockNumber::from(i)); } @@ -2470,8 +2475,8 @@ fn test_cleanup_old_account_storage_map_values() { let key = num_to_word(123); let value_base = num_to_word(456); - // Insert 70 storage map value entries for the same (slot_name, key) combination - for block_num in 1..=70 { + // Insert storage map value entries for blocks 1-TEST_CHAIN_LENGTH + for block_num in 1..=TEST_CHAIN_LENGTH { queries::insert_account_storage_map_value( &mut conn, account_id, @@ -2483,41 +2488,42 @@ fn test_cleanup_old_account_storage_map_values() { .unwrap(); } - // Verify we have 70 entries using Diesel API + // Verify we have TEST_CHAIN_LENGTH entries use schema::account_storage_map_values::dsl; let count = dsl::account_storage_map_values .filter(dsl::account_id.eq(account_id.to_bytes())) .count() .get_result::(&mut conn) .unwrap(); - assert_eq!(count, 70, "Should have 70 entries before cleanup"); - - // Run cleanup - let deleted = cleanup_old_account_storage_map_values(&mut conn, account_id).unwrap(); + assert_eq!(count, TEST_CHAIN_LENGTH as i64, "Should have TEST_CHAIN_LENGTH entries before cleanup"); - // We should have deleted entries beyond MAX_HISTORICAL_ENTRIES_PER_KEY - let expected_remaining = MAX_HISTORICAL_ENTRIES_PER_KEY + 1; // +1 for the latest - let expected_deleted = 70 - expected_remaining; + // Run cleanup with chain_tip at block TEST_CHAIN_LENGTH + let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); + let deleted = cleanup_old_account_storage_map_values(&mut conn, account_id, chain_tip).unwrap(); + // cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION, blocks 1 to cutoff-1 should be deleted + let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; + let expected_deleted = cutoff - 1; assert_eq!( - deleted, expected_deleted, - "Should have deleted {} old entries", - expected_deleted + deleted, expected_deleted as usize, + "Should have deleted entries older than block {}", + cutoff ); - // Verify remaining count using Diesel API + // Verify remaining count let remaining = dsl::account_storage_map_values .filter(dsl::account_id.eq(account_id.to_bytes())) .count() .get_result::(&mut conn) .unwrap(); + let expected_remaining = TEST_CHAIN_LENGTH - expected_deleted; assert_eq!( - remaining as usize, expected_remaining, + remaining as u32, expected_remaining, "Should have {} entries remaining", expected_remaining ); - // Verify the latest entry is still marked as latest using Diesel API + // Verify the latest entry is still marked as latest let latest_count = dsl::account_storage_map_values .filter(dsl::account_id.eq(account_id.to_bytes())) .filter(dsl::is_latest.eq(true)) @@ -2550,7 +2556,7 @@ fn test_cleanup_preserves_latest_state() { ) .unwrap(); - // Test with multiple vault keys to ensure per-key cleanup + // Test with multiple vault keys to ensure all latest entries are preserved let vault_key_1 = AssetVaultKey::new_unchecked(num_to_word(100)); let vault_key_2 = AssetVaultKey::new_unchecked(num_to_word(200)); let asset = Asset::Fungible(FungibleAsset::new(account_id, 1000).unwrap()); @@ -2576,11 +2582,12 @@ fn test_cleanup_preserves_latest_state() { .unwrap(); } - // Run cleanup (should not delete anything since we're under the limit) - let deleted = cleanup_old_account_vault_assets(&mut conn, account_id).unwrap(); - assert_eq!(deleted, 0, "Should not delete anything when under limit"); + // Run cleanup with chain_tip at 10 (cutoff = 10 - 50 = 0, so nothing should be deleted) + let chain_tip = BlockNumber::from(10); + let deleted = cleanup_old_account_vault_assets(&mut conn, account_id, chain_tip).unwrap(); + assert_eq!(deleted, 0, "Should not delete anything when chain is young"); - // Verify both latest entries exist using Diesel API + // Verify both latest entries exist use schema::account_vault_assets::dsl; let latest_entries: Vec<(Vec, i64)> = dsl::account_vault_assets .filter(dsl::account_id.eq(account_id.to_bytes())) @@ -2600,7 +2607,7 @@ fn test_cleanup_preserves_latest_state() { fn test_cleanup_all_accounts() { use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl}; - use crate::db::models::queries::cleanup_all_accounts; + use crate::db::models::queries::{HISTORICAL_BLOCK_RETENTION, cleanup_all_accounts}; let mut conn = create_db(); @@ -2609,8 +2616,8 @@ fn test_cleanup_all_accounts() { let account2 = AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE_2).unwrap(); let faucet_id = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap(); - // Create blocks - for i in 1..=60 { + // Create blocks 1-TEST_CHAIN_LENGTH + for i in 1..=TEST_CHAIN_LENGTH { create_block(&mut conn, BlockNumber::from(i)); } @@ -2632,7 +2639,7 @@ fn test_cleanup_all_accounts() { let vault_key = AssetVaultKey::new_unchecked(num_to_word(100)); let asset = Asset::Fungible(FungibleAsset::new(faucet_id, 1000).unwrap()); - for block_num in 1..=60 { + for block_num in 1..=TEST_CHAIN_LENGTH { queries::insert_account_vault_asset( &mut conn, account1, @@ -2653,14 +2660,20 @@ fn test_cleanup_all_accounts() { .unwrap(); } - // Run cleanup for all accounts - let (vault_deleted, _) = cleanup_all_accounts(&mut conn).unwrap(); + // Run cleanup for all accounts with chain_tip at TEST_CHAIN_LENGTH + let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); + let (vault_deleted, _) = cleanup_all_accounts(&mut conn, chain_tip).unwrap(); - // We should have deleted entries from both accounts - assert!(vault_deleted > 0, "Should have deleted some entries"); + // Each account should have deleted entries for blocks 1 to cutoff-1 + let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; + let expected_deleted_per_account = cutoff - 1; + assert_eq!( + vault_deleted, + (expected_deleted_per_account * 2) as usize, + "Should have deleted entries from both accounts" + ); - // Each account should have MAX_HISTORICAL_ENTRIES_PER_KEY + 1 entries remaining using - // Diesel API + // Each account should have entries for blocks cutoff to TEST_CHAIN_LENGTH use schema::account_vault_assets::dsl; let count1 = dsl::account_vault_assets .filter(dsl::account_id.eq(account1.to_bytes())) @@ -2674,12 +2687,7 @@ fn test_cleanup_all_accounts() { .get_result::(&mut conn) .unwrap(); - assert_eq!( - count1 as usize, - crate::db::models::queries::MAX_HISTORICAL_ENTRIES_PER_KEY + 1 - ); - assert_eq!( - count2 as usize, - crate::db::models::queries::MAX_HISTORICAL_ENTRIES_PER_KEY + 1 - ); + let expected_remaining = TEST_CHAIN_LENGTH - expected_deleted_per_account; + assert_eq!(count1 as u32, expected_remaining); + assert_eq!(count2 as u32, expected_remaining); } diff --git a/crates/store/src/inner_forest/mod.rs b/crates/store/src/inner_forest/mod.rs index 8c144b270..1000e356f 100644 --- a/crates/store/src/inner_forest/mod.rs +++ b/crates/store/src/inner_forest/mod.rs @@ -16,9 +16,10 @@ mod tests; // CONSTANTS // ================================================================================================ -/// Maximum number of historical block entries to retain per account key in the in-memory forest. +/// Number of historical blocks to retain in the in-memory forest. +/// Entries older than `chain_tip - HISTORICAL_BLOCK_RETENTION` will be pruned. /// This matches the database retention policy in [`crate::db::models::queries::accounts`]. -pub const MAX_HISTORICAL_BLOCKS_PER_KEY: usize = 50; +pub const HISTORICAL_BLOCK_RETENTION: u32 = 50; // ERRORS // ================================================================================================ @@ -455,13 +456,7 @@ impl InnerForest { /// Prunes old entries from the in-memory forest data structures. /// - /// Removes historical entries exceeding the retention limit while preserving the most recent - /// entries for each key: - /// - For vault_roots: keeps at most MAX_HISTORICAL_BLOCKS_PER_KEY entries per account_id - /// - For storage_map_roots: keeps at most MAX_HISTORICAL_BLOCKS_PER_KEY entries per - /// (account_id, slot_name) - /// - For storage_entries: keeps at most MAX_HISTORICAL_BLOCKS_PER_KEY entries per (account_id, - /// slot_name) + /// Removes entries where `block_num < chain_tip - HISTORICAL_BLOCK_RETENTION`. /// /// The SmtForest itself is not pruned directly as it uses structural sharing and old roots /// are naturally garbage-collected when they become unreachable. @@ -470,10 +465,16 @@ impl InnerForest { /// /// A tuple containing the number of entries removed from: /// (vault_roots_removed, storage_map_roots_removed, storage_entries_removed) - pub(crate) fn prune(&mut self) -> (usize, usize, usize) { - let vault_removed = Self::prune_map_by_account(&mut self.vault_roots); - let storage_roots_removed = Self::prune_map_by_account_slot(&mut self.storage_map_roots); - let storage_entries_removed = Self::prune_map_by_account_slot(&mut self.storage_entries); + pub(crate) fn prune(&mut self, chain_tip: BlockNumber) -> (usize, usize, usize) { + let cutoff_block = BlockNumber::from( + chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION), + ); + + let vault_removed = Self::prune_map_by_block(&mut self.vault_roots, cutoff_block); + let storage_roots_removed = + Self::prune_map_by_block_slot(&mut self.storage_map_roots, cutoff_block); + let storage_entries_removed = + Self::prune_map_by_block_slot(&mut self.storage_entries, cutoff_block); if vault_removed > 0 || storage_roots_removed > 0 || storage_entries_removed > 0 { tracing::info!( @@ -481,6 +482,7 @@ impl InnerForest { vault_roots_removed = vault_removed, storage_map_roots_removed = storage_roots_removed, storage_entries_removed = storage_entries_removed, + cutoff_block = cutoff_block.as_u32(), "Forest pruning completed" ); } @@ -498,87 +500,42 @@ impl InnerForest { (vault_removed, storage_roots_removed, storage_entries_removed) } - /// Prunes a map keyed by (AccountId, BlockNumber), keeping at most - /// MAX_HISTORICAL_BLOCKS_PER_KEY entries per account. - fn prune_map_by_account(map: &mut BTreeMap<(AccountId, BlockNumber), V>) -> usize { - // Group block numbers by account_id - let mut entries_by_account: BTreeMap> = BTreeMap::new(); - for (account_id, block_num) in map.keys() { - entries_by_account.entry(*account_id).or_default().push(*block_num); - } - - let mut removed = 0; - for (account_id, mut block_nums) in entries_by_account { - let len = block_nums.len(); - if len <= MAX_HISTORICAL_BLOCKS_PER_KEY { - continue; - } + /// Prunes entries from a map keyed by (AccountId, BlockNumber) where block_num < cutoff. + fn prune_map_by_block( + map: &mut BTreeMap<(AccountId, BlockNumber), V>, + cutoff_block: BlockNumber, + ) -> usize { + // BTreeMap iteration is sorted by key, so we can efficiently find all entries to remove. + // Keys are (AccountId, BlockNumber), but we need to remove across all accounts. + let keys_to_remove: Vec<_> = map + .keys() + .filter(|(_, block_num)| *block_num < cutoff_block) + .copied() + .collect(); - // Find the cutoff block number using O(n) selection algorithm. - // We want to keep the MAX_HISTORICAL_BLOCKS_PER_KEY largest block numbers. - // select_nth_unstable_by partitions so that element at index is in sorted position. - // After this, elements at indices >= cutoff_idx are the largest. - let cutoff_idx = len - MAX_HISTORICAL_BLOCKS_PER_KEY; - block_nums.select_nth_unstable(cutoff_idx); - let cutoff_block = block_nums[cutoff_idx]; - - // Remove entries with block_num < cutoff_block - // Since BTreeMap is sorted, we can use range to find entries to remove. - let keys_to_remove: Vec<_> = map - .range((account_id, BlockNumber::GENESIS)..(account_id, cutoff_block)) - .map(|(k, _)| *k) - .collect(); - - for key in keys_to_remove { - map.remove(&key); - removed += 1; - } + let removed = keys_to_remove.len(); + for key in keys_to_remove { + map.remove(&key); } removed } - /// Prunes a map keyed by (AccountId, StorageSlotName, BlockNumber), keeping at most - /// MAX_HISTORICAL_BLOCKS_PER_KEY entries per (account_id, slot_name). - fn prune_map_by_account_slot( + /// Prunes entries from a map keyed by (AccountId, StorageSlotName, BlockNumber) + /// where block_num < cutoff. + fn prune_map_by_block_slot( map: &mut BTreeMap<(AccountId, StorageSlotName, BlockNumber), V>, + cutoff_block: BlockNumber, ) -> usize { - // Group block numbers by (account_id, slot_name) - // Clone slot_name to avoid holding references to map during mutation - let mut entries_by_key: BTreeMap<(AccountId, StorageSlotName), Vec> = - BTreeMap::new(); - for (account_id, slot_name, block_num) in map.keys() { - entries_by_key - .entry((*account_id, slot_name.clone())) - .or_default() - .push(*block_num); - } - - let mut removed = 0; - for ((account_id, slot_name), mut block_nums) in entries_by_key { - let len = block_nums.len(); - if len <= MAX_HISTORICAL_BLOCKS_PER_KEY { - continue; - } + let keys_to_remove: Vec<_> = map + .keys() + .filter(|(_, _, block_num)| *block_num < cutoff_block) + .cloned() + .collect(); - // Find the cutoff block number using O(n) selection algorithm - let cutoff_idx = len - MAX_HISTORICAL_BLOCKS_PER_KEY; - block_nums.select_nth_unstable(cutoff_idx); - let cutoff_block = block_nums[cutoff_idx]; - - // Remove entries with block_num < cutoff_block - let keys_to_remove: Vec<_> = map - .range( - (account_id, slot_name.clone(), BlockNumber::GENESIS) - ..(account_id, slot_name.clone(), cutoff_block), - ) - .map(|(k, _)| k.clone()) - .collect(); - - for key in keys_to_remove { - map.remove(&key); - removed += 1; - } + let removed = keys_to_remove.len(); + for key in keys_to_remove { + map.remove(&key); } removed diff --git a/crates/store/src/inner_forest/tests.rs b/crates/store/src/inner_forest/tests.rs index f55bb3cd9..85589d86a 100644 --- a/crates/store/src/inner_forest/tests.rs +++ b/crates/store/src/inner_forest/tests.rs @@ -465,16 +465,20 @@ fn test_open_storage_map_returns_limit_exceeded_for_too_many_keys() { // PRUNING TESTS // ================================================================================================ +/// Chain length used in pruning tests - must be > HISTORICAL_BLOCK_RETENTION for meaningful tests. +const TEST_CHAIN_LENGTH: u32 = 100; + #[test] fn test_prune_vault_roots_removes_old_entries() { + use super::HISTORICAL_BLOCK_RETENTION; + let mut forest = InnerForest::new(); let account_id = dummy_account(); let faucet_id = dummy_faucet(); - // Add entries for MAX_HISTORICAL_BLOCKS_PER_KEY + 10 blocks - let num_blocks = MAX_HISTORICAL_BLOCKS_PER_KEY + 10; - for i in 1..=num_blocks { - let block_num = BlockNumber::from(i as u32); + // Add entries for blocks 1-TEST_CHAIN_LENGTH + for i in 1..=TEST_CHAIN_LENGTH { + let block_num = BlockNumber::from(i); let amount = (i * 100) as u64; let mut vault_delta = AccountVaultDelta::default(); vault_delta.add_asset(dummy_fungible_asset(faucet_id, amount)).unwrap(); @@ -483,22 +487,25 @@ fn test_prune_vault_roots_removes_old_entries() { } // Verify we have all entries before pruning - assert_eq!(forest.vault_roots.len(), num_blocks); + assert_eq!(forest.vault_roots.len(), TEST_CHAIN_LENGTH as usize); - // Prune - let (vault_removed, ..) = forest.prune(); + // Prune with chain_tip at TEST_CHAIN_LENGTH (cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION) + let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); + let (vault_removed, ..) = forest.prune(chain_tip); - // Should have removed 10 entries - assert_eq!(vault_removed, 10); + // Should have removed blocks 1 to cutoff-1 + let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; + let expected_removed = cutoff - 1; // blocks 1 to cutoff-1 + assert_eq!(vault_removed, expected_removed as usize); - // Should have exactly MAX_HISTORICAL_BLOCKS_PER_KEY entries remaining - assert_eq!(forest.vault_roots.len(), MAX_HISTORICAL_BLOCKS_PER_KEY); + // Should have blocks cutoff to TEST_CHAIN_LENGTH remaining + let expected_remaining = TEST_CHAIN_LENGTH - expected_removed; + assert_eq!(forest.vault_roots.len(), expected_remaining as usize); - // The remaining entries should be the most recent ones - let remaining_blocks: Vec<_> = forest.vault_roots.keys().map(|(_, b)| b.as_u32()).collect(); + // The remaining entries should be the most recent ones (>= cutoff) + let remaining_blocks = Vec::from_iter(forest.vault_roots.keys().map(|(_, b)| b.as_u32())); let oldest_remaining = *remaining_blocks.iter().min().unwrap(); - // Oldest remaining should be block 11 (since we kept the most recent 50 out of 60) - assert_eq!(oldest_remaining, 11); + assert_eq!(oldest_remaining, cutoff); } #[test] @@ -507,16 +514,17 @@ fn test_prune_storage_map_roots_removes_old_entries() { use miden_protocol::account::delta::{StorageMapDelta, StorageSlotDelta}; + use super::HISTORICAL_BLOCK_RETENTION; + let mut forest = InnerForest::new(); let account_id = dummy_account(); let slot_name = StorageSlotName::mock(3); - // Add entries for MAX_HISTORICAL_BLOCKS_PER_KEY + 5 blocks - let num_blocks = MAX_HISTORICAL_BLOCKS_PER_KEY + 5; - for i in 1..=num_blocks { - let block_num = BlockNumber::from(i as u32); - let key = Word::from([i as u32, 0, 0, 0]); - let value = Word::from([0, 0, 0, i as u32]); + // Add entries for blocks 1-TEST_CHAIN_LENGTH + for i in 1..=TEST_CHAIN_LENGTH { + let block_num = BlockNumber::from(i); + let key = Word::from([i, 0, 0, 0]); + let value = Word::from([0, 0, 0, i]); let mut map_delta = StorageMapDelta::default(); map_delta.insert(key, value); @@ -527,28 +535,32 @@ fn test_prune_storage_map_roots_removes_old_entries() { } // Verify we have all entries before pruning - assert_eq!(forest.storage_map_roots.len(), num_blocks); - - // Prune - let (_, storage_roots_removed, storage_entries_removed) = forest.prune(); - - // Should have removed 5 entries from each structure - assert_eq!(storage_roots_removed, 5); - assert_eq!(storage_entries_removed, 5); - - // Should have exactly MAX_HISTORICAL_BLOCKS_PER_KEY entries remaining - assert_eq!(forest.storage_map_roots.len(), MAX_HISTORICAL_BLOCKS_PER_KEY); - assert_eq!(forest.storage_entries.len(), MAX_HISTORICAL_BLOCKS_PER_KEY); + assert_eq!(forest.storage_map_roots.len(), TEST_CHAIN_LENGTH as usize); + + // Prune with chain_tip at TEST_CHAIN_LENGTH + let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); + let (_, storage_roots_removed, storage_entries_removed) = forest.prune(chain_tip); + + // Should have removed blocks 1 to cutoff-1 + let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; + let expected_removed = cutoff - 1; + assert_eq!(storage_roots_removed, expected_removed as usize); + assert_eq!(storage_entries_removed, expected_removed as usize); + + // Should have blocks cutoff to TEST_CHAIN_LENGTH remaining + let expected_remaining = TEST_CHAIN_LENGTH - expected_removed; + assert_eq!(forest.storage_map_roots.len(), expected_remaining as usize); + assert_eq!(forest.storage_entries.len(), expected_remaining as usize); } #[test] -fn test_prune_does_nothing_when_under_limit() { +fn test_prune_does_nothing_when_chain_is_young() { let mut forest = InnerForest::new(); let account_id = dummy_account(); let faucet_id = dummy_faucet(); - // Add fewer entries than the limit - let num_blocks = 10; + // Add entries for blocks 1-10 + let num_blocks = 10usize; for i in 1..=num_blocks { let block_num = BlockNumber::from(i as u32); let mut vault_delta = AccountVaultDelta::default(); @@ -559,8 +571,9 @@ fn test_prune_does_nothing_when_under_limit() { forest.update_account(block_num, &delta).unwrap(); } - // Prune - let (vault_removed, storage_roots_removed, storage_entries_removed) = forest.prune(); + // Prune with chain_tip at 10 (cutoff = 10 - 50 = 0, nothing to remove) + let chain_tip = BlockNumber::from(10); + let (vault_removed, storage_roots_removed, storage_entries_removed) = forest.prune(chain_tip); // Nothing should be removed assert_eq!(vault_removed, 0); @@ -573,16 +586,16 @@ fn test_prune_does_nothing_when_under_limit() { #[test] fn test_prune_handles_multiple_accounts() { + use super::HISTORICAL_BLOCK_RETENTION; + let mut forest = InnerForest::new(); let account1 = dummy_account(); let account2 = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap(); let faucet_id = dummy_faucet(); - // Add entries for multiple accounts, each exceeding the limit - let num_blocks = MAX_HISTORICAL_BLOCKS_PER_KEY + 5; - - for i in 1..=num_blocks { - let block_num = BlockNumber::from(i as u32); + // Add entries for blocks 1-TEST_CHAIN_LENGTH for both accounts + for i in 1..=TEST_CHAIN_LENGTH { + let block_num = BlockNumber::from(i); let amount = (i * 100) as u64; // Account 1 @@ -599,17 +612,21 @@ fn test_prune_handles_multiple_accounts() { } // Verify we have entries for both accounts before pruning - assert_eq!(forest.vault_roots.len(), num_blocks * 2); + assert_eq!(forest.vault_roots.len(), (TEST_CHAIN_LENGTH * 2) as usize); - // Prune - let (vault_removed, ..) = forest.prune(); + // Prune with chain_tip at TEST_CHAIN_LENGTH + let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); + let (vault_removed, ..) = forest.prune(chain_tip); - // Should have removed 5 entries from each account = 10 total - assert_eq!(vault_removed, 10); + // Should have removed blocks 1 to cutoff-1 from both accounts + let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; + let expected_removed_per_account = cutoff - 1; + assert_eq!(vault_removed, (expected_removed_per_account * 2) as usize); - // Each account should have exactly MAX_HISTORICAL_BLOCKS_PER_KEY entries + // Each account should have blocks cutoff to TEST_CHAIN_LENGTH remaining + let expected_remaining_per_account = TEST_CHAIN_LENGTH - expected_removed_per_account; let account1_entries = forest.vault_roots.keys().filter(|(id, _)| *id == account1).count(); let account2_entries = forest.vault_roots.keys().filter(|(id, _)| *id == account2).count(); - assert_eq!(account1_entries, MAX_HISTORICAL_BLOCKS_PER_KEY); - assert_eq!(account2_entries, MAX_HISTORICAL_BLOCKS_PER_KEY); + assert_eq!(account1_entries, expected_remaining_per_account as usize); + assert_eq!(account2_entries, expected_remaining_per_account as usize); } diff --git a/crates/store/src/state.rs b/crates/store/src/state.rs index 1172d152f..0d7e008b9 100644 --- a/crates/store/src/state.rs +++ b/crates/store/src/state.rs @@ -573,7 +573,8 @@ impl State { forest.apply_block_updates(block_num, account_deltas)?; // Prune old entries from the in-memory forest while holding both locks - let (vault_pruned, storage_roots_pruned, storage_entries_pruned) = forest.prune(); + let (vault_pruned, storage_roots_pruned, storage_entries_pruned) = + forest.prune(block_num); if vault_pruned > 0 || storage_roots_pruned > 0 || storage_entries_pruned > 0 { tracing::debug!( target: COMPONENT, From 6b0a8fcff0a76221b780618677b678cbc59d622c Mon Sep 17 00:00:00 2001 From: Bernhard Schuster Date: Tue, 20 Jan 2026 13:55:29 +0100 Subject: [PATCH 4/5] fmt clippy --- crates/store/src/db/models/queries/accounts.rs | 6 +++--- crates/store/src/db/tests.rs | 18 ++++++++++++++---- crates/store/src/inner_forest/mod.rs | 12 ++++-------- crates/store/src/inner_forest/tests.rs | 3 ++- 4 files changed, 23 insertions(+), 16 deletions(-) diff --git a/crates/store/src/db/models/queries/accounts.rs b/crates/store/src/db/models/queries/accounts.rs index ff0722e42..a19bc85bb 100644 --- a/crates/store/src/db/models/queries/accounts.rs +++ b/crates/store/src/db/models/queries/accounts.rs @@ -1190,7 +1190,7 @@ pub(crate) fn cleanup_old_account_vault_assets( chain_tip: BlockNumber, ) -> Result { let account_id_bytes = account_id.to_bytes(); - let cutoff_block = chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION) as i64; + let cutoff_block = i64::from(chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION)); diesel::sql_query( r#" @@ -1215,7 +1215,7 @@ pub(crate) fn cleanup_old_account_storage_map_values( chain_tip: BlockNumber, ) -> Result { let account_id_bytes = account_id.to_bytes(); - let cutoff_block = chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION) as i64; + let cutoff_block = i64::from(chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION)); diesel::sql_query( r#" @@ -1240,7 +1240,7 @@ pub fn cleanup_all_accounts( conn: &mut SqliteConnection, chain_tip: BlockNumber, ) -> Result<(usize, usize), DatabaseError> { - let cutoff_block = chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION) as i64; + let cutoff_block = i64::from(chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION)); let vault_deleted = diesel::sql_query( r#" diff --git a/crates/store/src/db/tests.rs b/crates/store/src/db/tests.rs index 217203af2..3f0120c38 100644 --- a/crates/store/src/db/tests.rs +++ b/crates/store/src/db/tests.rs @@ -2403,7 +2403,10 @@ fn test_cleanup_old_account_vault_assets() { .count() .get_result::(&mut conn) .unwrap(); - assert_eq!(count, TEST_CHAIN_LENGTH as i64, "Should have TEST_CHAIN_LENGTH entries before cleanup"); + assert_eq!( + count, TEST_CHAIN_LENGTH as i64, + "Should have TEST_CHAIN_LENGTH entries before cleanup" + ); // Run cleanup with chain_tip at block TEST_CHAIN_LENGTH // cutoff = TEST_CHAIN_LENGTH - 50, so blocks < cutoff should be deleted @@ -2450,7 +2453,10 @@ fn test_cleanup_old_account_vault_assets() { .select(dsl::block_num) .first::(&mut conn) .unwrap(); - assert_eq!(latest_block, TEST_CHAIN_LENGTH as i64, "Latest entry should be from block TEST_CHAIN_LENGTH"); + assert_eq!( + latest_block, TEST_CHAIN_LENGTH as i64, + "Latest entry should be from block TEST_CHAIN_LENGTH" + ); } #[test] @@ -2495,13 +2501,17 @@ fn test_cleanup_old_account_storage_map_values() { .count() .get_result::(&mut conn) .unwrap(); - assert_eq!(count, TEST_CHAIN_LENGTH as i64, "Should have TEST_CHAIN_LENGTH entries before cleanup"); + assert_eq!( + count, TEST_CHAIN_LENGTH as i64, + "Should have TEST_CHAIN_LENGTH entries before cleanup" + ); // Run cleanup with chain_tip at block TEST_CHAIN_LENGTH let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); let deleted = cleanup_old_account_storage_map_values(&mut conn, account_id, chain_tip).unwrap(); - // cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION, blocks 1 to cutoff-1 should be deleted + // cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION, blocks 1 to cutoff-1 should be + // deleted let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; let expected_deleted = cutoff - 1; assert_eq!( diff --git a/crates/store/src/inner_forest/mod.rs b/crates/store/src/inner_forest/mod.rs index 1000e356f..e4d4b03a7 100644 --- a/crates/store/src/inner_forest/mod.rs +++ b/crates/store/src/inner_forest/mod.rs @@ -466,9 +466,8 @@ impl InnerForest { /// A tuple containing the number of entries removed from: /// (vault_roots_removed, storage_map_roots_removed, storage_entries_removed) pub(crate) fn prune(&mut self, chain_tip: BlockNumber) -> (usize, usize, usize) { - let cutoff_block = BlockNumber::from( - chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION), - ); + let cutoff_block = + BlockNumber::from(chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION)); let vault_removed = Self::prune_map_by_block(&mut self.vault_roots, cutoff_block); let storage_roots_removed = @@ -507,11 +506,8 @@ impl InnerForest { ) -> usize { // BTreeMap iteration is sorted by key, so we can efficiently find all entries to remove. // Keys are (AccountId, BlockNumber), but we need to remove across all accounts. - let keys_to_remove: Vec<_> = map - .keys() - .filter(|(_, block_num)| *block_num < cutoff_block) - .copied() - .collect(); + let keys_to_remove: Vec<_> = + map.keys().filter(|(_, block_num)| *block_num < cutoff_block).copied().collect(); let removed = keys_to_remove.len(); for key in keys_to_remove { diff --git a/crates/store/src/inner_forest/tests.rs b/crates/store/src/inner_forest/tests.rs index 85589d86a..0a19c67a5 100644 --- a/crates/store/src/inner_forest/tests.rs +++ b/crates/store/src/inner_forest/tests.rs @@ -489,7 +489,8 @@ fn test_prune_vault_roots_removes_old_entries() { // Verify we have all entries before pruning assert_eq!(forest.vault_roots.len(), TEST_CHAIN_LENGTH as usize); - // Prune with chain_tip at TEST_CHAIN_LENGTH (cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION) + // Prune with chain_tip at TEST_CHAIN_LENGTH (cutoff = TEST_CHAIN_LENGTH - + // HISTORICAL_BLOCK_RETENTION) let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); let (vault_removed, ..) = forest.prune(chain_tip); From 2c589898c3b0e373c67a0420bc7714c148e12d57 Mon Sep 17 00:00:00 2001 From: Bernhard Schuster Date: Tue, 3 Feb 2026 09:48:57 +0100 Subject: [PATCH 5/5] remove any in memory cleanup --- CHANGELOG.md | 2 +- crates/store/src/db/mod.rs | 1 + crates/store/src/inner_forest/mod.rs | 94 +------------- crates/store/src/inner_forest/tests.rs | 170 ------------------------- crates/store/src/state.rs | 22 +--- 5 files changed, 5 insertions(+), 284 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1da045315..b35696bca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### Enhancements -- Added periodic cleanup of old account data from the database and in-memory forest ([#1304](https://github.com/0xMiden/miden-node/issues/1304)). +- Added periodic cleanup of old account data from the database ([#1304](https://github.com/0xMiden/miden-node/issues/1304)). - Added support for timeouts in the WASM remote prover clients ([#1383](https://github.com/0xMiden/miden-node/pull/1383)). - Added block validation endpoint to validator and integrated with block producer ([#1382](https://github.com/0xMiden/miden-node/pull/1381)). - Added support for caching mempool statistics in the block producer server ([#1388](https://github.com/0xMiden/miden-node/pull/1388)). diff --git a/crates/store/src/db/mod.rs b/crates/store/src/db/mod.rs index ff9300b11..e27d82a5d 100644 --- a/crates/store/src/db/mod.rs +++ b/crates/store/src/db/mod.rs @@ -328,6 +328,7 @@ impl Db { let me2 = me.clone(); // Spawn background cleanup task + // TODO: retain the join handle to coordinate shutdown or surface task failures. let _cleanup_task_handle = tokio::spawn(async move { Self::periodic_cleanup_task(me2, rx).await }); diff --git a/crates/store/src/inner_forest/mod.rs b/crates/store/src/inner_forest/mod.rs index e4d4b03a7..b9bba7aeb 100644 --- a/crates/store/src/inner_forest/mod.rs +++ b/crates/store/src/inner_forest/mod.rs @@ -13,14 +13,6 @@ use thiserror::Error; #[cfg(test)] mod tests; -// CONSTANTS -// ================================================================================================ - -/// Number of historical blocks to retain in the in-memory forest. -/// Entries older than `chain_tip - HISTORICAL_BLOCK_RETENTION` will be pruned. -/// This matches the database retention policy in [`crate::db::models::queries::accounts`]. -pub const HISTORICAL_BLOCK_RETENTION: u32 = 50; - // ERRORS // ================================================================================================ @@ -451,89 +443,5 @@ impl InnerForest { } } - // PRUNING - // -------------------------------------------------------------------------------------------- - - /// Prunes old entries from the in-memory forest data structures. - /// - /// Removes entries where `block_num < chain_tip - HISTORICAL_BLOCK_RETENTION`. - /// - /// The SmtForest itself is not pruned directly as it uses structural sharing and old roots - /// are naturally garbage-collected when they become unreachable. - /// - /// # Returns - /// - /// A tuple containing the number of entries removed from: - /// (vault_roots_removed, storage_map_roots_removed, storage_entries_removed) - pub(crate) fn prune(&mut self, chain_tip: BlockNumber) -> (usize, usize, usize) { - let cutoff_block = - BlockNumber::from(chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION)); - - let vault_removed = Self::prune_map_by_block(&mut self.vault_roots, cutoff_block); - let storage_roots_removed = - Self::prune_map_by_block_slot(&mut self.storage_map_roots, cutoff_block); - let storage_entries_removed = - Self::prune_map_by_block_slot(&mut self.storage_entries, cutoff_block); - - if vault_removed > 0 || storage_roots_removed > 0 || storage_entries_removed > 0 { - tracing::info!( - target: crate::COMPONENT, - vault_roots_removed = vault_removed, - storage_map_roots_removed = storage_roots_removed, - storage_entries_removed = storage_entries_removed, - cutoff_block = cutoff_block.as_u32(), - "Forest pruning completed" - ); - } - - // Prune the underlying SmtForest by telling it which roots are still referenced. - // Collect all roots that we're keeping from both vault_roots and storage_map_roots. - let roots_to_keep: Vec = self - .vault_roots - .values() - .chain(self.storage_map_roots.values()) - .copied() - .collect(); - self.forest.pop_smts(roots_to_keep); - - (vault_removed, storage_roots_removed, storage_entries_removed) - } - - /// Prunes entries from a map keyed by (AccountId, BlockNumber) where block_num < cutoff. - fn prune_map_by_block( - map: &mut BTreeMap<(AccountId, BlockNumber), V>, - cutoff_block: BlockNumber, - ) -> usize { - // BTreeMap iteration is sorted by key, so we can efficiently find all entries to remove. - // Keys are (AccountId, BlockNumber), but we need to remove across all accounts. - let keys_to_remove: Vec<_> = - map.keys().filter(|(_, block_num)| *block_num < cutoff_block).copied().collect(); - - let removed = keys_to_remove.len(); - for key in keys_to_remove { - map.remove(&key); - } - - removed - } - - /// Prunes entries from a map keyed by (AccountId, StorageSlotName, BlockNumber) - /// where block_num < cutoff. - fn prune_map_by_block_slot( - map: &mut BTreeMap<(AccountId, StorageSlotName, BlockNumber), V>, - cutoff_block: BlockNumber, - ) -> usize { - let keys_to_remove: Vec<_> = map - .keys() - .filter(|(_, _, block_num)| *block_num < cutoff_block) - .cloned() - .collect(); - - let removed = keys_to_remove.len(); - for key in keys_to_remove { - map.remove(&key); - } - - removed - } + // TODO: tie in-memory forest retention to DB pruning policy once forest queries rely on it. } diff --git a/crates/store/src/inner_forest/tests.rs b/crates/store/src/inner_forest/tests.rs index 0a19c67a5..da311db46 100644 --- a/crates/store/src/inner_forest/tests.rs +++ b/crates/store/src/inner_forest/tests.rs @@ -461,173 +461,3 @@ fn test_open_storage_map_returns_limit_exceeded_for_too_many_keys() { let details = result.expect("Should return Some").expect("Should not error"); assert_matches!(details.entries, StorageMapEntries::LimitExceeded); } - -// PRUNING TESTS -// ================================================================================================ - -/// Chain length used in pruning tests - must be > HISTORICAL_BLOCK_RETENTION for meaningful tests. -const TEST_CHAIN_LENGTH: u32 = 100; - -#[test] -fn test_prune_vault_roots_removes_old_entries() { - use super::HISTORICAL_BLOCK_RETENTION; - - let mut forest = InnerForest::new(); - let account_id = dummy_account(); - let faucet_id = dummy_faucet(); - - // Add entries for blocks 1-TEST_CHAIN_LENGTH - for i in 1..=TEST_CHAIN_LENGTH { - let block_num = BlockNumber::from(i); - let amount = (i * 100) as u64; - let mut vault_delta = AccountVaultDelta::default(); - vault_delta.add_asset(dummy_fungible_asset(faucet_id, amount)).unwrap(); - let delta = dummy_partial_delta(account_id, vault_delta, AccountStorageDelta::default()); - forest.update_account(block_num, &delta).unwrap(); - } - - // Verify we have all entries before pruning - assert_eq!(forest.vault_roots.len(), TEST_CHAIN_LENGTH as usize); - - // Prune with chain_tip at TEST_CHAIN_LENGTH (cutoff = TEST_CHAIN_LENGTH - - // HISTORICAL_BLOCK_RETENTION) - let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); - let (vault_removed, ..) = forest.prune(chain_tip); - - // Should have removed blocks 1 to cutoff-1 - let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; - let expected_removed = cutoff - 1; // blocks 1 to cutoff-1 - assert_eq!(vault_removed, expected_removed as usize); - - // Should have blocks cutoff to TEST_CHAIN_LENGTH remaining - let expected_remaining = TEST_CHAIN_LENGTH - expected_removed; - assert_eq!(forest.vault_roots.len(), expected_remaining as usize); - - // The remaining entries should be the most recent ones (>= cutoff) - let remaining_blocks = Vec::from_iter(forest.vault_roots.keys().map(|(_, b)| b.as_u32())); - let oldest_remaining = *remaining_blocks.iter().min().unwrap(); - assert_eq!(oldest_remaining, cutoff); -} - -#[test] -fn test_prune_storage_map_roots_removes_old_entries() { - use std::collections::BTreeMap; - - use miden_protocol::account::delta::{StorageMapDelta, StorageSlotDelta}; - - use super::HISTORICAL_BLOCK_RETENTION; - - let mut forest = InnerForest::new(); - let account_id = dummy_account(); - let slot_name = StorageSlotName::mock(3); - - // Add entries for blocks 1-TEST_CHAIN_LENGTH - for i in 1..=TEST_CHAIN_LENGTH { - let block_num = BlockNumber::from(i); - let key = Word::from([i, 0, 0, 0]); - let value = Word::from([0, 0, 0, i]); - - let mut map_delta = StorageMapDelta::default(); - map_delta.insert(key, value); - let raw = BTreeMap::from_iter([(slot_name.clone(), StorageSlotDelta::Map(map_delta))]); - let storage_delta = AccountStorageDelta::from_raw(raw); - let delta = dummy_partial_delta(account_id, AccountVaultDelta::default(), storage_delta); - forest.update_account(block_num, &delta).unwrap(); - } - - // Verify we have all entries before pruning - assert_eq!(forest.storage_map_roots.len(), TEST_CHAIN_LENGTH as usize); - - // Prune with chain_tip at TEST_CHAIN_LENGTH - let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); - let (_, storage_roots_removed, storage_entries_removed) = forest.prune(chain_tip); - - // Should have removed blocks 1 to cutoff-1 - let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; - let expected_removed = cutoff - 1; - assert_eq!(storage_roots_removed, expected_removed as usize); - assert_eq!(storage_entries_removed, expected_removed as usize); - - // Should have blocks cutoff to TEST_CHAIN_LENGTH remaining - let expected_remaining = TEST_CHAIN_LENGTH - expected_removed; - assert_eq!(forest.storage_map_roots.len(), expected_remaining as usize); - assert_eq!(forest.storage_entries.len(), expected_remaining as usize); -} - -#[test] -fn test_prune_does_nothing_when_chain_is_young() { - let mut forest = InnerForest::new(); - let account_id = dummy_account(); - let faucet_id = dummy_faucet(); - - // Add entries for blocks 1-10 - let num_blocks = 10usize; - for i in 1..=num_blocks { - let block_num = BlockNumber::from(i as u32); - let mut vault_delta = AccountVaultDelta::default(); - vault_delta - .add_asset(dummy_fungible_asset(faucet_id, (i * 100) as u64)) - .unwrap(); - let delta = dummy_partial_delta(account_id, vault_delta, AccountStorageDelta::default()); - forest.update_account(block_num, &delta).unwrap(); - } - - // Prune with chain_tip at 10 (cutoff = 10 - 50 = 0, nothing to remove) - let chain_tip = BlockNumber::from(10); - let (vault_removed, storage_roots_removed, storage_entries_removed) = forest.prune(chain_tip); - - // Nothing should be removed - assert_eq!(vault_removed, 0); - assert_eq!(storage_roots_removed, 0); - assert_eq!(storage_entries_removed, 0); - - // All entries should remain - assert_eq!(forest.vault_roots.len(), num_blocks); -} - -#[test] -fn test_prune_handles_multiple_accounts() { - use super::HISTORICAL_BLOCK_RETENTION; - - let mut forest = InnerForest::new(); - let account1 = dummy_account(); - let account2 = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap(); - let faucet_id = dummy_faucet(); - - // Add entries for blocks 1-TEST_CHAIN_LENGTH for both accounts - for i in 1..=TEST_CHAIN_LENGTH { - let block_num = BlockNumber::from(i); - let amount = (i * 100) as u64; - - // Account 1 - let mut vault_delta1 = AccountVaultDelta::default(); - vault_delta1.add_asset(dummy_fungible_asset(faucet_id, amount)).unwrap(); - let delta1 = dummy_partial_delta(account1, vault_delta1, AccountStorageDelta::default()); - forest.update_account(block_num, &delta1).unwrap(); - - // Account 2 (using a different faucet for variety) - let mut vault_delta2 = AccountVaultDelta::default(); - vault_delta2.add_asset(dummy_fungible_asset(account2, amount * 2)).unwrap(); - let delta2 = dummy_partial_delta(account2, vault_delta2, AccountStorageDelta::default()); - forest.update_account(block_num, &delta2).unwrap(); - } - - // Verify we have entries for both accounts before pruning - assert_eq!(forest.vault_roots.len(), (TEST_CHAIN_LENGTH * 2) as usize); - - // Prune with chain_tip at TEST_CHAIN_LENGTH - let chain_tip = BlockNumber::from(TEST_CHAIN_LENGTH); - let (vault_removed, ..) = forest.prune(chain_tip); - - // Should have removed blocks 1 to cutoff-1 from both accounts - let cutoff = TEST_CHAIN_LENGTH - HISTORICAL_BLOCK_RETENTION; - let expected_removed_per_account = cutoff - 1; - assert_eq!(vault_removed, (expected_removed_per_account * 2) as usize); - - // Each account should have blocks cutoff to TEST_CHAIN_LENGTH remaining - let expected_remaining_per_account = TEST_CHAIN_LENGTH - expected_removed_per_account; - let account1_entries = forest.vault_roots.keys().filter(|(id, _)| *id == account1).count(); - let account2_entries = forest.vault_roots.keys().filter(|(id, _)| *id == account2).count(); - assert_eq!(account1_entries, expected_remaining_per_account as usize); - assert_eq!(account2_entries, expected_remaining_per_account as usize); -} diff --git a/crates/store/src/state.rs b/crates/store/src/state.rs index 0d7e008b9..2ff1f887b 100644 --- a/crates/store/src/state.rs +++ b/crates/store/src/state.rs @@ -565,28 +565,10 @@ impl State { .apply_mutations(account_tree_update) .expect("Unreachable: old account tree root must be checked before this step"); inner.blockchain.push(block_commitment); - - // Update forest while still holding the inner lock to ensure consistency. - // Readers acquiring the inner read lock will see both account tree and forest - // updated atomically. - let mut forest = self.forest.write().await; - forest.apply_block_updates(block_num, account_deltas)?; - - // Prune old entries from the in-memory forest while holding both locks - let (vault_pruned, storage_roots_pruned, storage_entries_pruned) = - forest.prune(block_num); - if vault_pruned > 0 || storage_roots_pruned > 0 || storage_entries_pruned > 0 { - tracing::debug!( - target: COMPONENT, - block_num = block_num.as_u32(), - vault_pruned, - storage_roots_pruned, - storage_entries_pruned, - "Forest pruning applied" - ); - } } + self.forest.write().await.apply_block_updates(block_num, account_deltas)?; + info!(%block_commitment, block_num = block_num.as_u32(), COMPONENT, "apply_block successful"); Ok(())