mithril_cardano_node_chain/chain_importer/
importer_with_pruner.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use slog::{Logger, debug};
5
6use mithril_common::StdResult;
7use mithril_common::entities::BlockNumber;
8use mithril_common::logging::LoggerExtensions;
9
10use crate::chain_importer::ChainDataImporter;
11
12/// Cardano stored chain data pruner
13#[cfg_attr(test, mockall::automock)]
14#[async_trait]
15pub trait ChainDataPruner: Send + Sync {
16    /// Prune the chain data older than the given number of blocks.
17    async fn prune(&self, number_of_blocks_to_keep: BlockNumber) -> StdResult<()>;
18}
19
20/// A decorator of [ChainDataImporter] that prunes the chain data older than a given number of
21/// blocks after running the import.
22///
23/// If the number of blocks to keep is not provided, no pruning is performed.
24pub struct ChainDataImporterWithPruner {
25    number_of_blocks_to_keep: Option<BlockNumber>,
26    chain_data_pruner: Arc<dyn ChainDataPruner>,
27    wrapped_importer: Arc<dyn ChainDataImporter>,
28    logger: Logger,
29}
30
31impl ChainDataImporterWithPruner {
32    /// Create a new instance of [ChainDataImporterWithPruner].
33    pub fn new(
34        number_of_blocks_to_keep: Option<BlockNumber>,
35        chain_data_pruner: Arc<dyn ChainDataPruner>,
36        wrapped_importer: Arc<dyn ChainDataImporter>,
37        logger: Logger,
38    ) -> Self {
39        Self {
40            number_of_blocks_to_keep,
41            chain_data_pruner,
42            wrapped_importer,
43            logger: logger.new_with_component_name::<Self>(),
44        }
45    }
46}
47
48#[async_trait]
49impl ChainDataImporter for ChainDataImporterWithPruner {
50    async fn import(&self, up_to_beacon: BlockNumber) -> StdResult<()> {
51        self.wrapped_importer.import(up_to_beacon).await?;
52
53        if let Some(number_of_blocks_to_keep) = self.number_of_blocks_to_keep {
54            debug!(
55                self.logger,
56                "Chain data Import finished - Pruning data included in block range roots";
57                "number_of_blocks_to_keep" => *number_of_blocks_to_keep,
58            );
59            self.chain_data_pruner.prune(number_of_blocks_to_keep).await?;
60        }
61
62        Ok(())
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use mockall::predicate::eq;
69
70    use crate::chain_importer::MockChainDataImporter;
71    use crate::test::TestLogger;
72
73    use super::*;
74
75    impl ChainDataImporterWithPruner {
76        pub(crate) fn new_with_mock<P, I>(
77            number_of_blocks_to_keep: Option<BlockNumber>,
78            pruner_mock_config: P,
79            importer_mock_config: I,
80        ) -> Self
81        where
82            P: FnOnce(&mut MockChainDataPruner),
83            I: FnOnce(&mut MockChainDataImporter),
84        {
85            let mut chain_data_pruner = MockChainDataPruner::new();
86            pruner_mock_config(&mut chain_data_pruner);
87            let mut chain_data_importer = MockChainDataImporter::new();
88            importer_mock_config(&mut chain_data_importer);
89
90            Self::new(
91                number_of_blocks_to_keep,
92                Arc::new(chain_data_pruner),
93                Arc::new(chain_data_importer),
94                TestLogger::stdout(),
95            )
96        }
97    }
98
99    #[tokio::test]
100    async fn test_does_not_prune_if_none_is_configured() {
101        let importer = ChainDataImporterWithPruner::new_with_mock(
102            None,
103            |mock| {
104                mock.expect_prune().never();
105            },
106            |mock| {
107                mock.expect_import().once().returning(|_| Ok(()));
108            },
109        );
110
111        importer
112            .import(BlockNumber(100))
113            .await
114            .expect("Import should not fail");
115    }
116
117    #[tokio::test]
118    async fn test_does_prune_if_a_block_number_is_configured() {
119        let expected_block_number = BlockNumber(5);
120        let importer = ChainDataImporterWithPruner::new_with_mock(
121            Some(expected_block_number),
122            |mock| {
123                mock.expect_prune()
124                    .with(eq(expected_block_number))
125                    .once()
126                    .returning(|_| Ok(()));
127            },
128            |mock| {
129                mock.expect_import().once().returning(|_| Ok(()));
130            },
131        );
132
133        importer
134            .import(BlockNumber(100))
135            .await
136            .expect("Import should not fail");
137    }
138}