From de52d8204b08e52b4dc719d7b4e59b6c7ef93b2c Mon Sep 17 00:00:00 2001 From: Francis Li Date: Wed, 2 Sep 2026 13:00:35 -0500 Subject: [PATCH] backport(paradigmxyz/reth#26766): perf(storage): parallelize sharded history last-shard reads Squashed https://github.com/paradigmxyz/reth/pull/26766 at f47fd63a3762d8add2e07f1d4afa9363e60d0853. (cherry picked from commit 76a826109009c936c2ea98766e0b6b50d2803d26) --- .../src/stages/index_account_history.rs | 244 ++++++- .../src/stages/index_storage_history.rs | 253 +++++++- crates/stages/stages/src/stages/utils.rs | 599 ++++++++++++++---- crates/storage/provider/src/either_writer.rs | 30 +- crates/storage/provider/src/history_shards.rs | 471 ++++++++++++++ crates/storage/provider/src/lib.rs | 7 + .../src/providers/database/provider.rs | 22 +- .../src/providers/rocksdb/provider.rs | 155 ++--- 8 files changed, 1448 insertions(+), 333 deletions(-) create mode 100644 crates/storage/provider/src/history_shards.rs diff --git a/crates/stages/stages/src/stages/index_account_history.rs b/crates/stages/stages/src/stages/index_account_history.rs index 617be0a5ee7..5675da01d56 100644 --- a/crates/stages/stages/src/stages/index_account_history.rs +++ b/crates/stages/stages/src/stages/index_account_history.rs @@ -1,5 +1,8 @@ use super::collect_account_history_indices; -use crate::stages::utils::{collect_history_indices, load_account_history}; +use crate::stages::utils::{ + collect_history_indices, load_account_history_append, prepare_account_history_writes, + write_prepared_history_shards, +}; use reth_config::config::{EtlConfig, IndexHistoryConfig}; use reth_db_api::{models::ShardedKey, tables, transaction::DbTxMut, Tables}; use reth_provider::{ @@ -8,7 +11,8 @@ use reth_provider::{ }; use reth_prune_types::{PruneCheckpoint, PruneMode, PrunePurpose, PruneSegment}; use reth_stages_api::{ - ExecInput, ExecOutput, Stage, StageCheckpoint, StageError, StageId, UnwindInput, UnwindOutput, + BlockRangeOutput, ExecInput, ExecOutput, Stage, StageCheckpoint, StageError, StageId, + UnwindInput, UnwindOutput, }; use std::fmt::Debug; use tracing::info; @@ -66,6 +70,10 @@ where provider: &Provider, mut input: ExecInput, ) -> Result { + let initial_sync = input.checkpoint.is_none(); + let rebuild_from_empty = input.checkpoint().block_number == 0; + let mut prune_target_applied = false; + if let Some((target_prunable_block, prune_mode)) = self .prune_mode .map(|mode| { @@ -77,8 +85,10 @@ where }) .transpose()? .flatten() && - target_prunable_block > input.checkpoint().block_number + (target_prunable_block > input.checkpoint().block_number || + target_prunable_block == 0 && input.checkpoint().block_number == 0) { + prune_target_applied = true; input.checkpoint = Some(StageCheckpoint::new(target_prunable_block)); // Save prune checkpoint only if we don't have one already. @@ -95,38 +105,50 @@ where } } - if input.target_reached() { + let use_rocksdb = provider.cached_storage_settings().storage_v2; + if input.target_reached() && prune_target_applied { + if use_rocksdb { + provider.rocksdb_provider().clear::()?; + } else { + provider.tx_ref().clear::()?; + } + return Ok(ExecOutput::done(input.checkpoint())) + } + if input.target_reached() && !initial_sync { return Ok(ExecOutput::done(input.checkpoint())) } - let mut range = input.next_block_range(); - let first_sync = input.checkpoint().block_number == 0; - let use_rocksdb = provider.cached_storage_settings().storage_v2; + let BlockRangeOutput { mut block_range, is_final_range } = + if initial_sync && input.target() == 0 { + BlockRangeOutput { block_range: 0..=0, is_final_range: true } + } else { + input.next_block_range_with_threshold(self.commit_threshold.max(1)) + }; - // On first sync we might have history coming from genesis. We clear the table since it's - // faster to rebuild from scratch. - if first_sync { + // A zero checkpoint means no history range is durable. Clear stale rows before using the + // append-only loader, including when pruning advanced the in-memory checkpoint. + if rebuild_from_empty { if use_rocksdb { - // Note: RocksDB clear() executes immediately (not deferred to commit like MDBX), - // but this is safe for first_sync because if we crash before commit, the - // checkpoint stays at 0 and we'll just clear and rebuild again on restart. The - // source data (changesets) is intact. + // RocksDB clear executes immediately. A crash before commit leaves the durable + // checkpoint at zero, so the next attempt clears and rebuilds again. provider.rocksdb_provider().clear::()?; } else { provider.tx_ref().clear::()?; } - range = 0..=*input.next_block_range().end(); + if input.checkpoint().block_number == 0 && !prune_target_applied { + block_range = 0..=*block_range.end(); + } } - info!(target: "sync::stages::index_account_history::exec", ?first_sync, ?use_rocksdb, "Collecting indices"); + info!(target: "sync::stages::index_account_history::exec", ?rebuild_from_empty, ?use_rocksdb, "Collecting indices"); - let collector = if provider.cached_storage_settings().storage_v2 { + let collector = if use_rocksdb { // Use the provider-based collection that can read from static files. - collect_account_history_indices(provider, range.clone(), &self.etl_config)? + collect_account_history_indices(provider, block_range.clone(), &self.etl_config)? } else { collect_history_indices::<_, tables::AccountChangeSets, tables::AccountsHistory, _>( provider, - range.clone(), + block_range.clone(), ShardedKey::new, |(index, value)| (index, value.address), &self.etl_config, @@ -135,19 +157,35 @@ where info!(target: "sync::stages::index_account_history::exec", "Loading indices into database"); - provider.with_rocksdb_batch_auto_commit(|rocksdb_batch| { - let mut writer = EitherWriter::new_accounts_history(provider, rocksdb_batch)?; - load_account_history(collector, first_sync, &mut writer) - .map_err(|e| reth_provider::ProviderError::other(Box::new(e)))?; - Ok(((), writer.into_raw_rocksdb_batch())) - })?; + if rebuild_from_empty { + provider.with_rocksdb_batch_auto_commit(|rocksdb_batch| { + let mut writer = EitherWriter::new_accounts_history(provider, rocksdb_batch)?; + load_account_history_append(collector, &mut writer) + .map_err(|e| reth_provider::ProviderError::other(Box::new(e)))?; + Ok(((), writer.into_raw_rocksdb_batch())) + })?; + } else { + let prepared = + prepare_account_history_writes(collector, provider, use_rocksdb, &self.etl_config)?; + provider.with_rocksdb_batch_auto_commit(|rocksdb_batch| { + let mut writer = EitherWriter::new_accounts_history(provider, rocksdb_batch)?; + write_prepared_history_shards::( + prepared, + |key, value| writer.upsert_account_history(key, value), + )?; + Ok(((), writer.into_raw_rocksdb_batch())) + })?; + } if use_rocksdb { provider.commit_pending_rocksdb_batches()?; provider.rocksdb_provider().flush(&[Tables::AccountsHistory.name()])?; } - Ok(ExecOutput { checkpoint: StageCheckpoint::new(*range.end()), done: true }) + Ok(ExecOutput { + checkpoint: StageCheckpoint::new(*block_range.end()), + done: is_final_range, + }) } /// Unwind the stage. @@ -286,6 +324,56 @@ mod tests { assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![0])])); } + #[tokio::test] + async fn initial_target_zero_indexes_genesis() { + let db = TestStageDB::default(); + partial_setup(&db); + + run(&db, 0, None); + + let table = cast(db.table::().unwrap()); + assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![0])])); + } + + #[tokio::test] + async fn zero_commit_threshold_still_makes_progress() { + let db = TestStageDB::default(); + partial_setup(&db); + let input = ExecInput { target: Some(1), checkpoint: None }; + let mut stage = IndexAccountHistoryStage { commit_threshold: 0, ..Default::default() }; + let provider = db.factory.database_provider_rw().unwrap(); + + let output = stage.execute(&provider, input).unwrap(); + + assert_eq!(output, ExecOutput { checkpoint: StageCheckpoint::new(1), done: true }); + provider.commit().unwrap(); + let table = cast(db.table::().unwrap()); + assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![0, 1])])); + } + + #[tokio::test] + async fn execute_respects_commit_threshold() { + let db = TestStageDB::default(); + partial_setup(&db); + let mut stage = IndexAccountHistoryStage { commit_threshold: 2, ..Default::default() }; + + let mut checkpoint = None; + for (expected_checkpoint, done) in [(2, false), (4, false), (5, true)] { + let input = ExecInput { target: Some(5), checkpoint }; + let provider = db.factory.database_provider_rw().unwrap(); + let output = stage.execute(&provider, input).unwrap(); + assert_eq!( + output, + ExecOutput { checkpoint: StageCheckpoint::new(expected_checkpoint), done } + ); + provider.commit().unwrap(); + checkpoint = Some(output.checkpoint); + } + + let table = cast(db.table::().unwrap()); + assert_eq!(table, BTreeMap::from([(shard(u64::MAX), (0..=5).collect())])); + } + #[tokio::test] async fn insert_index_to_not_empty_shard() { // init @@ -314,6 +402,49 @@ mod tests { assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![1, 2, 3])])); } + #[tokio::test] + async fn insert_index_many_distinct_accounts_incremental() { + let db = TestStageDB::default(); + const KEYS: u8 = 32; + + db.commit(|tx| { + for block in 0..=MAX_BLOCK { + tx.put::( + block, + StoredBlockBodyIndices { tx_count: 3, ..Default::default() }, + )?; + } + for i in 1..=KEYS { + let address = Address::repeat_byte(i); + tx.put::( + ShardedKey::new(address, u64::MAX), + list(&[1, 2, 3]), + )?; + for block in 0..=MAX_BLOCK { + tx.put::( + block, + AccountBeforeTx { address, info: None }, + )?; + } + } + Ok(()) + }) + .unwrap(); + + run(&db, 5, Some(3)); + + let table = db.table::().unwrap(); + assert_eq!(table.len(), KEYS as usize); + for i in 1..=KEYS { + let address = Address::repeat_byte(i); + let shard = table + .iter() + .find(|(key, _)| key.key == address && key.highest_block_number == u64::MAX) + .map(|(_, list)| list.iter().collect::>()); + assert_eq!(shard, Some(vec![1, 2, 3, 4, 5]), "address {address}"); + } + } + #[tokio::test] async fn insert_index_to_full_shard() { // init @@ -494,6 +625,7 @@ mod tests { tx.put::(20, acc()).unwrap(); tx.put::(36, acc()).unwrap(); tx.put::(100, acc()).unwrap(); + tx.put::(shard(u64::MAX), list(&[1])).unwrap(); Ok(()) }) .unwrap(); @@ -521,6 +653,34 @@ mod tests { assert!(table.is_empty()); } + #[tokio::test] + async fn prune_target_zero_excludes_genesis() { + let db = TestStageDB::default(); + db.commit(|tx| { + tx.put::(0, acc())?; + tx.put::(1, acc())?; + Ok(()) + }) + .unwrap(); + + let input = ExecInput { target: Some(20_000), ..Default::default() }; + let mut stage = IndexAccountHistoryStage { + prune_mode: Some(PruneMode::Before(1)), + ..Default::default() + }; + let provider = db.factory.database_provider_rw().unwrap(); + let output = stage.execute(&provider, input).unwrap(); + assert_eq!(output, ExecOutput { checkpoint: StageCheckpoint::new(20_000), done: true }); + provider.commit().unwrap(); + + let table = cast(db.table::().unwrap()); + assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![1])])); + let provider = db.factory.provider().unwrap(); + let prune_checkpoint = + provider.get_prune_checkpoint(PruneSegment::AccountHistory).unwrap().unwrap(); + assert_eq!(prune_checkpoint.block_number, Some(0)); + } + stage_test_suite_ext!(IndexAccountHistoryTestRunner, index_account_history); struct IndexAccountHistoryTestRunner { @@ -724,6 +884,38 @@ mod tests { assert_eq!(blocks, (0..=10).collect::>()); } + #[tokio::test] + async fn prune_bump_rebuilds_empty_table_then_merges_later_chunks() { + let db = TestStageDB::default(); + setup_v2_account_data(&db, 0..=10); + let rocksdb = db.factory.rocksdb_provider(); + rocksdb.put::(shard(u64::MAX), &list(&[1])).unwrap(); + + let mut stage = IndexAccountHistoryStage { + commit_threshold: 2, + prune_mode: Some(PruneMode::Before(6)), + ..Default::default() + }; + let mut checkpoint = None; + for (threshold, expected_checkpoint, done) in + [(2, 7, false), (2, 9, false), (u64::MAX, 20_000, true)] + { + stage.commit_threshold = threshold; + let input = ExecInput { target: Some(20_000), checkpoint }; + let provider = db.factory.database_provider_rw().unwrap(); + let output = stage.execute(&provider, input).unwrap(); + assert_eq!( + output, + ExecOutput { checkpoint: StageCheckpoint::new(expected_checkpoint), done } + ); + provider.commit().unwrap(); + checkpoint = Some(output.checkpoint); + } + + let result = rocksdb.get::(shard(u64::MAX)).unwrap().unwrap(); + assert_eq!(result.iter().collect::>(), (6..=10).collect::>()); + } + /// Test that unwind works correctly when `account_history_in_rocksdb` is enabled. #[tokio::test] async fn unwind_works_when_rocksdb_enabled() { diff --git a/crates/stages/stages/src/stages/index_storage_history.rs b/crates/stages/stages/src/stages/index_storage_history.rs index 182bc1c96f5..f5a917de48b 100644 --- a/crates/stages/stages/src/stages/index_storage_history.rs +++ b/crates/stages/stages/src/stages/index_storage_history.rs @@ -1,5 +1,10 @@ use super::{collect_history_indices, collect_storage_history_indices}; -use crate::{stages::utils::load_storage_history, StageCheckpoint, StageId}; +use crate::{ + stages::utils::{ + load_storage_history_append, prepare_storage_history_writes, write_prepared_history_shards, + }, + StageCheckpoint, StageId, +}; use reth_config::config::{EtlConfig, IndexHistoryConfig}; use reth_db_api::{ models::{storage_sharded_key::StorageShardedKey, AddressStorageKey, BlockNumberAddress}, @@ -13,7 +18,9 @@ use reth_provider::{ StorageSettingsCache, }; use reth_prune_types::{PruneCheckpoint, PruneMode, PrunePurpose, PruneSegment}; -use reth_stages_api::{ExecInput, ExecOutput, Stage, StageError, UnwindInput, UnwindOutput}; +use reth_stages_api::{ + BlockRangeOutput, ExecInput, ExecOutput, Stage, StageError, UnwindInput, UnwindOutput, +}; use std::fmt::Debug; use tracing::info; @@ -71,6 +78,10 @@ where provider: &Provider, mut input: ExecInput, ) -> Result { + let initial_sync = input.checkpoint.is_none(); + let rebuild_from_empty = input.checkpoint().block_number == 0; + let mut prune_target_applied = false; + if let Some((target_prunable_block, prune_mode)) = self .prune_mode .map(|mode| { @@ -82,8 +93,10 @@ where }) .transpose()? .flatten() && - target_prunable_block > input.checkpoint().block_number + (target_prunable_block > input.checkpoint().block_number || + target_prunable_block == 0 && input.checkpoint().block_number == 0) { + prune_target_applied = true; input.checkpoint = Some(StageCheckpoint::new(target_prunable_block)); // Save prune checkpoint only if we don't have one already. @@ -100,36 +113,48 @@ where } } - if input.target_reached() { + let use_rocksdb = provider.cached_storage_settings().storage_v2; + if input.target_reached() && prune_target_applied { + if use_rocksdb { + provider.rocksdb_provider().clear::()?; + } else { + provider.tx_ref().clear::()?; + } + return Ok(ExecOutput::done(input.checkpoint())) + } + if input.target_reached() && !initial_sync { return Ok(ExecOutput::done(input.checkpoint())) } - let mut range = input.next_block_range(); - let first_sync = input.checkpoint().block_number == 0; - let use_rocksdb = provider.cached_storage_settings().storage_v2; + let BlockRangeOutput { mut block_range, is_final_range } = + if initial_sync && input.target() == 0 { + BlockRangeOutput { block_range: 0..=0, is_final_range: true } + } else { + input.next_block_range_with_threshold(self.commit_threshold.max(1)) + }; - // On first sync we might have history coming from genesis. We clear the table since it's - // faster to rebuild from scratch. - if first_sync { + // A zero checkpoint means no history range is durable. Clear stale rows before using the + // append-only loader, including when pruning advanced the in-memory checkpoint. + if rebuild_from_empty { if use_rocksdb { - // Note: RocksDB clear() executes immediately (not deferred to commit like MDBX), - // but this is safe for first_sync because if we crash before commit, the - // checkpoint stays at 0 and we'll just clear and rebuild again on restart. The - // source data (changesets) is intact. + // RocksDB clear executes immediately. A crash before commit leaves the durable + // checkpoint at zero, so the next attempt clears and rebuilds again. provider.rocksdb_provider().clear::()?; } else { provider.tx_ref().clear::()?; } - range = 0..=*input.next_block_range().end(); + if input.checkpoint().block_number == 0 && !prune_target_applied { + block_range = 0..=*block_range.end(); + } } - info!(target: "sync::stages::index_storage_history::exec", ?first_sync, ?use_rocksdb, "Collecting indices"); - let collector = if provider.cached_storage_settings().storage_v2 { - collect_storage_history_indices(provider, range.clone(), &self.etl_config)? + info!(target: "sync::stages::index_storage_history::exec", ?rebuild_from_empty, ?use_rocksdb, "Collecting indices"); + let collector = if use_rocksdb { + collect_storage_history_indices(provider, block_range.clone(), &self.etl_config)? } else { collect_history_indices::<_, tables::StorageChangeSets, tables::StoragesHistory, _>( provider, - BlockNumberAddress::range(range.clone()), + BlockNumberAddress::range(block_range.clone()), |AddressStorageKey((address, storage_key)), highest_block_number| { StorageShardedKey::new(address, storage_key, highest_block_number) }, @@ -140,19 +165,35 @@ where info!(target: "sync::stages::index_storage_history::exec", "Loading indices into database"); - provider.with_rocksdb_batch_auto_commit(|rocksdb_batch| { - let mut writer = EitherWriter::new_storages_history(provider, rocksdb_batch)?; - load_storage_history(collector, first_sync, &mut writer) - .map_err(|e| reth_provider::ProviderError::other(Box::new(e)))?; - Ok(((), writer.into_raw_rocksdb_batch())) - })?; + if rebuild_from_empty { + provider.with_rocksdb_batch_auto_commit(|rocksdb_batch| { + let mut writer = EitherWriter::new_storages_history(provider, rocksdb_batch)?; + load_storage_history_append(collector, &mut writer) + .map_err(|e| reth_provider::ProviderError::other(Box::new(e)))?; + Ok(((), writer.into_raw_rocksdb_batch())) + })?; + } else { + let prepared = + prepare_storage_history_writes(collector, provider, use_rocksdb, &self.etl_config)?; + provider.with_rocksdb_batch_auto_commit(|rocksdb_batch| { + let mut writer = EitherWriter::new_storages_history(provider, rocksdb_batch)?; + write_prepared_history_shards::( + prepared, + |key, value| writer.upsert_storage_history(key, value), + )?; + Ok(((), writer.into_raw_rocksdb_batch())) + })?; + } if use_rocksdb { provider.commit_pending_rocksdb_batches()?; provider.rocksdb_provider().flush(&[Tables::StoragesHistory.name()])?; } - Ok(ExecOutput { checkpoint: StageCheckpoint::new(*range.end()), done: true }) + Ok(ExecOutput { + checkpoint: StageCheckpoint::new(*block_range.end()), + done: is_final_range, + }) } /// Unwind the stage. @@ -304,6 +345,56 @@ mod tests { assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![0])])); } + #[tokio::test] + async fn initial_target_zero_indexes_genesis() { + let db = TestStageDB::default(); + partial_setup(&db); + + run(&db, 0, None); + + let table = cast(db.table::().unwrap()); + assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![0])])); + } + + #[tokio::test] + async fn zero_commit_threshold_still_makes_progress() { + let db = TestStageDB::default(); + partial_setup(&db); + let input = ExecInput { target: Some(1), checkpoint: None }; + let mut stage = IndexStorageHistoryStage { commit_threshold: 0, ..Default::default() }; + let provider = db.factory.database_provider_rw().unwrap(); + + let output = stage.execute(&provider, input).unwrap(); + + assert_eq!(output, ExecOutput { checkpoint: StageCheckpoint::new(1), done: true }); + provider.commit().unwrap(); + let table = cast(db.table::().unwrap()); + assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![0, 1])])); + } + + #[tokio::test] + async fn execute_respects_commit_threshold() { + let db = TestStageDB::default(); + partial_setup(&db); + let mut stage = IndexStorageHistoryStage { commit_threshold: 2, ..Default::default() }; + + let mut checkpoint = None; + for (expected_checkpoint, done) in [(2, false), (4, false), (5, true)] { + let input = ExecInput { target: Some(5), checkpoint }; + let provider = db.factory.database_provider_rw().unwrap(); + let output = stage.execute(&provider, input).unwrap(); + assert_eq!( + output, + ExecOutput { checkpoint: StageCheckpoint::new(expected_checkpoint), done } + ); + provider.commit().unwrap(); + checkpoint = Some(output.checkpoint); + } + + let table = cast(db.table::().unwrap()); + assert_eq!(table, BTreeMap::from([(shard(u64::MAX), (0..=5).collect())])); + } + #[tokio::test] async fn insert_index_to_not_empty_shard() { // init @@ -332,6 +423,55 @@ mod tests { assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![1, 2, 3])])); } + #[tokio::test] + async fn insert_index_many_distinct_slots_incremental() { + let db = TestStageDB::default(); + const KEYS: u8 = 32; + + db.commit(|tx| { + for block in 0..=MAX_BLOCK { + tx.put::( + block, + StoredBlockBodyIndices { tx_count: 3, ..Default::default() }, + )?; + } + for i in 1..=KEYS { + let address = Address::repeat_byte(i); + let slot = B256::repeat_byte(i); + tx.put::( + StorageShardedKey::new(address, slot, u64::MAX), + list(&[1, 2, 3]), + )?; + for block in 0..=MAX_BLOCK { + tx.put::( + BlockNumberAddress((block, address)), + storage(slot), + )?; + } + } + Ok(()) + }) + .unwrap(); + + run(&db, 5, Some(3)); + + let table = db.table::().unwrap(); + assert_eq!(table.len(), KEYS as usize); + for i in 1..=KEYS { + let address = Address::repeat_byte(i); + let slot = B256::repeat_byte(i); + let shard = table + .iter() + .find(|(key, _)| { + key.address == address && + key.sharded_key.key == slot && + key.sharded_key.highest_block_number == u64::MAX + }) + .map(|(_, list)| list.iter().collect::>()); + assert_eq!(shard, Some(vec![1, 2, 3, 4, 5]), "slot {address}/{slot}"); + } + } + #[tokio::test] async fn insert_index_to_full_shard() { // init @@ -515,6 +655,7 @@ mod tests { .unwrap(); tx.put::(block_number_address(100), storage(STORAGE_KEY)) .unwrap(); + tx.put::(shard(u64::MAX), list(&[1])).unwrap(); Ok(()) }) .unwrap(); @@ -542,6 +683,34 @@ mod tests { assert!(table.is_empty()); } + #[tokio::test] + async fn prune_target_zero_excludes_genesis() { + let db = TestStageDB::default(); + db.commit(|tx| { + tx.put::(block_number_address(0), storage(STORAGE_KEY))?; + tx.put::(block_number_address(1), storage(STORAGE_KEY))?; + Ok(()) + }) + .unwrap(); + + let input = ExecInput { target: Some(20_000), ..Default::default() }; + let mut stage = IndexStorageHistoryStage { + prune_mode: Some(PruneMode::Before(1)), + ..Default::default() + }; + let provider = db.factory.database_provider_rw().unwrap(); + let output = stage.execute(&provider, input).unwrap(); + assert_eq!(output, ExecOutput { checkpoint: StageCheckpoint::new(20_000), done: true }); + provider.commit().unwrap(); + + let table = cast(db.table::().unwrap()); + assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![1])])); + let provider = db.factory.provider().unwrap(); + let prune_checkpoint = + provider.get_prune_checkpoint(PruneSegment::StorageHistory).unwrap().unwrap(); + assert_eq!(prune_checkpoint.block_number, Some(0)); + } + stage_test_suite_ext!(IndexStorageHistoryTestRunner, index_storage_history); struct IndexStorageHistoryTestRunner { @@ -759,6 +928,38 @@ mod tests { assert_eq!(blocks, (0..=10).collect::>()); } + #[tokio::test] + async fn prune_bump_rebuilds_empty_table_then_merges_later_chunks() { + let db = TestStageDB::default(); + setup_v2_storage_data(&db, 0..=10); + let rocksdb = db.factory.rocksdb_provider(); + rocksdb.put::(shard(u64::MAX), &list(&[1])).unwrap(); + + let mut stage = IndexStorageHistoryStage { + commit_threshold: 2, + prune_mode: Some(PruneMode::Before(6)), + ..Default::default() + }; + let mut checkpoint = None; + for (threshold, expected_checkpoint, done) in + [(2, 7, false), (2, 9, false), (u64::MAX, 20_000, true)] + { + stage.commit_threshold = threshold; + let input = ExecInput { target: Some(20_000), checkpoint }; + let provider = db.factory.database_provider_rw().unwrap(); + let output = stage.execute(&provider, input).unwrap(); + assert_eq!( + output, + ExecOutput { checkpoint: StageCheckpoint::new(expected_checkpoint), done } + ); + provider.commit().unwrap(); + checkpoint = Some(output.checkpoint); + } + + let result = rocksdb.get::(shard(u64::MAX)).unwrap().unwrap(); + assert_eq!(result.iter().collect::>(), (6..=10).collect::>()); + } + /// Test that unwind works correctly when `storages_history_in_rocksdb` is enabled. #[tokio::test] async fn unwind_works_when_rocksdb_enabled() { diff --git a/crates/stages/stages/src/stages/utils.rs b/crates/stages/stages/src/stages/utils.rs index da8d3f1fe1b..d8b3f7870bc 100644 --- a/crates/stages/stages/src/stages/utils.rs +++ b/crates/stages/stages/src/stages/utils.rs @@ -11,23 +11,51 @@ use reth_db_api::{ AccountBeforeTx, AddressStorageKey, BlockNumberAddress, ShardedKey, }, table::{Decode, Decompress, Table}, + tables, transaction::DbTx, BlockNumberList, }; use reth_etl::Collector; use reth_primitives_traits::NodePrimitives; use reth_provider::{ - providers::StaticFileProvider, to_range, BlockReader, DBProvider, EitherWriter, ProviderError, - StaticFileProviderFactory, + prepare_history_shard_writes_parallel_vec, prepare_history_shard_writes_serial_vec, + providers::StaticFileProvider, to_range, BlockReader, DBProvider, EitherWriter, + PreparedHistoryShardWrites, ProviderError, ProviderResult, RocksDBProviderFactory, + ShardedHistoryTable, StaticFileProviderFactory, }; use reth_stages_api::StageError; use reth_static_file_types::StaticFileSegment; use reth_storage_api::{ChangeSetReader, StorageChangeSetReader}; -use std::{hash::Hash, ops::RangeBounds}; +use std::{hash::Hash, mem, ops::RangeBounds}; use tracing::info; /// Number of blocks before pushing indices from cache to [`Collector`] const DEFAULT_CACHE_THRESHOLD: u64 = 100_000; +/// Maximum number of logical keys held in a history collection cache. +const HISTORY_CACHE_KEY_LIMIT: usize = 500_000; +/// Maximum number of block numbers held in a history collection cache. +const HISTORY_CACHE_INDEX_LIMIT: usize = 8_000_000; +/// Maximum number of complete logical keys prepared in one batch. +const HISTORY_PREPARATION_KEY_LIMIT: usize = 65_536; +/// Maximum estimated decoded input size prepared in one batch. +const HISTORY_PREPARATION_BYTE_LIMIT: usize = 256 * 1024 * 1024; + +#[derive(Clone, Copy, Debug)] +struct HistoryPreparationLimits { + max_keys: usize, + max_bytes: usize, +} + +const HISTORY_PREPARATION_LIMITS: HistoryPreparationLimits = HistoryPreparationLimits { + max_keys: HISTORY_PREPARATION_KEY_LIMIT, + max_bytes: HISTORY_PREPARATION_BYTE_LIMIT, +}; + +const fn history_cache_limit_reached(keys: usize, indices: usize, blocks: u64) -> bool { + keys >= HISTORY_CACHE_KEY_LIMIT || + indices >= HISTORY_CACHE_INDEX_LIMIT || + blocks >= DEFAULT_CACHE_THRESHOLD +} /// Collects all history (`H`) indices for a range of changesets (`CS`) and stores them in a /// [`Collector`]. @@ -77,25 +105,31 @@ where let total_changesets = provider.tx_ref().entries::()?; let interval = (total_changesets / 1000).max(1); - let mut flush_counter = 0; - let mut current_block_number = u64::MAX; + let mut cached_blocks = 0; + let mut cached_indices = 0; + let mut current_block_number = None; for (idx, entry) in changeset_cursor.walk_range(range)?.enumerate() { let (block_number, key) = partial_key_factory(entry?); - cache.entry(key).or_default().push(block_number); if idx > 0 && idx.is_multiple_of(interval) && total_changesets > 1000 { info!(target: "sync::stages::index_history", progress = %format_args!("{:.4}%", (idx as f64 / total_changesets as f64) * 100.0), "Collecting indices"); } - // Make sure we only flush the cache every DEFAULT_CACHE_THRESHOLD blocks. - if current_block_number != block_number { - current_block_number = block_number; - flush_counter += 1; - if flush_counter > DEFAULT_CACHE_THRESHOLD { + // Check limits before the first row of a new block so a flush never splits one block. + if current_block_number != Some(block_number) { + if current_block_number.is_some() && + history_cache_limit_reached(cache.len(), cached_indices, cached_blocks) + { collect(&mut cache)?; - flush_counter = 0; + cached_blocks = 0; + cached_indices = 0; } + current_block_number = Some(block_number); + cached_blocks += 1; } + + cache.entry(key).or_default().push(block_number); + cached_indices += 1; } collect(&mut cache)?; @@ -144,28 +178,34 @@ where let walker = static_file_provider.walk_account_changeset_range(range); - let mut flush_counter = 0; - let mut current_block_number = u64::MAX; + let mut cached_blocks = 0; + let mut cached_indices = 0; + let mut current_block_number = None; for changeset_result in walker { let (block_number, AccountBeforeTx { address, .. }) = changeset_result?; - cache.entry(address).or_default().push(block_number); - if block_number != current_block_number { - current_block_number = block_number; - flush_counter += 1; + // Check limits before the first row of a new block so a flush never splits one block. + if current_block_number != Some(block_number) { + if let Some(completed_block) = current_block_number && + history_cache_limit_reached(cache.len(), cached_indices, cached_blocks) + { + info!( + target: "sync::stages::index_history", + processed_blocks = completed_block.saturating_sub(start_block) + 1, + current_block = completed_block, + "Collecting indices" + ); + collect_indices(cache.drain(), &mut insert_fn)?; + cached_blocks = 0; + cached_indices = 0; + } + current_block_number = Some(block_number); + cached_blocks += 1; } - if flush_counter > DEFAULT_CACHE_THRESHOLD { - info!( - target: "sync::stages::index_history", - processed_blocks = current_block_number.saturating_sub(start_block) + 1, - current_block = current_block_number, - "Collecting indices" - ); - collect_indices(cache.drain(), &mut insert_fn)?; - flush_counter = 0; - } + cache.entry(address).or_default().push(block_number); + cached_indices += 1; } collect_indices(cache.into_iter(), insert_fn)?; @@ -199,28 +239,34 @@ where let walker = static_file_provider.walk_storage_changeset_range(range); - let mut flush_counter = 0; - let mut current_block_number = u64::MAX; + let mut cached_blocks = 0; + let mut cached_indices = 0; + let mut current_block_number = None; for changeset_result in walker { let (BlockNumberAddress((block_number, address)), storage) = changeset_result?; - cache.entry(AddressStorageKey((address, storage.key))).or_default().push(block_number); - if block_number != current_block_number { - current_block_number = block_number; - flush_counter += 1; + // Check limits before the first row of a new block so a flush never splits one block. + if current_block_number != Some(block_number) { + if let Some(completed_block) = current_block_number && + history_cache_limit_reached(cache.len(), cached_indices, cached_blocks) + { + info!( + target: "sync::stages::index_history", + processed_blocks = completed_block.saturating_sub(start_block) + 1, + current_block = completed_block, + "Collecting indices" + ); + collect_indices(cache.drain(), &mut insert_fn)?; + cached_blocks = 0; + cached_indices = 0; + } + current_block_number = Some(block_number); + cached_blocks += 1; } - if flush_counter > DEFAULT_CACHE_THRESHOLD { - info!( - target: "sync::stages::index_history", - processed_blocks = current_block_number.saturating_sub(start_block) + 1, - current_block = current_block_number, - "Collecting indices" - ); - collect_indices(cache.drain(), &mut insert_fn)?; - flush_counter = 0; - } + cache.entry(AddressStorageKey((address, storage.key))).or_default().push(block_number); + cached_indices += 1; } collect_indices(cache.into_iter(), insert_fn)?; @@ -228,26 +274,242 @@ where Ok(collector) } -/// Loads account history indices into the database via `EitherWriter`. +fn emit_grouped_history_key( + key: T::PartialKey, + indices: Vec, + grouped: &mut Vec<(T::PartialKey, Vec)>, + grouped_bytes: &mut usize, + limits: HistoryPreparationLimits, + emit: &mut F, +) -> Result<(), StageError> +where + T: ShardedHistoryTable, + F: FnMut(Vec<(T::PartialKey, Vec)>) -> Result<(), StageError>, +{ + debug_assert!( + indices.windows(2).all(|window| window[0] < window[1]), + "indices must be strictly increasing: {indices:?}" + ); + + let group_bytes = history_group_bytes::(&indices); + let next_key_exceeds_limit = grouped.len() >= limits.max_keys; + let next_key_exceeds_bytes = group_bytes > limits.max_bytes.saturating_sub(*grouped_bytes); + + if !grouped.is_empty() && (next_key_exceeds_limit || next_key_exceeds_bytes) { + emit(mem::take(grouped))?; + *grouped_bytes = 0; + } + + *grouped_bytes = (*grouped_bytes).saturating_add(group_bytes); + grouped.push((key, indices)); + Ok(()) +} + +const fn history_group_bytes(indices: &Vec) -> usize { + mem::size_of::<(T::PartialKey, Vec)>() + .saturating_add(indices.capacity().saturating_mul(mem::size_of::())) +} + +/// Streams sorted ETL rows into bounded batches of complete logical keys. +fn for_each_grouped_history_chunk( + mut collector: Collector, + limits: HistoryPreparationLimits, + mut emit: F, +) -> Result<(), StageError> +where + T: ShardedHistoryTable, + F: FnMut(Vec<(T::PartialKey, Vec)>) -> Result<(), StageError>, +{ + let mut grouped = Vec::new(); + let mut grouped_bytes = 0; + let mut current_key = None; + let mut current_indices = Vec::new(); + let total_entries = collector.len(); + let interval = (total_entries / 10).max(1); + + for (index, element) in collector.iter()?.enumerate() { + let (k, v) = element?; + let key = T::Key::decode_owned(k)?; + + if index > 0 && index.is_multiple_of(interval) && total_entries > 10 { + info!(target: "sync::stages::index_history", progress = %format!("{:.2}%", (index as f64 / total_entries as f64) * 100.0), "Grouping indices"); + } + + let partial_key = T::partial_key(&key); + if current_key != Some(partial_key) && + let Some(previous_key) = current_key.replace(partial_key) + { + emit_grouped_history_key::( + previous_key, + mem::take(&mut current_indices), + &mut grouped, + &mut grouped_bytes, + limits, + &mut emit, + )?; + } + + if !grouped.is_empty() && + (grouped.len() >= limits.max_keys || grouped_bytes >= limits.max_bytes) + { + emit(mem::take(&mut grouped))?; + grouped_bytes = 0; + } + + let new_list = BlockNumberList::decompress_owned(v)?; + let projected_len = + current_indices.len().saturating_add(new_list.len().try_into().unwrap_or(usize::MAX)); + let projected_capacity = current_indices.capacity().max(projected_len); + let projected_bytes = mem::size_of::<(T::PartialKey, Vec)>() + .saturating_add(projected_capacity.saturating_mul(mem::size_of::())); + if !grouped.is_empty() && projected_bytes > limits.max_bytes.saturating_sub(grouped_bytes) { + emit(mem::take(&mut grouped))?; + grouped_bytes = 0; + } + current_indices.extend(new_list.iter()); + + if !grouped.is_empty() && + history_group_bytes::(¤t_indices) > + limits.max_bytes.saturating_sub(grouped_bytes) + { + emit(mem::take(&mut grouped))?; + grouped_bytes = 0; + } + } + + if let Some(key) = current_key { + emit_grouped_history_key::( + key, + current_indices, + &mut grouped, + &mut grouped_bytes, + limits, + &mut emit, + )?; + } + if !grouped.is_empty() { + emit(grouped)?; + } + + Ok(()) +} + +fn prepare_grouped_history_writes( + grouped: Vec<(T::PartialKey, Vec)>, + provider: &Provider, + use_rocksdb: bool, +) -> Result, StageError> +where + T: ShardedHistoryTable, + Provider: DBProvider + RocksDBProviderFactory, +{ + if use_rocksdb { + let rocksdb = provider.rocksdb_provider(); + Ok(prepare_history_shard_writes_parallel_vec::(grouped, |key| rocksdb.get::(key))?) + } else { + Ok(prepare_history_shard_writes_serial_vec::(grouped, |key| { + provider.tx_ref().get::(key).map_err(Into::into) + })?) + } +} + +fn prepare_history_writes( + collector: Collector, + provider: &Provider, + use_rocksdb: bool, + etl_config: &EtlConfig, +) -> Result, StageError> +where + T: ShardedHistoryTable, + Provider: DBProvider + RocksDBProviderFactory, +{ + let mut prepared = Collector::new(etl_config.file_size, etl_config.dir.clone()); + for_each_grouped_history_chunk::(collector, HISTORY_PREPARATION_LIMITS, |grouped| { + let writes = prepare_grouped_history_writes::(grouped, provider, use_rocksdb)?; + for (key, value) in writes.into_writes() { + prepared.insert(key, value)?; + } + Ok(()) + })?; + Ok(prepared) +} + +/// Spools prepared account-history shards after merging each key's committed last shard. /// -/// Works with [`EitherWriter`] to support both MDBX and `RocksDB` backends. +/// Call this **before** opening a `RocksDB` write batch so last-shard gets observe committed state +/// only and do not overlap auto-commit `write_opt`. +pub(crate) fn prepare_account_history_writes( + collector: Collector, BlockNumberList>, + provider: &Provider, + use_rocksdb: bool, + etl_config: &EtlConfig, +) -> Result, BlockNumberList>, StageError> +where + Provider: DBProvider + RocksDBProviderFactory, +{ + prepare_history_writes::( + collector, + provider, + use_rocksdb, + etl_config, + ) +} + +/// Spools prepared storage-history shards after merging each key's committed last shard. /// -/// ## Process -/// Iterates over elements, grouping indices by their address. It flushes indices to disk -/// when reaching a shard's max length (`NUM_OF_INDICES_IN_SHARD`) or when the address changes, -/// ensuring the last previous address shard is stored. +/// Call this **before** opening a `RocksDB` write batch so last-shard gets observe committed state +/// only and do not overlap auto-commit `write_opt`. +pub(crate) fn prepare_storage_history_writes( + collector: Collector, + provider: &Provider, + use_rocksdb: bool, + etl_config: &EtlConfig, +) -> Result, StageError> +where + Provider: DBProvider + RocksDBProviderFactory, +{ + prepare_history_writes::( + collector, + provider, + use_rocksdb, + etl_config, + ) +} + +/// Streams prepared history shards into serial puts, logging progress every 10%. +pub(crate) fn write_prepared_history_shards( + mut prepared: Collector, + mut write: impl FnMut(T::Key, &BlockNumberList) -> ProviderResult<()>, +) -> ProviderResult<()> +where + T: ShardedHistoryTable, +{ + let total_writes = prepared.len(); + let interval = (total_writes / 10).max(1); + + for (index, element) in prepared.iter().map_err(ProviderError::other)?.enumerate() { + if index > 0 && index.is_multiple_of(interval) && total_writes > 10 { + info!(target: "sync::stages::index_history", progress = %format!("{:.2}%", (index as f64 / total_writes as f64) * 100.0), "Writing indices"); + } + let (key, value) = element.map_err(ProviderError::other)?; + let key = T::Key::decode_owned(key)?; + let value = BlockNumberList::decompress_owned(value)?; + write(key, &value)?; + } + + Ok(()) +} + +/// Append-only empty-table loader for account history. /// -/// Uses `Option
` instead of `Address::default()` as the sentinel to avoid -/// incorrectly treating `Address::ZERO` as "no previous address". -pub(crate) fn load_account_history( +/// Streams the collector into `append_*` and never reads last shards. +pub(crate) fn load_account_history_append( mut collector: Collector, BlockNumberList>, - append_only: bool, writer: &mut EitherWriter<'_, CURSOR, N>, ) -> Result<(), StageError> where N: NodePrimitives, - CURSOR: DbCursorRW - + DbCursorRO, + CURSOR: DbCursorRW + DbCursorRO, { let mut current_address: Option
= None; // Accumulator for block numbers where the current address changed. @@ -271,31 +533,23 @@ where if current_address != Some(address) { // Flush all remaining shards for the previous address (uses u64::MAX for last shard). if let Some(prev_addr) = current_address { - flush_account_history_shards(prev_addr, &mut current_list, append_only, writer)?; + flush_account_history_shards(prev_addr, &mut current_list, writer)?; } current_address = Some(address); current_list.clear(); - - // On incremental sync, merge with the existing last shard from the database. - // The last shard is stored with key (address, u64::MAX) so we can find it. - if !append_only && - let Some(last_shard) = writer.get_last_account_history_shard(address)? - { - current_list.extend(last_shard.iter()); - } } // Append new block numbers to the accumulator. current_list.extend(new_list.iter()); // Flush complete shards, keeping the last (partial) shard buffered. - flush_account_history_shards_partial(address, &mut current_list, append_only, writer)?; + flush_account_history_shards_partial(address, &mut current_list, writer)?; } // Flush the final address's remaining shard. if let Some(addr) = current_address { - flush_account_history_shards(addr, &mut current_list, append_only, writer)?; + flush_account_history_shards(addr, &mut current_list, writer)?; } Ok(()) @@ -309,13 +563,11 @@ where fn flush_account_history_shards_partial( address: Address, list: &mut Vec, - append_only: bool, writer: &mut EitherWriter<'_, CURSOR, N>, ) -> Result<(), StageError> where N: NodePrimitives, - CURSOR: DbCursorRW - + DbCursorRO, + CURSOR: DbCursorRW + DbCursorRO, { // Nothing to flush if we haven't filled a complete shard yet. if list.len() <= NUM_OF_INDICES_IN_SHARD { @@ -345,12 +597,7 @@ where let highest = *chunk.last().expect("chunk is non-empty"); let key = ShardedKey::new(address, highest); let value = BlockNumberList::new_pre_sorted(chunk.iter().copied()); - - if append_only { - writer.append_account_history(key, &value)?; - } else { - writer.upsert_account_history(key, &value)?; - } + writer.append_account_history(key, &value)?; } // Keep the remaining indices for the next iteration. @@ -365,13 +612,11 @@ where fn flush_account_history_shards( address: Address, list: &mut Vec, - append_only: bool, writer: &mut EitherWriter<'_, CURSOR, N>, ) -> Result<(), StageError> where N: NodePrimitives, - CURSOR: DbCursorRW - + DbCursorRO, + CURSOR: DbCursorRW + DbCursorRO, { if list.is_empty() { return Ok(()); @@ -388,12 +633,7 @@ where let key = ShardedKey::new(address, highest); let value = BlockNumberList::new_pre_sorted(chunk.iter().copied()); - - if append_only { - writer.append_account_history(key, &value)?; - } else { - writer.upsert_account_history(key, &value)?; - } + writer.append_account_history(key, &value)?; } list.clear(); @@ -436,26 +676,16 @@ where }) } -/// Loads storage history indices into the database via `EitherWriter`. -/// -/// Works with [`EitherWriter`] to support both MDBX and `RocksDB` backends. -/// -/// ## Process -/// Iterates over elements, grouping indices by their (address, `storage_key`) pairs. It flushes -/// indices to disk when reaching a shard's max length (`NUM_OF_INDICES_IN_SHARD`) or when the -/// (address, `storage_key`) pair changes, ensuring the last previous shard is stored. +/// Append-only empty-table loader for storage history. /// -/// Uses `Option<(Address, B256)>` instead of default values as the sentinel to avoid -/// incorrectly treating `(Address::ZERO, B256::ZERO)` as "no previous key". -pub(crate) fn load_storage_history( +/// Streams the collector into `append_*` and never reads last shards. +pub(crate) fn load_storage_history_append( mut collector: Collector, - append_only: bool, writer: &mut EitherWriter<'_, CURSOR, N>, ) -> Result<(), StageError> where N: NodePrimitives, - CURSOR: DbCursorRW - + DbCursorRO, + CURSOR: DbCursorRW + DbCursorRO, { let mut current_key: Option<(Address, B256)> = None; // Accumulator for block numbers where the current (address, storage_key) changed. @@ -483,22 +713,12 @@ where prev_addr, prev_storage_key, &mut current_list, - append_only, writer, )?; } current_key = Some(partial_key); current_list.clear(); - - // On incremental sync, merge with the existing last shard from the database. - // The last shard is stored with key (address, storage_key, u64::MAX) so we can find it. - if !append_only && - let Some(last_shard) = - writer.get_last_storage_history_shard(partial_key.0, partial_key.1)? - { - current_list.extend(last_shard.iter()); - } } // Append new block numbers to the accumulator. @@ -509,14 +729,13 @@ where partial_key.0, partial_key.1, &mut current_list, - append_only, writer, )?; } // Flush the final key's remaining shard. if let Some((addr, storage_key)) = current_key { - flush_storage_history_shards(addr, storage_key, &mut current_list, append_only, writer)?; + flush_storage_history_shards(addr, storage_key, &mut current_list, writer)?; } Ok(()) @@ -531,13 +750,11 @@ fn flush_storage_history_shards_partial( address: Address, storage_key: B256, list: &mut Vec, - append_only: bool, writer: &mut EitherWriter<'_, CURSOR, N>, ) -> Result<(), StageError> where N: NodePrimitives, - CURSOR: DbCursorRW - + DbCursorRO, + CURSOR: DbCursorRW + DbCursorRO, { // Nothing to flush if we haven't filled a complete shard yet. if list.len() <= NUM_OF_INDICES_IN_SHARD { @@ -567,12 +784,7 @@ where let highest = *chunk.last().expect("chunk is non-empty"); let key = StorageShardedKey::new(address, storage_key, highest); let value = BlockNumberList::new_pre_sorted(chunk.iter().copied()); - - if append_only { - writer.append_storage_history(key, &value)?; - } else { - writer.upsert_storage_history(key, &value)?; - } + writer.append_storage_history(key, &value)?; } // Keep the remaining indices for the next iteration. @@ -589,13 +801,11 @@ fn flush_storage_history_shards( address: Address, storage_key: B256, list: &mut Vec, - append_only: bool, writer: &mut EitherWriter<'_, CURSOR, N>, ) -> Result<(), StageError> where N: NodePrimitives, - CURSOR: DbCursorRW - + DbCursorRO, + CURSOR: DbCursorRW + DbCursorRO, { if list.is_empty() { return Ok(()); @@ -612,14 +822,145 @@ where let key = StorageShardedKey::new(address, storage_key, highest); let value = BlockNumberList::new_pre_sorted(chunk.iter().copied()); - - if append_only { - writer.append_storage_history(key, &value)?; - } else { - writer.upsert_storage_history(key, &value)?; - } + writer.append_storage_history(key, &value)?; } list.clear(); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use alloy_primitives::{address, b256, BlockNumber}; + + fn list(indices: &[u64]) -> BlockNumberList { + BlockNumberList::new(indices.iter().copied()).unwrap() + } + + fn grouped_account_chunks( + collector: Collector, BlockNumberList>, + limits: HistoryPreparationLimits, + ) -> Vec)>> { + let mut chunks = Vec::new(); + for_each_grouped_history_chunk::(collector, limits, |chunk| { + chunks.push(chunk); + Ok(()) + }) + .unwrap(); + chunks + } + + #[test] + fn grouping_coalesces_one_logical_key_across_etl_files() { + let address = address!("0x0000000000000000000000000000000000000001"); + // Every entry exceeds this tiny buffer, forcing it into a separate ETL file. + let mut collector = Collector::new(1, None); + collector.insert(ShardedKey::new(address, 2), list(&[1, 2])).unwrap(); + collector.insert(ShardedKey::new(address, 4), list(&[3, 4])).unwrap(); + + let chunks = grouped_account_chunks( + collector, + HistoryPreparationLimits { max_keys: 10, max_bytes: usize::MAX }, + ); + assert_eq!(chunks, vec![vec![(address, vec![1, 2, 3, 4])]]); + } + + #[test] + fn storage_grouping_coalesces_one_logical_key_across_etl_files() { + let address = address!("0x0000000000000000000000000000000000000001"); + let slot = b256!("0x0000000000000000000000000000000000000000000000000000000000000002"); + let mut collector = Collector::new(1, None); + collector.insert(StorageShardedKey::new(address, slot, 2), list(&[1, 2])).unwrap(); + collector.insert(StorageShardedKey::new(address, slot, 4), list(&[3, 4])).unwrap(); + + let mut chunks = Vec::new(); + for_each_grouped_history_chunk::( + collector, + HistoryPreparationLimits { max_keys: 10, max_bytes: usize::MAX }, + |chunk| { + chunks.push(chunk); + Ok(()) + }, + ) + .unwrap(); + assert_eq!(chunks, vec![vec![((address, slot), vec![1, 2, 3, 4])]]); + } + + #[test] + fn grouping_never_splits_a_key_at_key_limit() { + let mut collector = Collector::new(1, None); + for byte in 1u8..=5 { + let address = Address::repeat_byte(byte); + collector.insert(ShardedKey::new(address, u64::MAX), list(&[u64::from(byte)])).unwrap(); + } + + let chunks = grouped_account_chunks( + collector, + HistoryPreparationLimits { max_keys: 2, max_bytes: usize::MAX }, + ); + assert_eq!(chunks.iter().map(Vec::len).collect::>(), vec![2, 2, 1]); + assert!(chunks.into_iter().flatten().all(|(_, indices)| indices.len() == 1)); + } + + #[test] + fn grouping_processes_an_over_budget_key_alone() { + let mut collector = Collector::new(1, None); + for byte in 1u8..=2 { + let address = Address::repeat_byte(byte); + collector.insert(ShardedKey::new(address, u64::MAX), list(&[u64::from(byte)])).unwrap(); + } + + let chunks = grouped_account_chunks( + collector, + HistoryPreparationLimits { max_keys: usize::MAX, max_bytes: 1 }, + ); + assert_eq!(chunks.iter().map(Vec::len).collect::>(), vec![1, 1]); + } + + #[test] + fn grouping_keeps_a_coalesced_key_together_when_byte_limit_is_exceeded() { + let address_a = Address::repeat_byte(1); + let address_b = Address::repeat_byte(2); + let first_row: Vec = (1..=8).collect(); + let second_row: Vec = (9..=16).collect(); + let combined: Vec = (1..=16).collect(); + + let mut collector = Collector::new(1, None); + collector + .insert(ShardedKey::new(address_a, *first_row.last().unwrap()), list(&first_row)) + .unwrap(); + collector + .insert(ShardedKey::new(address_a, *second_row.last().unwrap()), list(&second_row)) + .unwrap(); + collector.insert(ShardedKey::new(address_b, u64::MAX), list(&[1])).unwrap(); + + let first_row_bytes = super::history_group_bytes::(&first_row); + let combined_bytes = super::history_group_bytes::(&combined); + assert!( + combined_bytes > first_row_bytes, + "combined={combined_bytes} first={first_row_bytes}" + ); + + let chunks = grouped_account_chunks( + collector, + HistoryPreparationLimits { max_keys: usize::MAX, max_bytes: first_row_bytes }, + ); + + let a_indices: Vec<_> = chunks + .iter() + .flat_map(|chunk| chunk.iter().filter(|(addr, _)| *addr == address_a)) + .map(|(_, indices)| indices.clone()) + .collect(); + assert_eq!( + a_indices, + vec![combined], + "coalesced key must occupy one prepare batch: {chunks:?}" + ); + assert!(chunks.len() >= 2, "neighbor key must not share the over-budget batch: {chunks:?}"); + assert_eq!( + chunks.iter().filter(|chunk| chunk.iter().any(|(addr, _)| *addr == address_b)).count(), + 1 + ); + } +} diff --git a/crates/storage/provider/src/either_writer.rs b/crates/storage/provider/src/either_writer.rs index 0a448d1f19b..965def02161 100644 --- a/crates/storage/provider/src/either_writer.rs +++ b/crates/storage/provider/src/either_writer.rs @@ -14,7 +14,7 @@ use crate::{ }, StaticFileProviderFactory, }; -use alloy_primitives::{map::HashMap, Address, BlockNumber, TxHash, TxNumber, B256}; +use alloy_primitives::{map::HashMap, Address, BlockNumber, TxHash, TxNumber}; use rayon::slice::ParallelSliceMut; use reth_db::{ cursor::{DbCursorRO, DbDupCursorRW}, @@ -533,20 +533,6 @@ where Self::RocksDB(batch) => batch.put::(key, value), } } - - /// Gets the last shard for an address and storage key (keyed with `u64::MAX`). - pub fn get_last_storage_history_shard( - &mut self, - address: Address, - storage_key: B256, - ) -> ProviderResult> { - let key = StorageShardedKey::last(address, storage_key); - match self { - Self::Database(cursor) => Ok(cursor.seek_exact(key)?.map(|(_, v)| v)), - Self::StaticFile(_) => Err(ProviderError::UnsupportedProvider), - Self::RocksDB(batch) => batch.get::(key), - } - } } impl<'a, CURSOR, N: NodePrimitives> EitherWriter<'a, CURSOR, N> @@ -579,20 +565,6 @@ where } } - /// Gets the last shard for an address (keyed with `u64::MAX`). - pub fn get_last_account_history_shard( - &mut self, - address: Address, - ) -> ProviderResult> { - match self { - Self::Database(cursor) => { - Ok(cursor.seek_exact(ShardedKey::last(address))?.map(|(_, v)| v)) - } - Self::StaticFile(_) => Err(ProviderError::UnsupportedProvider), - Self::RocksDB(batch) => batch.get::(ShardedKey::last(address)), - } - } - /// Deletes an account history entry. pub fn delete_account_history(&mut self, key: ShardedKey
) -> ProviderResult<()> { match self { diff --git a/crates/storage/provider/src/history_shards.rs b/crates/storage/provider/src/history_shards.rs new file mode 100644 index 00000000000..2a22c94116f --- /dev/null +++ b/crates/storage/provider/src/history_shards.rs @@ -0,0 +1,471 @@ +//! Shared last-shard merge/rechunk for sharded history tables. +//! +//! Writers of [`tables::AccountsHistory`] and [`tables::StoragesHistory`] must: +//! 1. Group new indices by logical key. +//! 2. Finish all committed last-shard reads (and join Rayon workers) **before** opening a write +//! batch. +//! 3. Create the `RocksDB` batch / [`EitherWriter`](crate::EitherWriter). +//! 4. Serially put the prepared physical shards. +//! +//! Last-shard `get`s read committed state only, not pending batch writes. Getting a key after +//! putting it in the same batch would miss the merge (one logical key per batch). Do not overlap +//! parallel [`RocksDBProvider::get`](crate::providers::RocksDBProvider::get) with auto-commit +//! `write_opt`. MDBX `DbTx` is not safe for concurrent `get`s — use the serial prepare path. + +use alloy_primitives::{Address, BlockNumber, B256}; +use itertools::Itertools; +use rayon::prelude::*; +use reth_db_api::{ + models::{ + sharded_key::NUM_OF_INDICES_IN_SHARD, storage_sharded_key::StorageShardedKey, ShardedKey, + }, + table::Table, + tables, BlockNumberList, +}; +use reth_storage_errors::provider::{ProviderError, ProviderResult}; +use std::collections::BTreeMap; + +/// A history table whose keys are sharded by highest block number. +/// +/// The last shard for each logical key is stored with `highest_block = u64::MAX` so writers can +/// point-read it, merge new indices, and rechunk at [`NUM_OF_INDICES_IN_SHARD`]. +pub trait ShardedHistoryTable: Table { + /// Logical key that groups all shards of one history entry (address, or address+slot). + type PartialKey: Copy + Ord + Send + Sync; + + /// Returns the logical key for a physical shard key. + fn partial_key(key: &Self::Key) -> Self::PartialKey; + + /// Builds a physical shard key from the logical key and highest block number. + fn shard_key(key: Self::PartialKey, highest_block: BlockNumber) -> Self::Key; + + /// Physical key of the last shard (`highest_block = u64::MAX`). + fn last_shard_key(key: Self::PartialKey) -> Self::Key { + Self::shard_key(key, u64::MAX) + } +} + +impl ShardedHistoryTable for tables::AccountsHistory { + type PartialKey = Address; + + fn partial_key(key: &Self::Key) -> Self::PartialKey { + key.key + } + + fn shard_key(key: Self::PartialKey, highest_block: BlockNumber) -> Self::Key { + ShardedKey::new(key, highest_block) + } +} + +impl ShardedHistoryTable for tables::StoragesHistory { + type PartialKey = (Address, B256); + + fn partial_key(key: &Self::Key) -> Self::PartialKey { + (key.address, key.sharded_key.key) + } + + fn shard_key(key: Self::PartialKey, highest_block: BlockNumber) -> Self::Key { + StorageShardedKey::new(key.0, key.1, highest_block) + } +} + +/// Physical shard puts produced by merging new indices with each key's committed last shard. +/// +/// One inner vec per logical key, preserving that key's shard order. Outer order follows the +/// parallel (or serial) preparation walk and is not required to be sorted. +#[must_use = "prepared shard writes must be put"] +#[derive(Debug)] +pub struct PreparedHistoryShardWrites { + per_key: Vec>, +} + +impl PreparedHistoryShardWrites { + /// Flattened iterator of physical `(key, shard)` puts. + pub fn into_writes(self) -> impl Iterator { + self.per_key.into_iter().flatten() + } +} + +/// Prepares history shard writes, reading last shards in parallel. +/// +/// `get_last` must read **committed** state only (for example [`RocksDBProvider::get`]). Join all +/// reads before opening a `RocksDB` write batch: gets do not see pending puts, and parallel gets +/// must not overlap auto-commit `write_opt`. +pub fn prepare_history_shard_writes_parallel( + grouped: BTreeMap>, + get_last: F, +) -> ProviderResult> +where + T: ShardedHistoryTable, + F: Fn(T::Key) -> ProviderResult> + Send + Sync, +{ + prepare_history_shard_writes_parallel_vec(grouped.into_iter().collect(), get_last) +} + +/// Prepares history shard writes from an owned vector, reading last shards in parallel. +/// +/// `get_last` must read **committed** state only. Join all reads before opening a `RocksDB` write +/// batch. `grouped` must be strictly ordered by logical key so every key occurs exactly once. +pub fn prepare_history_shard_writes_parallel_vec( + grouped: Vec<(T::PartialKey, Vec)>, + get_last: F, +) -> ProviderResult> +where + T: ShardedHistoryTable, + F: Fn(T::Key) -> ProviderResult> + Send + Sync, +{ + validate_grouped_keys::(&grouped)?; + let per_key = grouped + .into_par_iter() + .with_min_len(1) + .map(|(partial_key, indices)| prepare_one::(partial_key, indices, &get_last)) + .collect::>>()?; + + Ok(PreparedHistoryShardWrites { per_key }) +} + +/// Prepares history shard writes, reading last shards serially. +/// +/// Use this for MDBX: a `DbTx` is not safe for concurrent `get`s. +pub fn prepare_history_shard_writes_serial( + grouped: BTreeMap>, + get_last: F, +) -> ProviderResult> +where + T: ShardedHistoryTable, + F: FnMut(T::Key) -> ProviderResult>, +{ + prepare_history_shard_writes_serial_vec(grouped.into_iter().collect(), get_last) +} + +/// Prepares history shard writes from an owned vector, reading last shards serially. +/// +/// Use this for MDBX: a `DbTx` is not safe for concurrent `get`s. `grouped` must be strictly +/// ordered by logical key so every key occurs exactly once. +pub fn prepare_history_shard_writes_serial_vec( + grouped: Vec<(T::PartialKey, Vec)>, + mut get_last: F, +) -> ProviderResult> +where + T: ShardedHistoryTable, + F: FnMut(T::Key) -> ProviderResult>, +{ + validate_grouped_keys::(&grouped)?; + let mut per_key = Vec::with_capacity(grouped.len()); + for (partial_key, indices) in grouped { + per_key.push(prepare_one::(partial_key, indices, &mut get_last)?); + } + Ok(PreparedHistoryShardWrites { per_key }) +} + +fn validate_grouped_keys( + grouped: &[(T::PartialKey, Vec)], +) -> ProviderResult<()> { + if grouped.windows(2).any(|window| window[0].0 >= window[1].0) { + return Err(ProviderError::other(InvalidGroupedHistoryKeys)) + } + Ok(()) +} + +#[derive(Debug)] +struct InvalidGroupedHistoryKeys; + +impl core::fmt::Display for InvalidGroupedHistoryKeys { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str("history shard preparation keys must be strictly ordered and unique") + } +} + +impl std::error::Error for InvalidGroupedHistoryKeys {} + +/// Merges `indices` onto the committed last shard and rechunks at [`NUM_OF_INDICES_IN_SHARD`]. +/// +/// Matches `BlockNumberList::append` + last-shard `u64::MAX` key encoding used by live persist. +fn prepare_one( + partial_key: T::PartialKey, + indices: Vec, + mut get_last: F, +) -> ProviderResult> +where + T: ShardedHistoryTable, + F: FnMut(T::Key) -> ProviderResult>, +{ + if indices.is_empty() { + return Ok(Vec::new()); + } + + debug_assert!( + indices.windows(2).all(|w| w[0] < w[1]), + "indices must be strictly increasing: {indices:?}" + ); + + let last_shard_opt = get_last(T::last_shard_key(partial_key))?; + let mut last_shard = last_shard_opt.unwrap_or_else(BlockNumberList::empty); + + last_shard.append(indices).map_err(ProviderError::other)?; + + if last_shard.len() <= NUM_OF_INDICES_IN_SHARD as u64 { + return Ok(vec![(T::last_shard_key(partial_key), last_shard)]); + } + + let chunks = last_shard.iter().chunks(NUM_OF_INDICES_IN_SHARD); + let mut chunks_peekable = chunks.into_iter().peekable(); + let mut shards = Vec::new(); + + while let Some(chunk) = chunks_peekable.next() { + let shard = BlockNumberList::new_pre_sorted(chunk); + let highest_block_number = if chunks_peekable.peek().is_some() { + shard.iter().next_back().expect("`chunks` does not return empty list") + } else { + u64::MAX + }; + + shards.push((T::shard_key(partial_key, highest_block_number), shard)); + } + + Ok(shards) +} + +#[cfg(test)] +mod tests { + use super::*; + use alloy_primitives::{address, b256}; + use std::{ + collections::HashMap, + sync::atomic::{AtomicUsize, Ordering}, + thread, + time::Duration, + }; + + fn list(blocks: &[u64]) -> BlockNumberList { + BlockNumberList::new(blocks.iter().copied()).unwrap() + } + + fn blocks(shard: &BlockNumberList) -> Vec { + shard.iter().collect() + } + + fn account_map_getter( + existing: HashMap, BlockNumberList>, + ) -> impl Fn(ShardedKey
) -> ProviderResult> { + move |key| Ok(existing.get(&key).cloned()) + } + + fn storage_map_getter( + existing: HashMap, + ) -> impl Fn(StorageShardedKey) -> ProviderResult> { + move |key| Ok(existing.get(&key).cloned()) + } + + #[test] + fn prepare_account_without_existing_shard() { + let addr = address!("0x0000000000000000000000000000000000000001"); + let grouped = BTreeMap::from([(addr, vec![1, 2, 3])]); + let prepared = prepare_history_shard_writes_serial::( + grouped, + account_map_getter(HashMap::new()), + ) + .unwrap(); + + let writes: Vec<_> = prepared.into_writes().collect(); + assert_eq!(writes.len(), 1); + assert_eq!(writes[0].0, ShardedKey::new(addr, u64::MAX)); + assert_eq!(blocks(&writes[0].1), vec![1, 2, 3]); + } + + #[test] + fn prepare_account_merges_getter_last_shard() { + let addr = address!("0x0000000000000000000000000000000000000002"); + let existing = HashMap::from([(ShardedKey::new(addr, u64::MAX), list(&[1, 2, 3]))]); + let grouped = BTreeMap::from([(addr, vec![4, 5])]); + let prepared = prepare_history_shard_writes_serial::( + grouped, + account_map_getter(existing), + ) + .unwrap(); + + let writes: Vec<_> = prepared.into_writes().collect(); + assert_eq!(writes.len(), 1); + assert_eq!(writes[0].0, ShardedKey::new(addr, u64::MAX)); + assert_eq!(blocks(&writes[0].1), vec![1, 2, 3, 4, 5]); + } + + #[test] + fn prepare_account_rechunks_at_shard_boundary() { + let addr = address!("0x0000000000000000000000000000000000000003"); + let limit = NUM_OF_INDICES_IN_SHARD as u64; + let existing_indices: Vec = (0..limit).collect(); + let existing = HashMap::from([(ShardedKey::new(addr, u64::MAX), list(&existing_indices))]); + let grouped = BTreeMap::from([(addr, vec![limit])]); + let prepared = prepare_history_shard_writes_serial::( + grouped, + account_map_getter(existing), + ) + .unwrap(); + + let writes: Vec<_> = prepared.into_writes().collect(); + assert_eq!(writes.len(), 2); + assert_eq!(writes[0].0, ShardedKey::new(addr, limit - 1)); + assert_eq!(blocks(&writes[0].1), existing_indices); + assert_eq!(writes[1].0, ShardedKey::new(addr, u64::MAX)); + assert_eq!(blocks(&writes[1].1), vec![limit]); + } + + #[test] + fn prepare_storage_without_existing_shard() { + let addr = address!("0x0000000000000000000000000000000000000010"); + let slot = b256!("0x0000000000000000000000000000000000000000000000000000000000000001"); + let grouped = BTreeMap::from([((addr, slot), vec![7, 8])]); + let prepared = prepare_history_shard_writes_serial::( + grouped, + storage_map_getter(HashMap::new()), + ) + .unwrap(); + + let writes: Vec<_> = prepared.into_writes().collect(); + assert_eq!(writes.len(), 1); + assert_eq!(writes[0].0, StorageShardedKey::new(addr, slot, u64::MAX)); + assert_eq!(blocks(&writes[0].1), vec![7, 8]); + } + + #[test] + fn prepare_storage_rechunks_at_shard_boundary() { + let addr = address!("0x0000000000000000000000000000000000000011"); + let slot = b256!("0x0000000000000000000000000000000000000000000000000000000000000002"); + let limit = NUM_OF_INDICES_IN_SHARD as u64; + let existing_indices: Vec = (0..limit).collect(); + let existing = HashMap::from([( + StorageShardedKey::new(addr, slot, u64::MAX), + list(&existing_indices), + )]); + let grouped = BTreeMap::from([((addr, slot), vec![limit, limit + 1])]); + let prepared = prepare_history_shard_writes_serial::( + grouped, + storage_map_getter(existing), + ) + .unwrap(); + + let writes: Vec<_> = prepared.into_writes().collect(); + assert_eq!(writes.len(), 2); + assert_eq!(writes[0].0, StorageShardedKey::new(addr, slot, limit - 1)); + assert_eq!(blocks(&writes[0].1).len(), NUM_OF_INDICES_IN_SHARD); + assert_eq!(writes[1].0, StorageShardedKey::new(addr, slot, u64::MAX)); + assert_eq!(blocks(&writes[1].1), vec![limit, limit + 1]); + } + + #[test] + fn prepare_skips_empty_indices() { + let addr = address!("0x0000000000000000000000000000000000000004"); + let grouped = BTreeMap::from([(addr, Vec::new())]); + let prepared = prepare_history_shard_writes_serial::( + grouped, + account_map_getter(HashMap::new()), + ) + .unwrap(); + assert!(prepared.into_writes().next().is_none()); + } + + #[test] + fn prepare_parallel_matches_serial() { + let addr_a = address!("0x00000000000000000000000000000000000000aa"); + let addr_b = address!("0x00000000000000000000000000000000000000bb"); + let existing = HashMap::from([ + (ShardedKey::new(addr_a, u64::MAX), list(&[1, 2])), + (ShardedKey::new(addr_b, u64::MAX), list(&[10])), + ]); + let grouped = BTreeMap::from([(addr_a, vec![3, 4]), (addr_b, vec![11])]); + + let serial = prepare_history_shard_writes_serial::( + grouped.clone(), + account_map_getter(existing.clone()), + ) + .unwrap(); + let parallel = prepare_history_shard_writes_parallel::( + grouped, + account_map_getter(existing), + ) + .unwrap(); + + let mut serial_writes: Vec<_> = + serial.into_writes().map(|(k, v)| (k, blocks(&v))).collect(); + let mut parallel_writes: Vec<_> = + parallel.into_writes().map(|(k, v)| (k, blocks(&v))).collect(); + serial_writes.sort_by(|a, b| a.0.cmp(&b.0)); + parallel_writes.sort_by(|a, b| a.0.cmp(&b.0)); + assert_eq!(serial_writes, parallel_writes); + } + + #[test] + fn prepare_vector_parallel_matches_serial() { + let addr_a = address!("0x00000000000000000000000000000000000000aa"); + let addr_b = address!("0x00000000000000000000000000000000000000bb"); + let existing = HashMap::from([ + (ShardedKey::new(addr_a, u64::MAX), list(&[1, 2])), + (ShardedKey::new(addr_b, u64::MAX), list(&[10])), + ]); + let grouped = vec![(addr_a, vec![3, 4]), (addr_b, vec![11])]; + + let serial = prepare_history_shard_writes_serial_vec::( + grouped.clone(), + account_map_getter(existing.clone()), + ) + .unwrap(); + let parallel = prepare_history_shard_writes_parallel_vec::( + grouped, + account_map_getter(existing), + ) + .unwrap(); + + let serial_writes: Vec<_> = + serial.into_writes().map(|(key, value)| (key, blocks(&value))).collect(); + let parallel_writes: Vec<_> = + parallel.into_writes().map(|(key, value)| (key, blocks(&value))).collect(); + assert_eq!(serial_writes, parallel_writes); + } + + #[test] + fn prepare_vector_rejects_duplicate_keys() { + let address = address!("0x00000000000000000000000000000000000000aa"); + let grouped = vec![(address, vec![1]), (address, vec![2])]; + + let error = prepare_history_shard_writes_serial_vec::( + grouped, + account_map_getter(HashMap::new()), + ) + .unwrap_err(); + + assert!(error.is_other::()); + } + + #[test] + fn prepare_parallel_overlaps_gets_on_dedicated_pool() { + let pool = rayon::ThreadPoolBuilder::new().num_threads(4).build().unwrap(); + let inflight = AtomicUsize::new(0); + let max_inflight = AtomicUsize::new(0); + + let grouped: Vec<_> = + (1u8..=8).map(|i| (Address::repeat_byte(i), vec![u64::from(i)])).collect(); + + let prepared = pool + .install(|| { + prepare_history_shard_writes_parallel_vec::( + grouped, + |_| { + let current = inflight.fetch_add(1, Ordering::SeqCst) + 1; + max_inflight.fetch_max(current, Ordering::SeqCst); + thread::sleep(Duration::from_millis(50)); + inflight.fetch_sub(1, Ordering::SeqCst); + Ok(None) + }, + ) + }) + .unwrap(); + assert_eq!(prepared.into_writes().count(), 8); + + assert!( + max_inflight.load(Ordering::SeqCst) >= 2, + "parallel last-shard gets should overlap on a 4-thread pool, max inflight was {}", + max_inflight.load(Ordering::SeqCst) + ); + } +} diff --git a/crates/storage/provider/src/lib.rs b/crates/storage/provider/src/lib.rs index ce9f8d7ba44..e549f1b361f 100644 --- a/crates/storage/provider/src/lib.rs +++ b/crates/storage/provider/src/lib.rs @@ -37,6 +37,13 @@ pub mod test_utils; pub mod either_writer; pub use either_writer::*; +pub mod history_shards; +pub use history_shards::{ + prepare_history_shard_writes_parallel, prepare_history_shard_writes_parallel_vec, + prepare_history_shard_writes_serial, prepare_history_shard_writes_serial_vec, + PreparedHistoryShardWrites, ShardedHistoryTable, +}; + mod bal; pub use bal::{BalConfig, InMemoryBalStore, RocksDBBalStore}; diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index 4aee893a1bd..405fdb07e14 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/crates/storage/provider/src/providers/database/provider.rs @@ -1,6 +1,7 @@ use super::SaveBlocksInput; use crate::{ changesets_utils::StorageRevertsIter, + prepare_history_shard_writes_parallel, providers::{ database::{chain::ChainStorage, metrics, DatabaseProviderMetrics}, rocksdb::{ @@ -3818,13 +3819,17 @@ impl BlockWriter // Use pre-computed transitions for history indices since static file // writes aren't visible until commit. - // Note: For MDBX we use insert_*_history_index. For RocksDB we use - // append_*_history_shard which handles read-merge-write internally. + // RocksDB: committed last-shard reads finish before the write batch is created. let storage_settings = self.cached_storage_settings(); if storage_settings.storage_v2 { + let rocksdb = self.rocksdb_provider(); + let prepared = prepare_history_shard_writes_parallel::( + account_transitions, + |key| rocksdb.get::(key), + )?; self.with_rocksdb_batch(|mut batch| { - for (address, blocks) in account_transitions { - batch.append_account_history_shard(address, blocks)?; + for (key, shard) in prepared.into_writes() { + batch.put::(key, &shard)?; } Ok(((), Some(batch.into_inner()))) })?; @@ -3832,9 +3837,14 @@ impl BlockWriter self.insert_account_history_index(account_transitions)?; } if storage_settings.storage_v2 { + let rocksdb = self.rocksdb_provider(); + let prepared = prepare_history_shard_writes_parallel::( + storage_transitions, + |key| rocksdb.get::(key), + )?; self.with_rocksdb_batch(|mut batch| { - for ((address, key), blocks) in storage_transitions { - batch.append_storage_history_shard(address, key, blocks)?; + for (key, shard) in prepared.into_writes() { + batch.put::(key, &shard)?; } Ok(((), Some(batch.into_inner()))) })?; diff --git a/crates/storage/provider/src/providers/rocksdb/provider.rs b/crates/storage/provider/src/providers/rocksdb/provider.rs index b681c382c64..a738f2c63ed 100644 --- a/crates/storage/provider/src/providers/rocksdb/provider.rs +++ b/crates/storage/provider/src/providers/rocksdb/provider.rs @@ -1,21 +1,22 @@ use super::metrics::{RocksDBMetrics, RocksDBOperation, ROCKSDB_TABLES}; -use crate::providers::{compute_history_rank, needs_prev_shard_check, HistoryInfo}; +use crate::{ + history_shards::{ + prepare_history_shard_writes_parallel, prepare_history_shard_writes_serial, + ShardedHistoryTable, + }, + providers::{compute_history_rank, needs_prev_shard_check, HistoryInfo}, +}; use alloy_consensus::transaction::TxHashRef; use alloy_primitives::{ map::{AddressMap, HashMap}, Address, BlockNumber, TxNumber, B256, }; -use itertools::Itertools; use metrics::Label; use parking_lot::Mutex; -use rayon::prelude::*; use reth_chain_state::ExecutedBlock; use reth_db_api::{ database_metrics::DatabaseMetrics, - models::{ - sharded_key::NUM_OF_INDICES_IN_SHARD, storage_sharded_key::StorageShardedKey, ShardedKey, - StorageSettings, - }, + models::{storage_sharded_key::StorageShardedKey, ShardedKey, StorageSettings}, table::{Compress, Decode, Decompress, Encode, Table}, tables, BlockNumberList, DatabaseError, }; @@ -911,7 +912,7 @@ impl RocksDBProvider { /// Creates a new batch with auto-commit enabled. /// - /// When the batch size exceeds the threshold (4 GiB), the batch is automatically + /// When the batch size exceeds the threshold (512 MiB), the batch is automatically /// committed and reset. This prevents OOM during large bulk writes while maintaining /// crash-safety via the consistency check on startup. pub fn batch_with_auto_commit(&self) -> RocksDBBatch<'_> { @@ -1523,7 +1524,6 @@ impl RocksDBProvider { blocks: &[ExecutedBlock], ctx: &RocksDBWriteCtx, ) -> ProviderResult<()> { - let mut batch = self.batch(); let mut account_history: BTreeMap> = BTreeMap::new(); for (block_idx, block) in blocks.iter().enumerate() { @@ -1539,12 +1539,7 @@ impl RocksDBProvider { } } - // Write account history using proper shard append logic - for (address, indices) in account_history { - batch.append_account_history_shard(address, indices)?; - } - ctx.pending_batches.lock().push(batch.into_inner()); - Ok(()) + self.write_history::(account_history, ctx) } /// Writes storage history indices for the given blocks. @@ -1577,68 +1572,28 @@ impl RocksDBProvider { } } - let shard_puts = storage_history - .into_par_iter() - .map(|((address, slot), indices)| { - self.storage_history_shards_to_put(address, slot, indices) - }) - .collect::>>()?; - - let mut batch = self.batch(); - for shards in shard_puts { - for (key, shard) in shards { - batch.put::(key, &shard)?; - } - } - ctx.pending_batches.lock().push(batch.into_inner()); - Ok(()) + self.write_history::(storage_history, ctx) } - /// Prepares storage history shard writes by reading the current last shard and appending - /// indices. - fn storage_history_shards_to_put( + /// Parallel last-shard reads, then a serial batch of prepared puts. + /// + /// Last-shard `get`s read committed state only. All of them (and Rayon workers) finish before + /// [`Self::batch`] is created so a key is never read after it has been put, and so parallel + /// gets do not overlap `write_opt`. + fn write_history( &self, - address: Address, - storage_key: B256, - indices: Vec, - ) -> ProviderResult> { - if indices.is_empty() { - return Ok(Vec::new()); - } - - debug_assert!( - indices.windows(2).all(|w| w[0] < w[1]), - "indices must be strictly increasing: {:?}", - indices - ); - - let last_key = StorageShardedKey::last(address, storage_key); - let last_shard_opt = self.get::(last_key.clone())?; - let mut last_shard = last_shard_opt.unwrap_or_else(BlockNumberList::empty); - - last_shard.append(indices).map_err(ProviderError::other)?; - - if last_shard.len() <= NUM_OF_INDICES_IN_SHARD as u64 { - return Ok(vec![(last_key, last_shard)]); - } - - let chunks = last_shard.iter().chunks(NUM_OF_INDICES_IN_SHARD); - let mut chunks_peekable = chunks.into_iter().peekable(); - let mut shards = Vec::new(); - - while let Some(chunk) = chunks_peekable.next() { - let shard = BlockNumberList::new_pre_sorted(chunk); - let highest_block_number = if chunks_peekable.peek().is_some() { - shard.iter().next_back().expect("`chunks` does not return empty list") - } else { - u64::MAX - }; + grouped: BTreeMap>, + ctx: &RocksDBWriteCtx, + ) -> ProviderResult<()> { + let prepared = + prepare_history_shard_writes_parallel::(grouped, |key| self.get::(key))?; - shards - .push((StorageShardedKey::new(address, storage_key, highest_block_number), shard)); + let mut batch = self.batch(); + for (key, shard) in prepared.into_writes() { + batch.put::(key, &shard)?; } - - Ok(shards) + ctx.pending_batches.lock().push(batch.into_inner()); + Ok(()) } } @@ -2109,47 +2064,13 @@ impl<'a> RocksDBBatch<'a> { indices: impl IntoIterator, ) -> ProviderResult<()> { let indices: Vec = indices.into_iter().collect(); - - if indices.is_empty() { - return Ok(()); - } - - debug_assert!( - indices.windows(2).all(|w| w[0] < w[1]), - "indices must be strictly increasing: {:?}", - indices - ); - - let last_key = ShardedKey::new(address, u64::MAX); - let last_shard_opt = self.provider.get::(last_key.clone())?; - let mut last_shard = last_shard_opt.unwrap_or_else(BlockNumberList::empty); - - last_shard.append(indices).map_err(ProviderError::other)?; - - // Fast path: all indices fit in one shard - if last_shard.len() <= NUM_OF_INDICES_IN_SHARD as u64 { - self.put::(last_key, &last_shard)?; - return Ok(()); - } - - // Slow path: rechunk into multiple shards - let chunks = last_shard.iter().chunks(NUM_OF_INDICES_IN_SHARD); - let mut chunks_peekable = chunks.into_iter().peekable(); - - while let Some(chunk) = chunks_peekable.next() { - let shard = BlockNumberList::new_pre_sorted(chunk); - let highest_block_number = if chunks_peekable.peek().is_some() { - shard.iter().next_back().expect("`chunks` does not return empty list") - } else { - u64::MAX - }; - - self.put::( - ShardedKey::new(address, highest_block_number), - &shard, - )?; + let prepared = prepare_history_shard_writes_serial::( + BTreeMap::from([(address, indices)]), + |key| self.provider.get::(key), + )?; + for (key, shard) in prepared.into_writes() { + self.put::(key, &shard)?; } - Ok(()) } @@ -2171,13 +2092,13 @@ impl<'a> RocksDBBatch<'a> { indices: impl IntoIterator, ) -> ProviderResult<()> { let indices: Vec = indices.into_iter().collect(); - - for (key, shard) in - self.provider.storage_history_shards_to_put(address, storage_key, indices)? - { + let prepared = prepare_history_shard_writes_serial::( + BTreeMap::from([((address, storage_key), indices)]), + |key| self.provider.get::(key), + )?; + for (key, shard) in prepared.into_writes() { self.put::(key, &shard)?; } - Ok(()) }