Add a decorator of TransactionImporter to add tx pruning capability

Author: Alenar

mithril-signer/src/database/repository/cardano_transaction_repository.rs

@@ -7,7 +7,7 @@ use mithril_common::entities::{BlockNumber, BlockRange, CardanoTransaction, Immu
 use mithril_common::StdResult;
 use mithril_persistence::database::repository::CardanoTransactionRepository;
 
-use crate::TransactionStore;
+use crate::{TransactionPruner, TransactionStore};
 
 #[async_trait]
 impl TransactionStore for CardanoTransactionRepository {
@@ -46,3 +46,10 @@ impl TransactionStore for CardanoTransactionRepository {
         Ok(())
     }
 }
+
+#[async_trait]
+impl TransactionPruner for CardanoTransactionRepository {
+    async fn prune(&self, number_of_blocks_to_keep: BlockNumber) -> StdResult<()> {
+        self.prune_transaction(number_of_blocks_to_keep).await
+    }
+}

mithril-signer/src/lib.rs

@@ -15,6 +15,7 @@ pub mod metrics;
 mod protocol_initializer_store;
 mod runtime;
 mod single_signer;
+mod transactions_importer_with_prune_decorator;
 
 #[cfg(test)]
 pub use aggregator_client::dumb::DumbAggregatorClient;
@@ -28,6 +29,7 @@ pub use metrics::*;
 pub use protocol_initializer_store::{ProtocolInitializerStore, ProtocolInitializerStorer};
 pub use runtime::*;
 pub use single_signer::*;
+pub use transactions_importer_with_prune_decorator::*;
 
 /// HTTP request timeout duration in milliseconds
 const HTTP_REQUEST_TIMEOUT_DURATION: u64 = 30000;

mithril-signer/src/runtime/signer_services.rs

@@ -265,17 +265,23 @@ impl<'a> ServiceBuilder for ProductionServiceBuilder<'a> {
         let transaction_store = Arc::new(CardanoTransactionRepository::new(
             transaction_sqlite_connection,
         ));
-        let transactions_importer = CardanoTransactionsImporter::new(
+        let transactions_importer = Arc::new(CardanoTransactionsImporter::new(
             block_scanner,
             transaction_store.clone(),
             &self.config.db_directory,
             // Rescan the last immutable when importing transactions, it may have been partially imported
             Some(1),
             slog_scope::logger(),
-        );
+        ));
+        // Wrap the transaction importer with decorator to prune the transactions after import
+        let transactions_importer = Arc::new(crate::TransactionsImporterWithPruneDecorator::new(
+            None,
+            transaction_store.clone(),
+            transactions_importer,
+        ));
         let block_range_root_retriever = transaction_store.clone();
         let cardano_transactions_builder = Arc::new(CardanoTransactionsSignableBuilder::new(
-            Arc::new(transactions_importer),
+            transactions_importer,
             block_range_root_retriever,
             slog_scope::logger(),
         ));

mithril-signer/src/transactions_importer_with_prune_decorator.rs

@@ -0,0 +1,130 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+
+use mithril_common::entities::{BlockNumber, ImmutableFileNumber};
+use mithril_common::signable_builder::TransactionsImporter;
+use mithril_common::StdResult;
+
+/// Cardano transactions pruner
+#[cfg_attr(test, mockall::automock)]
+#[async_trait]
+pub trait TransactionPruner: Send + Sync {
+    /// Prune the transactions older than the given number of blocks (based on the block range root
+    /// stored).
+    async fn prune(&self, number_of_blocks_to_keep: BlockNumber) -> StdResult<()>;
+}
+
+/// A decorator of [TransactionsImporter] that prunes the transactions older than a given number of
+/// blocks after running the import.
+///
+/// If the number of blocks to keep is not provided, no pruning is performed.
+pub struct TransactionsImporterWithPruneDecorator {
+    number_of_blocks_to_keep: Option<BlockNumber>,
+    transaction_pruner: Arc<dyn TransactionPruner>,
+    wrapped_importer: Arc<dyn TransactionsImporter>,
+}
+
+impl TransactionsImporterWithPruneDecorator {
+    /// Create a new instance of [TransactionsImporterWithPruneDecorator].
+    pub fn new(
+        number_of_blocks_to_keep: Option<BlockNumber>,
+        transaction_pruner: Arc<dyn TransactionPruner>,
+        wrapped_importer: Arc<dyn TransactionsImporter>,
+    ) -> Self {
+        Self {
+            number_of_blocks_to_keep,
+            transaction_pruner,
+            wrapped_importer,
+        }
+    }
+}
+
+#[async_trait]
+impl TransactionsImporter for TransactionsImporterWithPruneDecorator {
+    async fn import(&self, up_to_beacon: ImmutableFileNumber) -> StdResult<()> {
+        self.wrapped_importer.import(up_to_beacon).await?;
+
+        if let Some(number_of_blocks_to_keep) = self.number_of_blocks_to_keep {
+            self.transaction_pruner
+                .prune(number_of_blocks_to_keep)
+                .await?;
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use mockall::mock;
+    use mockall::predicate::eq;
+
+    use super::*;
+
+    mock! {
+        pub TransactionImporterImpl {}
+
+        #[async_trait]
+        impl TransactionsImporter for TransactionImporterImpl {
+            async fn import(&self, up_to_beacon: ImmutableFileNumber) -> StdResult<()>;
+        }
+    }
+
+    impl TransactionsImporterWithPruneDecorator {
+        pub fn new_with_mock<Ti, Pr>(
+            number_of_blocks_to_keep: Option<BlockNumber>,
+            transaction_pruner_mock_config: Pr,
+            importer_mock_config: Ti,
+        ) -> Self
+        where
+            Pr: FnOnce(&mut MockTransactionPruner),
+            Ti: FnOnce(&mut MockTransactionImporterImpl),
+        {
+            let mut transaction_pruner = MockTransactionPruner::new();
+            transaction_pruner_mock_config(&mut transaction_pruner);
+            let mut transaction_importer = MockTransactionImporterImpl::new();
+            importer_mock_config(&mut transaction_importer);
+
+            Self::new(
+                number_of_blocks_to_keep,
+                Arc::new(transaction_pruner),
+                Arc::new(transaction_importer),
+            )
+        }
+    }
+
+    #[tokio::test]
+    async fn test_does_not_prune_if_none_is_configured() {
+        let importer = TransactionsImporterWithPruneDecorator::new_with_mock(
+            None,
+            |mock| {
+                mock.expect_prune().never();
+            },
+            |mock| {
+                mock.expect_import().once().returning(|_| Ok(()));
+            },
+        );
+
+        importer.import(10).await.expect("Import should not fail");
+    }
+
+    #[tokio::test]
+    async fn test_does_prune_if_a_block_number_is_configured() {
+        let expected_block_number: BlockNumber = 5;
+        let importer = TransactionsImporterWithPruneDecorator::new_with_mock(
+            Some(expected_block_number),
+            |mock| {
+                mock.expect_prune()
+                    .with(eq(expected_block_number))
+                    .once()
+                    .returning(|_| Ok(()));
+            },
+            |mock| {
+                mock.expect_import().once().returning(|_| Ok(()));
+            },
+        );
+
+        importer.import(10).await.expect("Import should not fail");
+    }
+}