feat: implement simple cache for transaction prover

Author: jpraynaud

mithril-aggregator/src/runtime/runner.rs

@@ -399,6 +399,14 @@ impl AggregatorRunnerTrait for AggregatorRunner {
                 )
             })?;
 
+        if let SignedEntityType::CardanoTransactions(up_to_beacon) = signed_entity_type {
+            self.dependencies.prover_service.clear_cache().await?;
+            self.dependencies
+                .prover_service
+                .compute_cache(up_to_beacon)
+                .await?;
+        }
+
         Ok(())
     }
 

mithril-aggregator/src/services/prover.rs

@@ -6,14 +6,15 @@ use std::{
 use async_trait::async_trait;
 
 use mithril_common::{
-    crypto_helper::MKTree,
+    crypto_helper::{MKMap, MKMapNode, MKTree},
     entities::{
         BlockRange, CardanoDbBeacon, CardanoTransaction, CardanoTransactionsSetProof,
         TransactionHash,
     },
     signable_builder::BlockRangeRootRetriever,
     StdResult,
 };
+use tokio::sync::Mutex;
 
 /// Prover service is the cryptographic engine in charge of producing cryptographic proofs for transactions
 #[cfg_attr(test, mockall::automock)]
@@ -25,6 +26,16 @@ pub trait ProverService: Sync + Send {
         up_to: &CardanoDbBeacon,
         transaction_hashes: &[TransactionHash],
     ) -> StdResult<Vec<CardanoTransactionsSetProof>>;
+
+    /// Compute the cache
+    async fn compute_cache(&self, _up_to: &CardanoDbBeacon) -> StdResult<()> {
+        Ok(())
+    }
+
+    /// Clear the cache
+    async fn clear_cache(&self) -> StdResult<()> {
+        Ok(())
+    }
 }
 
 /// Transactions retriever
@@ -51,6 +62,7 @@ pub trait TransactionsRetriever: Sync + Send {
 pub struct MithrilProverService {
     transaction_retriever: Arc<dyn TransactionsRetriever>,
     block_range_root_retriever: Arc<dyn BlockRangeRootRetriever>,
+    mk_map_cache: Mutex<Option<MKMap<BlockRange, MKMapNode<BlockRange>>>>,
 }
 
 impl MithrilProverService {
@@ -62,6 +74,7 @@ impl MithrilProverService {
         Self {
             transaction_retriever,
             block_range_root_retriever,
+            mk_map_cache: Mutex::new(None),
         }
     }
 
@@ -123,10 +136,9 @@ impl ProverService for MithrilProverService {
         }
 
         // 3 - Compute block range roots Merkle map
-        let mut mk_map = self
-            .block_range_root_retriever
-            .compute_merkle_map_from_block_range_roots(up_to.immutable_file_number)
-            .await?;
+        self.compute_cache(up_to).await?;
+        let mut mk_map = self.mk_map_cache.lock().await;
+        let mk_map = mk_map.as_mut().unwrap();
 
         // 4 - Enrich the Merkle map with the block ranges Merkle trees
         for (block_range, mk_tree) in mk_trees {
@@ -149,6 +161,27 @@ impl ProverService for MithrilProverService {
             Ok(vec![])
         }
     }
+
+    async fn compute_cache(&self, up_to: &CardanoDbBeacon) -> StdResult<()> {
+        let mut mk_map = self.mk_map_cache.lock().await;
+        if mk_map.is_none() {
+            println!("Computing Merkle map from block range roots");
+            let mk_map_cache = self
+                .block_range_root_retriever
+                .compute_merkle_map_from_block_range_roots(up_to.immutable_file_number)
+                .await?;
+            mk_map.replace(mk_map_cache);
+        }
+
+        Ok(())
+    }
+
+    async fn clear_cache(&self) -> StdResult<()> {
+        let mut mk_map = self.mk_map_cache.lock().await;
+        mk_map.take();
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]