reth_prune_types/
mode.rs

1use crate::{segment::PrunePurpose, PruneSegment, PruneSegmentError};
2use alloy_primitives::BlockNumber;
3use reth_codecs::{add_arbitrary_tests, Compact};
4use serde::{Deserialize, Serialize};
5
6/// Prune mode.
7#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, Compact)]
8#[serde(rename_all = "lowercase")]
9#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
10#[add_arbitrary_tests(compact)]
11pub enum PruneMode {
12    /// Prune all blocks.
13    Full,
14    /// Prune blocks before the `head-N` block number. In other words, keep last N + 1 blocks.
15    Distance(u64),
16    /// Prune blocks before the specified block number. The specified block number is not pruned.
17    Before(BlockNumber),
18}
19
20#[cfg(any(test, feature = "test-utils"))]
21impl Default for PruneMode {
22    fn default() -> Self {
23        Self::Full
24    }
25}
26
27impl PruneMode {
28    /// Prune blocks up to the specified block number. The specified block number is also pruned.
29    ///
30    /// This acts as `PruneMode::Before(block_number + 1)`.
31    pub const fn before_inclusive(block_number: BlockNumber) -> Self {
32        Self::Before(block_number + 1)
33    }
34
35    /// Returns block up to which variant pruning needs to be done, inclusive, according to the
36    /// provided tip.
37    pub fn prune_target_block(
38        &self,
39        tip: BlockNumber,
40        segment: PruneSegment,
41        purpose: PrunePurpose,
42    ) -> Result<Option<(BlockNumber, Self)>, PruneSegmentError> {
43        let result = match self {
44            Self::Full if segment.min_blocks(purpose) == 0 => Some((tip, *self)),
45            Self::Distance(distance) if *distance > tip => None, // Nothing to prune yet
46            Self::Distance(distance) if *distance >= segment.min_blocks(purpose) => {
47                Some((tip - distance, *self))
48            }
49            Self::Before(n) if *n == tip + 1 && purpose.is_static_file() => Some((tip, *self)),
50            Self::Before(n) if *n > tip => None, // Nothing to prune yet
51            Self::Before(n) if tip - n >= segment.min_blocks(purpose) => {
52                Some(((*n).saturating_sub(1), *self))
53            }
54            _ => return Err(PruneSegmentError::Configuration(segment)),
55        };
56        Ok(result)
57    }
58
59    /// Check if target block should be pruned according to the provided prune mode and tip.
60    pub const fn should_prune(&self, block: BlockNumber, tip: BlockNumber) -> bool {
61        match self {
62            Self::Full => true,
63            Self::Distance(distance) => {
64                if *distance > tip {
65                    return false
66                }
67                block < tip - *distance
68            }
69            Self::Before(n) => *n > block,
70        }
71    }
72
73    /// Returns true if the prune mode is [`PruneMode::Full`].
74    pub const fn is_full(&self) -> bool {
75        matches!(self, Self::Full)
76    }
77
78    /// Returns true if the prune mode is [`PruneMode::Distance`].
79    pub const fn is_distance(&self) -> bool {
80        matches!(self, Self::Distance(_))
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use crate::{
87        PruneMode, PrunePurpose, PruneSegment, PruneSegmentError, MINIMUM_PRUNING_DISTANCE,
88    };
89    use assert_matches::assert_matches;
90    use serde::Deserialize;
91
92    #[test]
93    fn test_prune_target_block() {
94        let tip = 20000;
95        let segment = PruneSegment::Receipts;
96
97        let tests = vec![
98            // MINIMUM_PRUNING_DISTANCE makes this impossible
99            (PruneMode::Full, Err(PruneSegmentError::Configuration(segment))),
100            // Nothing to prune
101            (PruneMode::Distance(tip + 1), Ok(None)),
102            (
103                PruneMode::Distance(segment.min_blocks(PrunePurpose::User) + 1),
104                Ok(Some(tip - (segment.min_blocks(PrunePurpose::User) + 1))),
105            ),
106            // Nothing to prune
107            (PruneMode::Before(tip + 1), Ok(None)),
108            (
109                PruneMode::Before(tip - MINIMUM_PRUNING_DISTANCE),
110                Ok(Some(tip - MINIMUM_PRUNING_DISTANCE - 1)),
111            ),
112            (
113                PruneMode::Before(tip - MINIMUM_PRUNING_DISTANCE - 1),
114                Ok(Some(tip - MINIMUM_PRUNING_DISTANCE - 2)),
115            ),
116            (PruneMode::Before(tip - 1), Err(PruneSegmentError::Configuration(segment))),
117        ];
118
119        for (index, (mode, expected_result)) in tests.into_iter().enumerate() {
120            assert_eq!(
121                mode.prune_target_block(tip, segment, PrunePurpose::User),
122                expected_result.map(|r| r.map(|b| (b, mode))),
123                "Test {} failed",
124                index + 1,
125            );
126        }
127
128        // Test for a scenario where there are no minimum blocks and Full can be used
129        assert_eq!(
130            PruneMode::Full.prune_target_block(tip, PruneSegment::Transactions, PrunePurpose::User),
131            Ok(Some((tip, PruneMode::Full))),
132        );
133    }
134
135    #[test]
136    fn test_should_prune() {
137        let tip = 20000;
138        let should_prune = true;
139
140        let tests = vec![
141            (PruneMode::Distance(tip + 1), 1, !should_prune),
142            (
143                PruneMode::Distance(MINIMUM_PRUNING_DISTANCE + 1),
144                tip - MINIMUM_PRUNING_DISTANCE - 1,
145                !should_prune,
146            ),
147            (
148                PruneMode::Distance(MINIMUM_PRUNING_DISTANCE + 1),
149                tip - MINIMUM_PRUNING_DISTANCE - 2,
150                should_prune,
151            ),
152            (PruneMode::Before(tip + 1), 1, should_prune),
153            (PruneMode::Before(tip + 1), tip + 1, !should_prune),
154        ];
155
156        for (index, (mode, block, expected_result)) in tests.into_iter().enumerate() {
157            assert_eq!(mode.should_prune(block, tip), expected_result, "Test {} failed", index + 1,);
158        }
159    }
160
161    #[test]
162    fn prune_mode_deserialize() {
163        #[derive(Debug, Deserialize)]
164        struct Config {
165            a: Option<PruneMode>,
166            b: Option<PruneMode>,
167            c: Option<PruneMode>,
168            d: Option<PruneMode>,
169        }
170
171        let toml_str = r#"
172        a = "full"
173        b = { distance = 10 }
174        c = { before = 20 }
175    "#;
176
177        assert_matches!(
178            toml::from_str(toml_str),
179            Ok(Config {
180                a: Some(PruneMode::Full),
181                b: Some(PruneMode::Distance(10)),
182                c: Some(PruneMode::Before(20)),
183                d: None
184            })
185        );
186    }
187}
reth_prune_types/mode.rs

reth_prune_types/
mode.rs
