Skip to main content

reth_prune/
db_ext.rs

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/// Result of a single prune step in [`DbTxPruneExt::prune_table_with_range_step`].
12#[derive(Debug, Clone, Copy)]
13pub(crate) struct PruneStepResult {
14    /// `true` if the walker is finished, `false` if it may have more data to prune.
15    done: bool,
16    /// `true` if the current entry was deleted, `false` if it was skipped.
17    deleted: bool,
18}
19
20pub(crate) trait DbTxPruneExt: DbTxMut + DbTx {
21    /// Clear the entire table in a single operation.
22    ///
23    /// This is much faster than iterating entry-by-entry for `PruneMode::Full`.
24    /// Returns the number of entries that were in the table.
25    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    /// Prune the table for the specified pre-sorted key iterator.
32    ///
33    /// Returns number of rows pruned.
34    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    /// Prune the table for the specified key range.
74    ///
75    /// Returns number of rows pruned.
76    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            // check for time out must be done in this scope since it's not done in
90            // `prune_table_with_range_step`
91            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    /// Steps once with the given walker and prunes the entry in the table.
123    ///
124    /// CAUTION: Pruner limits are not checked. This allows for a clean exit of a prune run that's
125    /// pruning different tables concurrently, by letting them step to the same height before
126    /// timing out.
127    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    /// Prune a DUPSORT table for the specified key range.
151    ///
152    /// Returns number of rows pruned.
153    #[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    /// Prune duplicate entries for a single DUPSORT key.
198    ///
199    /// Returns the number of rows pruned and whether all duplicate entries for the key were
200    /// deleted.
201    #[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        // an entry remaining means the loop stopped because a limit was reached
220        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}