mithril_cardano_node_chain/chain_importer/
importer_with_pruner.rs1use 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#[cfg_attr(test, mockall::automock)]
14#[async_trait]
15pub trait ChainDataPruner: Send + Sync {
16 async fn prune(&self, number_of_blocks_to_keep: BlockNumber) -> StdResult<()>;
18}
19
20pub 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 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}