diff --git a/crates/prune/prune/src/db_ext.rs b/crates/prune/prune/src/db_ext.rs index a14127af20..143cb5e277 100644 --- a/crates/prune/prune/src/db_ext.rs +++ b/crates/prune/prune/src/db_ext.rs @@ -1,12 +1,12 @@ use std::{fmt::Debug, ops::RangeBounds}; +use crate::PruneLimiter; use reth_db::{ cursor::{DbCursorRO, DbCursorRW, RangeWalker}, table::{Table, TableRow}, transaction::DbTxMut, DatabaseError, }; -use reth_prune_types::PruneLimiter; use tracing::debug; pub(crate) trait DbTxPruneExt: DbTxMut { diff --git a/crates/prune/prune/src/lib.rs b/crates/prune/prune/src/lib.rs index e6bcbe5e81..ef3ee0de2d 100644 --- a/crates/prune/prune/src/lib.rs +++ b/crates/prune/prune/src/lib.rs @@ -12,6 +12,7 @@ mod builder; mod db_ext; mod error; +mod limiter; mod metrics; mod pruner; pub mod segments; @@ -19,6 +20,7 @@ pub mod segments; use crate::metrics::Metrics; pub use builder::PrunerBuilder; pub use error::PrunerError; +pub use limiter::PruneLimiter; pub use pruner::{Pruner, PrunerResult, PrunerWithFactory, PrunerWithResult}; // Re-export prune types diff --git a/crates/prune/types/src/limiter.rs b/crates/prune/prune/src/limiter.rs similarity index 94% rename from crates/prune/types/src/limiter.rs rename to crates/prune/prune/src/limiter.rs index d555db2573..654eed04f2 100644 --- a/crates/prune/types/src/limiter.rs +++ b/crates/prune/prune/src/limiter.rs @@ -1,3 +1,4 @@ +use reth_prune_types::{PruneInterruptReason, PruneProgress}; use std::{ num::NonZeroUsize, time::{Duration, Instant}, @@ -119,6 +120,30 @@ impl PruneLimiter { pub fn is_limit_reached(&self) -> bool { self.is_deleted_entries_limit_reached() || self.is_time_limit_reached() } + + /// Creates new [`PruneInterruptReason`] based on the limiter's state. + pub fn interrupt_reason(&self) -> PruneInterruptReason { + if self.is_time_limit_reached() { + PruneInterruptReason::Timeout + } else if self.is_deleted_entries_limit_reached() { + PruneInterruptReason::DeletedEntriesLimitReached + } else { + PruneInterruptReason::Unknown + } + } + + /// Creates new [`PruneProgress`]. + /// + /// If `done == true`, returns [`PruneProgress::Finished`], otherwise + /// [`PruneProgress::HasMoreData`] is returned with [`PruneInterruptReason`] according to the + /// limiter's state. + pub fn progress(&self, done: bool) -> PruneProgress { + if done { + PruneProgress::Finished + } else { + PruneProgress::HasMoreData(self.interrupt_reason()) + } + } } #[cfg(test)] diff --git a/crates/prune/prune/src/pruner.rs b/crates/prune/prune/src/pruner.rs index 0ad149bb65..2344578bd0 100644 --- a/crates/prune/prune/src/pruner.rs +++ b/crates/prune/prune/src/pruner.rs @@ -2,14 +2,14 @@ use crate::{ segments::{PruneInput, Segment}, - Metrics, PrunerError, PrunerEvent, + Metrics, PruneLimiter, PrunerError, PrunerEvent, }; use alloy_primitives::BlockNumber; use reth_exex_types::FinishedExExHeight; use reth_provider::{ DBProvider, DatabaseProviderFactory, PruneCheckpointReader, PruneCheckpointWriter, }; -use reth_prune_types::{PruneLimiter, PruneProgress, PrunedSegmentInfo, PrunerOutput}; +use reth_prune_types::{PruneProgress, PrunedSegmentInfo, PrunerOutput}; use reth_tokio_util::{EventSender, EventStream}; use std::time::{Duration, Instant}; use tokio::sync::watch; diff --git a/crates/prune/prune/src/segments/mod.rs b/crates/prune/prune/src/segments/mod.rs index e828512fa8..ae18bcb3c6 100644 --- a/crates/prune/prune/src/segments/mod.rs +++ b/crates/prune/prune/src/segments/mod.rs @@ -3,12 +3,10 @@ mod set; mod static_file; mod user; -use crate::PrunerError; +use crate::{PruneLimiter, PrunerError}; use alloy_primitives::{BlockNumber, TxNumber}; use reth_provider::{errors::provider::ProviderResult, BlockReader, PruneCheckpointWriter}; -use reth_prune_types::{ - PruneCheckpoint, PruneLimiter, PruneMode, PrunePurpose, PruneSegment, SegmentOutput, -}; +use reth_prune_types::{PruneCheckpoint, PruneMode, PrunePurpose, PruneSegment, SegmentOutput}; pub use set::SegmentSet; pub use static_file::{ Headers as StaticFileHeaders, Receipts as StaticFileReceipts, diff --git a/crates/prune/prune/src/segments/receipts.rs b/crates/prune/prune/src/segments/receipts.rs index a365738a77..dbea32c47f 100644 --- a/crates/prune/prune/src/segments/receipts.rs +++ b/crates/prune/prune/src/segments/receipts.rs @@ -12,9 +12,7 @@ use reth_provider::{ errors::provider::ProviderResult, BlockReader, DBProvider, NodePrimitivesProvider, PruneCheckpointWriter, TransactionsProvider, }; -use reth_prune_types::{ - PruneCheckpoint, PruneProgress, PruneSegment, SegmentOutput, SegmentOutputCheckpoint, -}; +use reth_prune_types::{PruneCheckpoint, PruneSegment, SegmentOutput, SegmentOutputCheckpoint}; use tracing::trace; pub(crate) fn prune( @@ -56,7 +54,7 @@ where // so we could finish pruning its receipts on the next run. .checked_sub(if done { 0 } else { 1 }); - let progress = PruneProgress::new(done, &limiter); + let progress = limiter.progress(done); Ok(SegmentOutput { progress, @@ -83,7 +81,7 @@ pub(crate) fn save_checkpoint( #[cfg(test)] mod tests { - use crate::segments::{PruneInput, SegmentOutput}; + use crate::segments::{PruneInput, PruneLimiter, SegmentOutput}; use alloy_primitives::{BlockNumber, TxNumber, B256}; use assert_matches::assert_matches; use itertools::{ @@ -93,7 +91,7 @@ mod tests { use reth_db::tables; use reth_provider::{DatabaseProviderFactory, PruneCheckpointReader}; use reth_prune_types::{ - PruneCheckpoint, PruneInterruptReason, PruneLimiter, PruneMode, PruneProgress, PruneSegment, + PruneCheckpoint, PruneInterruptReason, PruneMode, PruneProgress, PruneSegment, }; use reth_stages::test_utils::{StorageKind, TestStageDB}; use reth_testing_utils::generators::{ diff --git a/crates/prune/prune/src/segments/static_file/headers.rs b/crates/prune/prune/src/segments/static_file/headers.rs index 5cd6f62643..7d100f4e28 100644 --- a/crates/prune/prune/src/segments/static_file/headers.rs +++ b/crates/prune/prune/src/segments/static_file/headers.rs @@ -3,7 +3,7 @@ use std::num::NonZeroUsize; use crate::{ db_ext::DbTxPruneExt, segments::{PruneInput, Segment}, - PrunerError, + PruneLimiter, PrunerError, }; use alloy_primitives::BlockNumber; use itertools::Itertools; @@ -14,8 +14,7 @@ use reth_db::{ }; use reth_provider::{providers::StaticFileProvider, DBProvider, StaticFileProviderFactory}; use reth_prune_types::{ - PruneLimiter, PruneMode, PruneProgress, PrunePurpose, PruneSegment, SegmentOutput, - SegmentOutputCheckpoint, + PruneMode, PrunePurpose, PruneSegment, SegmentOutput, SegmentOutputCheckpoint, }; use reth_static_file_types::StaticFileSegment; use tracing::trace; @@ -92,7 +91,7 @@ impl> Segment Self { - if limiter.is_time_limit_reached() { - Self::Timeout - } else if limiter.is_deleted_entries_limit_reached() { - Self::DeletedEntriesLimitReached - } else { - Self::Unknown - } - } - /// Returns `true` if the reason is timeout. pub const fn is_timeout(&self) -> bool { matches!(self, Self::Timeout) @@ -124,19 +113,6 @@ impl PruneInterruptReason { } impl PruneProgress { - /// Creates new [`PruneProgress`]. - /// - /// If `done == true`, returns [`PruneProgress::Finished`], otherwise - /// [`PruneProgress::HasMoreData`] is returned with [`PruneInterruptReason`] according to the - /// passed limiter. - pub fn new(done: bool, limiter: &PruneLimiter) -> Self { - if done { - Self::Finished - } else { - Self::HasMoreData(PruneInterruptReason::new(limiter)) - } - } - /// Returns `true` if prune run is finished. pub const fn is_finished(&self) -> bool { matches!(self, Self::Finished)