1use crate::PruneLimiter;
2use reth_db_api::{
3 cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO, RangeWalker},
4 table::{DupSort, Table, TableRow},
5 transaction::{DbTx, DbTxMut},
6 DatabaseError,
7};
8use std::{fmt::Debug, ops::RangeBounds};
9use tracing::debug;
10
11#[derive(Debug, Clone, Copy)]
13pub(crate) struct PruneStepResult {
14 done: bool,
16 deleted: bool,
18}
19
20pub(crate) trait DbTxPruneExt: DbTxMut + DbTx {
21 fn clear_table<T: Table>(&self) -> Result<usize, DatabaseError> {
26 let count = self.entries::<T>()?;
27 <Self as DbTxMut>::clear::<T>(self)?;
28 Ok(count)
29 }
30
31 fn prune_table_with_iterator<T: Table>(
35 &self,
36 keys: impl IntoIterator<Item = T::Key>,
37 limiter: &mut PruneLimiter,
38 mut delete_callback: impl FnMut(TableRow<T>),
39 ) -> Result<(usize, bool), DatabaseError> {
40 let mut cursor = self.cursor_write::<T>()?;
41 let mut keys = keys.into_iter().peekable();
42
43 let mut deleted_entries = 0;
44
45 let mut done = true;
46 while keys.peek().is_some() {
47 if limiter.is_limit_reached() {
48 debug!(
49 target: "providers::db",
50 ?limiter,
51 deleted_entries_limit = %limiter.is_deleted_entries_limit_reached(),
52 time_limit = %limiter.is_time_limit_reached(),
53 table = %T::NAME,
54 "Pruning limit reached"
55 );
56 done = false;
57 break
58 }
59
60 let key = keys.next().expect("peek() said Some");
61 let row = cursor.seek_exact(key)?;
62 if let Some(row) = row {
63 cursor.delete_current()?;
64 limiter.increment_deleted_entries_count();
65 deleted_entries += 1;
66 delete_callback(row);
67 }
68 }
69
70 Ok((deleted_entries, done))
71 }
72
73 fn prune_table_with_range<T: Table>(
77 &self,
78 keys: impl RangeBounds<T::Key> + Clone + Debug,
79 limiter: &mut PruneLimiter,
80 mut skip_filter: impl FnMut(&TableRow<T>) -> bool,
81 mut delete_callback: impl FnMut(TableRow<T>),
82 ) -> Result<(usize, bool), DatabaseError> {
83 let mut cursor = self.cursor_write::<T>()?;
84 let mut walker = cursor.walk_range(keys)?;
85
86 let mut deleted_entries = 0;
87
88 let done = loop {
89 if limiter.is_limit_reached() {
92 debug!(
93 target: "providers::db",
94 ?limiter,
95 deleted_entries_limit = %limiter.is_deleted_entries_limit_reached(),
96 time_limit = %limiter.is_time_limit_reached(),
97 table = %T::NAME,
98 "Pruning limit reached"
99 );
100 break false
101 }
102
103 let result = self.prune_table_with_range_step(
104 &mut walker,
105 limiter,
106 &mut skip_filter,
107 &mut delete_callback,
108 )?;
109
110 if result.deleted {
111 deleted_entries += 1;
112 }
113
114 if result.done {
115 break true
116 }
117 };
118
119 Ok((deleted_entries, done))
120 }
121
122 fn prune_table_with_range_step<T: Table>(
128 &self,
129 walker: &mut RangeWalker<'_, T, Self::CursorMut<T>>,
130 limiter: &mut PruneLimiter,
131 skip_filter: &mut impl FnMut(&TableRow<T>) -> bool,
132 delete_callback: &mut impl FnMut(TableRow<T>),
133 ) -> Result<PruneStepResult, DatabaseError> {
134 let Some(res) = walker.next() else {
135 return Ok(PruneStepResult { done: true, deleted: false })
136 };
137
138 let row = res?;
139
140 if skip_filter(&row) {
141 Ok(PruneStepResult { done: false, deleted: false })
142 } else {
143 walker.delete_current()?;
144 limiter.increment_deleted_entries_count();
145 delete_callback(row);
146 Ok(PruneStepResult { done: false, deleted: true })
147 }
148 }
149
150 #[expect(unused)]
154 fn prune_dupsort_table_with_range<T: DupSort>(
155 &self,
156 keys: impl RangeBounds<T::Key> + Clone + Debug,
157 limiter: &mut PruneLimiter,
158 mut delete_callback: impl FnMut(TableRow<T>),
159 ) -> Result<(usize, bool), DatabaseError> {
160 let starting_entries = self.entries::<T>()?;
161 let mut cursor = self.cursor_dup_write::<T>()?;
162 let mut walker = cursor.walk_range(keys)?;
163
164 let done = loop {
165 if limiter.is_limit_reached() {
166 debug!(
167 target: "providers::db",
168 ?limiter,
169 deleted_entries_limit = %limiter.is_deleted_entries_limit_reached(),
170 time_limit = %limiter.is_time_limit_reached(),
171 table = %T::NAME,
172 "Pruning limit reached"
173 );
174 break false
175 }
176
177 let Some(res) = walker.next() else { break true };
178 let row = res?;
179
180 walker.delete_current_duplicates()?;
181 limiter.increment_deleted_entries_count();
182 delete_callback(row);
183 };
184
185 debug!(
186 target: "providers::db",
187 table=?T::NAME,
188 cursor_current=?cursor.current(),
189 "done walking",
190 );
191
192 let ending_entries = self.entries::<T>()?;
193
194 Ok((starting_entries - ending_entries, done))
195 }
196
197 #[allow(dead_code)]
202 fn prune_dupsort_key_entries<T: DupSort>(
203 &self,
204 key: T::Key,
205 limiter: &mut PruneLimiter,
206 ) -> Result<(usize, bool), DatabaseError> {
207 let mut cursor = self.cursor_dup_write::<T>()?;
208 let mut entry = cursor.seek_exact(key)?;
209
210 let mut deleted_entries = 0;
211
212 while entry.is_some() && !limiter.is_limit_reached() {
213 cursor.delete_current()?;
214 limiter.increment_deleted_entries_count();
215 deleted_entries += 1;
216 entry = cursor.next_dup()?;
217 }
218
219 let done = entry.is_none();
221 if !done {
222 debug!(
223 target: "providers::db",
224 ?limiter,
225 deleted_entries_limit = %limiter.is_deleted_entries_limit_reached(),
226 time_limit = %limiter.is_time_limit_reached(),
227 table = %T::NAME,
228 "Pruning limit reached"
229 );
230 }
231
232 Ok((deleted_entries, done))
233 }
234}
235
236impl<Tx> DbTxPruneExt for Tx where Tx: DbTxMut + DbTx {}
237
238#[cfg(test)]
239mod tests {
240 use super::DbTxPruneExt;
241 use crate::PruneLimiter;
242 use alloy_primitives::{B256, U256};
243 use reth_db_api::{tables, transaction::DbTxMut};
244 use reth_primitives_traits::{SignerRecoverable, StorageEntry};
245 use reth_provider::{DBProvider, DatabaseProviderFactory};
246 use reth_stages::test_utils::{StorageKind, TestStageDB};
247 use reth_testing_utils::generators::{self, random_block_range, BlockRangeParams};
248 use std::{
249 sync::{
250 atomic::{AtomicUsize, Ordering},
251 Arc,
252 },
253 time::Duration,
254 };
255
256 struct CountingIter {
257 data: Vec<u64>,
258 calls: Arc<AtomicUsize>,
259 }
260
261 impl CountingIter {
262 fn new(data: Vec<u64>, calls: Arc<AtomicUsize>) -> Self {
263 Self { data, calls }
264 }
265 }
266
267 struct CountingIntoIter {
268 inner: std::vec::IntoIter<u64>,
269 calls: Arc<AtomicUsize>,
270 }
271
272 impl Iterator for CountingIntoIter {
273 type Item = u64;
274 fn next(&mut self) -> Option<Self::Item> {
275 let res = self.inner.next();
276 self.calls.fetch_add(1, Ordering::SeqCst);
277 res
278 }
279 }
280
281 impl IntoIterator for CountingIter {
282 type Item = u64;
283 type IntoIter = CountingIntoIter;
284 fn into_iter(self) -> Self::IntoIter {
285 CountingIntoIter { inner: self.data.into_iter(), calls: self.calls }
286 }
287 }
288
289 fn storage_entry(slot_byte: u8) -> StorageEntry {
290 StorageEntry { key: B256::with_last_byte(slot_byte), value: U256::from(slot_byte) }
291 }
292
293 fn insert_hashed_storages(db: &TestStageDB, rows: impl IntoIterator<Item = (u8, u8)>) {
294 let provider = db.factory.database_provider_rw().unwrap();
295 for (address_byte, slot_byte) in rows {
296 provider
297 .tx_ref()
298 .put::<tables::HashedStorages>(
299 B256::with_last_byte(address_byte),
300 storage_entry(slot_byte),
301 )
302 .expect("insert hashed storage");
303 }
304 provider.commit().expect("commit");
305 }
306
307 fn hashed_storage_slots(db: &TestStageDB, address_byte: u8) -> Vec<B256> {
308 db.table::<tables::HashedStorages>()
309 .unwrap()
310 .into_iter()
311 .filter_map(|(key, entry)| {
312 (key == B256::with_last_byte(address_byte)).then_some(entry.key)
313 })
314 .collect()
315 }
316
317 fn prune_hashed_storage_key(
318 db: &TestStageDB,
319 address_byte: u8,
320 limiter: &mut PruneLimiter,
321 ) -> (usize, bool) {
322 let provider = db.factory.database_provider_rw().unwrap();
323 let result = provider
324 .tx_ref()
325 .prune_dupsort_key_entries::<tables::HashedStorages>(
326 B256::with_last_byte(address_byte),
327 limiter,
328 )
329 .expect("prune hashed storages");
330 provider.commit().expect("commit");
331 result
332 }
333
334 #[test]
335 fn prune_table_with_iterator_early_exit_does_not_overconsume() {
336 let db = TestStageDB::default();
337 let mut rng = generators::rng();
338
339 let blocks = random_block_range(
340 &mut rng,
341 1..=3,
342 BlockRangeParams {
343 parent: Some(alloy_primitives::B256::ZERO),
344 tx_count: 2..3,
345 ..Default::default()
346 },
347 );
348 db.insert_blocks(blocks.iter(), StorageKind::Database(None)).expect("insert blocks");
349
350 let mut tx_senders = Vec::new();
351 for block in &blocks {
352 tx_senders.reserve_exact(block.transaction_count());
353 for transaction in &block.body().transactions {
354 tx_senders.push((
355 tx_senders.len() as u64,
356 transaction.recover_signer().expect("recover signer"),
357 ));
358 }
359 }
360 let total = tx_senders.len();
361 db.insert_transaction_senders(tx_senders).expect("insert transaction senders");
362
363 let provider = db.factory.database_provider_rw().unwrap();
364
365 let calls = Arc::new(AtomicUsize::new(0));
366 let keys: Vec<u64> = (0..total as u64).collect();
367 let counting_iter = CountingIter::new(keys, calls.clone());
368
369 let mut limiter = PruneLimiter::default().set_deleted_entries_limit(2);
370
371 let (pruned, done) = provider
372 .tx_ref()
373 .prune_table_with_iterator::<tables::TransactionSenders>(
374 counting_iter,
375 &mut limiter,
376 |_| {},
377 )
378 .expect("prune");
379
380 assert_eq!(pruned, 2);
381 assert!(!done);
382 assert_eq!(calls.load(Ordering::SeqCst), pruned + 1);
383
384 provider.commit().expect("commit");
385 assert_eq!(db.table::<tables::TransactionSenders>().unwrap().len(), total - 2);
386 }
387
388 #[test]
389 fn prune_table_with_iterator_consumes_to_end_reports_done() {
390 let db = TestStageDB::default();
391 let mut rng = generators::rng();
392
393 let blocks = random_block_range(
394 &mut rng,
395 1..=2,
396 BlockRangeParams {
397 parent: Some(alloy_primitives::B256::ZERO),
398 tx_count: 1..2,
399 ..Default::default()
400 },
401 );
402 db.insert_blocks(blocks.iter(), StorageKind::Database(None)).expect("insert blocks");
403
404 let mut tx_senders = Vec::new();
405 for block in &blocks {
406 for transaction in &block.body().transactions {
407 tx_senders.push((
408 tx_senders.len() as u64,
409 transaction.recover_signer().expect("recover signer"),
410 ));
411 }
412 }
413 let total = tx_senders.len();
414 db.insert_transaction_senders(tx_senders).expect("insert transaction senders");
415
416 let provider = db.factory.database_provider_rw().unwrap();
417
418 let calls = Arc::new(AtomicUsize::new(0));
419 let keys: Vec<u64> = (0..total as u64).collect();
420 let counting_iter = CountingIter::new(keys, calls.clone());
421
422 let mut limiter = PruneLimiter::default().set_deleted_entries_limit(usize::MAX);
423
424 let (pruned, done) = provider
425 .tx_ref()
426 .prune_table_with_iterator::<tables::TransactionSenders>(
427 counting_iter,
428 &mut limiter,
429 |_| {},
430 )
431 .expect("prune");
432
433 assert_eq!(pruned, total);
434 assert!(done);
435 assert_eq!(calls.load(Ordering::SeqCst), total + 1);
436
437 provider.commit().expect("commit");
438 assert_eq!(db.table::<tables::TransactionSenders>().unwrap().len(), 0);
439 }
440
441 #[test]
442 fn prune_dupsort_key_entries_resumes_with_deleted_entries_budget() {
443 let db = TestStageDB::default();
444 insert_hashed_storages(&db, (0..5).map(|slot| (1, slot)));
445 insert_hashed_storages(&db, (10..12).map(|slot| (2, slot)));
446
447 let mut total_deleted = 0;
448
449 let mut limiter = PruneLimiter::default().set_deleted_entries_limit(2);
450 let (deleted, done) = prune_hashed_storage_key(&db, 1, &mut limiter);
451 total_deleted += deleted;
452 assert_eq!((deleted, done), (2, false));
453 assert_eq!(
454 hashed_storage_slots(&db, 1),
455 vec![B256::with_last_byte(2), B256::with_last_byte(3), B256::with_last_byte(4),]
456 );
457 assert_eq!(
458 hashed_storage_slots(&db, 2),
459 vec![B256::with_last_byte(10), B256::with_last_byte(11)]
460 );
461
462 let mut limiter = PruneLimiter::default().set_deleted_entries_limit(2);
463 let (deleted, done) = prune_hashed_storage_key(&db, 1, &mut limiter);
464 total_deleted += deleted;
465 assert_eq!((deleted, done), (2, false));
466 assert_eq!(hashed_storage_slots(&db, 1), vec![B256::with_last_byte(4)]);
467 assert_eq!(
468 hashed_storage_slots(&db, 2),
469 vec![B256::with_last_byte(10), B256::with_last_byte(11)]
470 );
471
472 let mut limiter = PruneLimiter::default().set_deleted_entries_limit(2);
473 let (deleted, done) = prune_hashed_storage_key(&db, 1, &mut limiter);
474 total_deleted += deleted;
475 assert_eq!((deleted, done), (1, true));
476 assert!(hashed_storage_slots(&db, 1).is_empty());
477 assert_eq!(
478 hashed_storage_slots(&db, 2),
479 vec![B256::with_last_byte(10), B256::with_last_byte(11)]
480 );
481 assert_eq!(total_deleted, 5);
482 }
483
484 #[test]
485 fn prune_dupsort_key_entries_stops_on_time_limit() {
486 let db = TestStageDB::default();
487 insert_hashed_storages(&db, (0..3).map(|slot| (1, slot)));
488
489 let mut limiter = PruneLimiter::default().set_time_limit(Duration::from_nanos(1));
490 std::thread::sleep(Duration::from_millis(1));
491
492 let (deleted, done) = prune_hashed_storage_key(&db, 1, &mut limiter);
493
494 assert_eq!((deleted, done), (0, false));
495 assert!(limiter.is_time_limit_reached());
496 assert_eq!(
497 hashed_storage_slots(&db, 1),
498 vec![B256::with_last_byte(0), B256::with_last_byte(1), B256::with_last_byte(2)]
499 );
500 }
501
502 #[test]
503 fn prune_dupsort_key_entries_missing_key_is_done() {
504 let db = TestStageDB::default();
505 insert_hashed_storages(&db, (0..3).map(|slot| (1, slot)));
506
507 let mut limiter = PruneLimiter::default().set_deleted_entries_limit(2);
508 let result = prune_hashed_storage_key(&db, 2, &mut limiter);
509
510 assert_eq!(result, (0, true));
511 assert_eq!(
512 hashed_storage_slots(&db, 1),
513 vec![B256::with_last_byte(0), B256::with_last_byte(1), B256::with_last_byte(2)]
514 );
515 }
516
517 #[test]
518 fn prune_dupsort_key_entries_leaves_other_keys_untouched() {
519 let db = TestStageDB::default();
520 insert_hashed_storages(&db, (0..3).map(|slot| (1, slot)));
521 insert_hashed_storages(&db, (10..13).map(|slot| (2, slot)));
522
523 let mut limiter = PruneLimiter::default();
524 let result = prune_hashed_storage_key(&db, 1, &mut limiter);
525
526 assert_eq!(result, (3, true));
527 assert!(hashed_storage_slots(&db, 1).is_empty());
528 assert_eq!(
529 hashed_storage_slots(&db, 2),
530 vec![B256::with_last_byte(10), B256::with_last_byte(11), B256::with_last_byte(12),]
531 );
532 }
533}