diff --git a/crates/storage/db-api/src/models/integer_list.rs b/crates/storage/db-api/src/models/integer_list.rs index 3a6d1738a82..06a993e75e8 100644 --- a/crates/storage/db-api/src/models/integer_list.rs +++ b/crates/storage/db-api/src/models/integer_list.rs @@ -2,7 +2,7 @@ use crate::table::{Compress, Decompress}; use bytes::BufMut; -use core::fmt; +use core::{fmt, ops::RangeBounds}; use derive_more::Deref; use reth_codecs::DecompressError; use roaring::RoaringTreemap; @@ -68,6 +68,11 @@ impl IntegerList { self.0.clear(); } + /// Removes the integers in the given range, returning how many were removed. + pub fn remove_range>(&mut self, range: R) -> u64 { + self.0.remove_range(range) + } + /// Serializes an [`IntegerList`] into a sequence of bytes. pub fn to_bytes(&self) -> Vec { let mut vec = Vec::with_capacity(self.0.serialized_size()); @@ -199,4 +204,27 @@ mod tests { let blist = ef_list.to_bytes(); assert_eq!(IntegerList::from_bytes(&blist).unwrap(), ef_list) } + + #[test] + fn remove_range_matches_filtering() { + // Spans more than one 2^16 roaring container so multi-container removal is covered. + let values = [1u64, 2, 100, 65_535, 65_536, 70_000, 200_000]; + + for to_block in [0u64, 1, 99, 100, 65_535, 65_536, 199_999, 200_000, 200_001] { + let mut list = IntegerList::new(values).unwrap(); + let removed = list.remove_range(0..=to_block); + + let expected = values.into_iter().filter(|value| *value > to_block).collect::>(); + assert_eq!(list.iter().collect::>(), expected, "to_block {to_block}"); + assert_eq!(removed, (values.len() - expected.len()) as u64, "to_block {to_block}"); + assert_eq!(list.is_empty(), expected.is_empty(), "to_block {to_block}"); + } + } + + #[test] + fn remove_range_on_empty_list_removes_nothing() { + let mut list = IntegerList::empty(); + assert_eq!(list.remove_range(0..=100), 0); + assert!(list.is_empty()); + } } diff --git a/crates/storage/provider/src/providers/rocksdb/provider.rs b/crates/storage/provider/src/providers/rocksdb/provider.rs index 812b3938536..90ea1a0db1c 100644 --- a/crates/storage/provider/src/providers/rocksdb/provider.rs +++ b/crates/storage/provider/src/providers/rocksdb/provider.rs @@ -2044,10 +2044,14 @@ impl<'a> RocksDBBatch<'a> { /// Generic implementation for both account and storage history pruning. /// Mirrors MDBX `prune_shard` semantics. After pruning, the last remaining shard /// (if any) will have the sentinel key (`u64::MAX`). + /// + /// `shards_complete` must be `false` when `shards` stops short of the key's last shard, so the + /// survivor is not re-keyed over a sentinel that is still on disk. #[expect(clippy::too_many_arguments)] fn prune_history_shards_inner( &mut self, shards: Vec<(K, BlockNumberList)>, + shards_complete: bool, to_block: BlockNumber, get_highest: impl Fn(&K) -> u64, is_sentinel: impl Fn(&K) -> bool, @@ -2066,21 +2070,19 @@ impl<'a> RocksDBBatch<'a> { let mut updated = false; let mut last_remaining: Option<(K, BlockNumberList)> = None; - for (key, block_list) in shards { + for (key, mut block_list) in shards { if !is_sentinel(&key) && get_highest(&key) <= to_block { delete_shard(self, key)?; deleted = true; } else { - let original_len = block_list.len(); - let filtered = - BlockNumberList::new_pre_sorted(block_list.iter().filter(|&b| b > to_block)); + let removed = block_list.remove_range(0..=to_block); - if filtered.is_empty() { + if block_list.is_empty() { delete_shard(self, key)?; deleted = true; - } else if filtered.len() < original_len { - put_shard(self, key.clone(), &filtered)?; - last_remaining = Some((key, filtered)); + } else if removed > 0 { + put_shard(self, key.clone(), &block_list)?; + last_remaining = Some((key, block_list)); updated = true; } else { last_remaining = Some((key, block_list)); @@ -2088,7 +2090,8 @@ impl<'a> RocksDBBatch<'a> { } } - if let Some((last_key, last_value)) = last_remaining && + if shards_complete && + let Some((last_key, last_value)) = last_remaining && !is_sentinel(&last_key) { delete_shard(self, last_key)?; @@ -2117,6 +2120,7 @@ impl<'a> RocksDBBatch<'a> { let shards = self.provider.account_history_shards(address)?; self.prune_history_shards_inner( shards, + true, to_block, |key| key.highest_block_number, |key| key.highest_block_number == u64::MAX, @@ -2132,8 +2136,8 @@ impl<'a> RocksDBBatch<'a> { /// because it reuses a single raw iterator and skips seeks when the iterator is already /// positioned correctly (which happens when targets are sorted and adjacent in key order). /// - /// `targets` MUST be sorted by address for correctness and optimal performance - /// (matches on-disk key order). + /// `targets` MUST be sorted by address and contain each address at most once, for + /// correctness and optimal performance (matches on-disk key order). pub fn prune_account_history_batch( &mut self, targets: &[(Address, BlockNumber)], @@ -2143,8 +2147,8 @@ impl<'a> RocksDBBatch<'a> { } debug_assert!( - targets.windows(2).all(|w| w[0].0 <= w[1].0), - "prune_account_history_batch: targets must be sorted by address" + targets.windows(2).all(|w| w[0].0 < w[1].0), + "prune_account_history_batch: targets must be sorted and unique" ); // ShardedKey
layout: [address: 20][block: 8] = 28 bytes @@ -2163,7 +2167,9 @@ impl<'a> RocksDBBatch<'a> { // Check if we need to seek or if the iterator is already positioned correctly. // After processing the previous target, the iterator is either: // 1. Positioned at a key with a different prefix (we iterated past our shards) - // 2. Invalid (no more keys) + // 2. Positioned on a later shard of the previous target (we stopped early), whose + // prefix is below ours because targets are sorted and unique + // 3. Invalid (no more keys) // If the current key's prefix >= our target prefix, we may be able to skip the seek. let needs_seek = if iter.valid() { if let Some(current_key) = iter.key() { @@ -2188,8 +2194,10 @@ impl<'a> RocksDBBatch<'a> { })?; } - // Collect all shards for this address using raw prefix comparison + // Collect the shards for this address that pruning can touch, using raw prefix + // comparison let mut shards = Vec::new(); + let mut shards_complete = true; while iter.valid() { let Some(key_bytes) = iter.key() else { break }; @@ -2207,12 +2215,36 @@ impl<'a> RocksDBBatch<'a> { let value = BlockNumberList::decompress(value_bytes) .map_err(|_| ProviderError::Database(DatabaseError::Decode))?; + let highest = key.highest_block_number; shards.push((key, value)); + iter.next(); + + // Shards are ordered by their highest block and their contents partition the + // address's history, so this is the last shard holding anything at or below the + // target. Peek past it only to tell whether it was the address's last shard, + // which decides whether a survivor may be re-keyed to the sentinel. + if highest > *to_block { + shards_complete = iter.key().and_then(|next_key| next_key.get(..PREFIX_LEN)) != + Some(target_prefix); + break; + } + } + + // The iterator also goes invalid on a read error, which would otherwise pass a + // truncated shard list off as the key's complete one. + if !iter.valid() { + iter.status().map_err(|e| { + ProviderError::Database(DatabaseError::Read(DatabaseErrorInfo { + message: e.to_string().into(), + code: -1, + })) + })?; } match self.prune_history_shards_inner( shards, + shards_complete, *to_block, |key| key.highest_block_number, |key| key.highest_block_number == u64::MAX, @@ -2243,6 +2275,7 @@ impl<'a> RocksDBBatch<'a> { let shards = self.provider.storage_history_shards(address, storage_key)?; self.prune_history_shards_inner( shards, + true, to_block, |key| key.sharded_key.highest_block_number, |key| key.sharded_key.highest_block_number == u64::MAX, @@ -2259,8 +2292,8 @@ impl<'a> RocksDBBatch<'a> { /// because it reuses a single raw iterator and skips seeks when the iterator is already /// positioned correctly (which happens when targets are sorted and adjacent in key order). /// - /// `targets` MUST be sorted by (address, `storage_key`) for correctness and optimal - /// performance (matches on-disk key order). + /// `targets` MUST be sorted by (address, `storage_key`) and contain each pair at most once, + /// for correctness and optimal performance (matches on-disk key order). pub fn prune_storage_history_batch( &mut self, targets: &[((Address, B256), BlockNumber)], @@ -2270,8 +2303,8 @@ impl<'a> RocksDBBatch<'a> { } debug_assert!( - targets.windows(2).all(|w| w[0].0 <= w[1].0), - "prune_storage_history_batch: targets must be sorted by (address, storage_key)" + targets.windows(2).all(|w| w[0].0 < w[1].0), + "prune_storage_history_batch: targets must be sorted and unique" ); // StorageShardedKey layout: [address: 20][storage_key: 32][block: 8] = 60 bytes @@ -2290,7 +2323,9 @@ impl<'a> RocksDBBatch<'a> { // Check if we need to seek or if the iterator is already positioned correctly. // After processing the previous target, the iterator is either: // 1. Positioned at a key with a different prefix (we iterated past our shards) - // 2. Invalid (no more keys) + // 2. Positioned on a later shard of the previous target (we stopped early), whose + // prefix is below ours because targets are sorted and unique + // 3. Invalid (no more keys) // If the current key's prefix >= our target prefix, we may be able to skip the seek. let needs_seek = if iter.valid() { if let Some(current_key) = iter.key() { @@ -2315,8 +2350,10 @@ impl<'a> RocksDBBatch<'a> { })?; } - // Collect all shards for this (address, storage_key) pair using prefix comparison + // Collect the shards for this (address, storage_key) pair that pruning can touch, + // using prefix comparison let mut shards = Vec::new(); + let mut shards_complete = true; while iter.valid() { let Some(key_bytes) = iter.key() else { break }; @@ -2334,13 +2371,37 @@ impl<'a> RocksDBBatch<'a> { let value = BlockNumberList::decompress(value_bytes) .map_err(|_| ProviderError::Database(DatabaseError::Decode))?; + let highest = key.sharded_key.highest_block_number; shards.push((key, value)); + iter.next(); + + // Shards are ordered by their highest block and their contents partition the + // key's history, so this is the last shard holding anything at or below the + // target. Peek past it only to tell whether it was the key's last shard, which + // decides whether a survivor may be re-keyed to the sentinel. + if highest > *to_block { + shards_complete = iter.key().and_then(|next_key| next_key.get(..PREFIX_LEN)) != + Some(target_prefix); + break; + } + } + + // The iterator also goes invalid on a read error, which would otherwise pass a + // truncated shard list off as the key's complete one. + if !iter.valid() { + iter.status().map_err(|e| { + ProviderError::Database(DatabaseError::Read(DatabaseErrorInfo { + message: e.to_string().into(), + code: -1, + })) + })?; } // Use existing prune_history_shards_inner logic match self.prune_history_shards_inner( shards, + shards_complete, *to_block, |key| key.sharded_key.highest_block_number, |key| key.sharded_key.highest_block_number == u64::MAX, @@ -4533,4 +4594,204 @@ mod tests { let shards2 = provider.storage_history_shards(addr, slot2).unwrap(); assert_eq!(shards2[0].1.iter().collect::>(), vec![15, 25]); } + + /// Shards for one address, keyed by highest block, as `(highest, blocks)`. + fn account_shard_layout(provider: &RocksDBProvider, address: Address) -> Vec<(u64, Vec)> { + provider + .account_history_shards(address) + .unwrap() + .into_iter() + .map(|(key, list)| (key.highest_block_number, list.iter().collect::>())) + .collect() + } + + /// Shards for one storage slot, keyed by highest block, as `(highest, blocks)`. + fn storage_shard_layout( + provider: &RocksDBProvider, + address: Address, + storage_key: B256, + ) -> Vec<(u64, Vec)> { + provider + .storage_history_shards(address, storage_key) + .unwrap() + .into_iter() + .map(|(key, list)| { + (key.sharded_key.highest_block_number, list.iter().collect::>()) + }) + .collect() + } + + fn seed_three_storage_shards(provider: &RocksDBProvider, address: Address, storage_key: B256) { + let mut batch = provider.batch(); + batch + .put::( + StorageShardedKey::new(address, storage_key, 100), + &BlockNumberList::new_pre_sorted([10, 50, 100]), + ) + .unwrap(); + batch + .put::( + StorageShardedKey::new(address, storage_key, 200), + &BlockNumberList::new_pre_sorted([150, 200]), + ) + .unwrap(); + batch + .put::( + StorageShardedKey::last(address, storage_key), + &BlockNumberList::new_pre_sorted([250, 300]), + ) + .unwrap(); + batch.commit().unwrap(); + } + + #[test] + fn test_prune_storage_history_batch_leaves_shards_above_target_untouched() { + let temp_dir = TempDir::new().unwrap(); + let provider = RocksDBBuilder::new(temp_dir.path()).with_default_tables().build().unwrap(); + + let addr = Address::from([0x42; 20]); + let slot = B256::from([0x01; 32]); + seed_three_storage_shards(&provider, addr, slot); + + // Only the oldest shard holds blocks at or below the target. + let mut batch = provider.batch(); + let outcomes = batch.prune_storage_history_batch(&[((addr, slot), 50)]).unwrap(); + batch.commit().unwrap(); + + assert_eq!(outcomes.updated, 1); + // The trimmed shard keeps its own key. Re-keying it to the sentinel here would overwrite + // the sentinel's blocks. + assert_eq!( + storage_shard_layout(&provider, addr, slot), + vec![(100, vec![100]), (200, vec![150, 200]), (u64::MAX, vec![250, 300])] + ); + } + + #[test] + fn test_prune_storage_history_batch_trims_sentinel_once_earlier_shards_expire() { + let temp_dir = TempDir::new().unwrap(); + let provider = RocksDBBuilder::new(temp_dir.path()).with_default_tables().build().unwrap(); + + let addr = Address::from([0x42; 20]); + let slot = B256::from([0x01; 32]); + seed_three_storage_shards(&provider, addr, slot); + + // Every non-sentinel shard expires whole and the sentinel loses its lowest block. + let mut batch = provider.batch(); + let outcomes = batch.prune_storage_history_batch(&[((addr, slot), 250)]).unwrap(); + batch.commit().unwrap(); + + assert_eq!(outcomes.deleted, 1); + assert_eq!(storage_shard_layout(&provider, addr, slot), vec![(u64::MAX, vec![300])]); + } + + #[test] + fn test_prune_account_history_batch_leaves_shards_above_target_untouched() { + let temp_dir = TempDir::new().unwrap(); + let provider = RocksDBBuilder::new(temp_dir.path()).with_default_tables().build().unwrap(); + + let addr = Address::from([0x42; 20]); + + let mut batch = provider.batch(); + batch + .put::( + ShardedKey::new(addr, 100), + &BlockNumberList::new_pre_sorted([10, 50, 100]), + ) + .unwrap(); + batch + .put::( + ShardedKey::new(addr, u64::MAX), + &BlockNumberList::new_pre_sorted([250, 300]), + ) + .unwrap(); + batch.commit().unwrap(); + + let mut batch = provider.batch(); + let outcomes = batch.prune_account_history_batch(&[(addr, 50)]).unwrap(); + batch.commit().unwrap(); + + assert_eq!(outcomes.updated, 1); + assert_eq!( + account_shard_layout(&provider, addr), + vec![(100, vec![100]), (u64::MAX, vec![250, 300])] + ); + } + + #[test] + fn test_prune_account_history_batch_seeks_after_stopping_early() { + let temp_dir = TempDir::new().unwrap(); + let provider = RocksDBBuilder::new(temp_dir.path()).with_default_tables().build().unwrap(); + + let addr1 = Address::from([0x01; 20]); + let addr2 = Address::from([0x02; 20]); + + let mut batch = provider.batch(); + batch + .put::( + ShardedKey::new(addr1, 100), + &BlockNumberList::new_pre_sorted([10, 50, 100]), + ) + .unwrap(); + batch + .put::( + ShardedKey::new(addr1, u64::MAX), + &BlockNumberList::new_pre_sorted([250, 300]), + ) + .unwrap(); + batch + .put::( + ShardedKey::new(addr2, u64::MAX), + &BlockNumberList::new_pre_sorted([5, 10, 15]), + ) + .unwrap(); + batch.commit().unwrap(); + + // The first target stops on addr1's oldest shard, leaving the iterator on addr1's + // sentinel. The second target must seek past it instead of skipping addr2. + let mut batch = provider.batch(); + let outcomes = batch.prune_account_history_batch(&[(addr1, 50), (addr2, 10)]).unwrap(); + batch.commit().unwrap(); + + assert_eq!(outcomes.updated, 2); + assert_eq!( + account_shard_layout(&provider, addr1), + vec![(100, vec![100]), (u64::MAX, vec![250, 300])] + ); + assert_eq!(account_shard_layout(&provider, addr2), vec![(u64::MAX, vec![15])]); + } + + #[test] + fn test_prune_storage_history_batch_seeks_after_stopping_early() { + let temp_dir = TempDir::new().unwrap(); + let provider = RocksDBBuilder::new(temp_dir.path()).with_default_tables().build().unwrap(); + + let addr = Address::from([0x42; 20]); + let slot1 = B256::from([0x01; 32]); + let slot2 = B256::from([0x02; 32]); + seed_three_storage_shards(&provider, addr, slot1); + + let mut batch = provider.batch(); + batch + .put::( + StorageShardedKey::last(addr, slot2), + &BlockNumberList::new_pre_sorted([20, 40]), + ) + .unwrap(); + batch.commit().unwrap(); + + // The first target stops on slot1's oldest shard, leaving the iterator on slot1's next + // shard. The second target must seek past it instead of skipping slot2. + let mut batch = provider.batch(); + let outcomes = + batch.prune_storage_history_batch(&[((addr, slot1), 50), ((addr, slot2), 30)]).unwrap(); + batch.commit().unwrap(); + + assert_eq!(outcomes.updated, 2); + assert_eq!( + storage_shard_layout(&provider, addr, slot1), + vec![(100, vec![100]), (200, vec![150, 200]), (u64::MAX, vec![250, 300])] + ); + assert_eq!(storage_shard_layout(&provider, addr, slot2), vec![(u64::MAX, vec![40])]); + } }