feat: implement a resource pool

Author: jpraynaud

mithril-aggregator/src/dependency_injection/builder.rs

@@ -24,7 +24,7 @@ use mithril_common::{
         CardanoImmutableDigester, DumbImmutableFileObserver, ImmutableDigester,
         ImmutableFileObserver, ImmutableFileSystemObserver,
     },
-    entities::{CertificatePending, CompressionAlgorithm, Epoch},
+    entities::{CardanoDbBeacon, CertificatePending, CompressionAlgorithm, Epoch},
     era::{
         adapters::{EraReaderAdapterBuilder, EraReaderDummyAdapter},
         EraChecker, EraMarker, EraReader, EraReaderAdapter, SupportedEra,
@@ -1369,6 +1369,14 @@ impl DependenciesBuilder {
         let block_range_root_retriever = self.get_transaction_repository().await?;
         let service = MithrilProverService::new(transaction_retriever, block_range_root_retriever);
 
+        service
+            .compute_cache(&CardanoDbBeacon::new(
+                self.configuration.get_network()?.to_string(),
+                0,
+                u64::MAX,
+            ))
+            .await?;
+
         Ok(Arc::new(service))
     }
 

mithril-aggregator/src/services/prover.rs

@@ -3,15 +3,16 @@ use rayon::prelude::*;
 use std::{
     collections::{BTreeMap, BTreeSet, HashMap},
     sync::Arc,
+    time::Duration,
 };
-use tokio::sync::Mutex;
 
 use mithril_common::{
     crypto_helper::{MKMap, MKMapNode, MKTree},
     entities::{
         BlockRange, CardanoDbBeacon, CardanoTransaction, CardanoTransactionsSetProof,
         TransactionHash,
     },
+    resource_pool::ResourcePool,
     signable_builder::BlockRangeRootRetriever,
     StdResult,
 };
@@ -28,14 +29,10 @@ pub trait ProverService: Sync + Send {
     ) -> StdResult<Vec<CardanoTransactionsSetProof>>;
 
     /// Compute the cache
-    async fn compute_cache(&self, _up_to: &CardanoDbBeacon) -> StdResult<()> {
-        Ok(())
-    }
+    async fn compute_cache(&self, _up_to: &CardanoDbBeacon) -> StdResult<()>;
 
     /// Clear the cache
-    async fn clear_cache(&self) -> StdResult<()> {
-        Ok(())
-    }
+    async fn clear_cache(&self) -> StdResult<()>;
 }
 
 /// Transactions retriever
@@ -62,7 +59,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>>>>,
+    mk_map_cache: ResourcePool<MKMap<BlockRange, MKMapNode<BlockRange>>>,
 }
 
 impl MithrilProverService {
@@ -74,7 +71,7 @@ impl MithrilProverService {
         Self {
             transaction_retriever,
             block_range_root_retriever,
-            mk_map_cache: Mutex::new(None),
+            mk_map_cache: ResourcePool::default(),
         }
     }
 
@@ -119,7 +116,7 @@ impl MithrilProverService {
 impl ProverService for MithrilProverService {
     async fn compute_transactions_proofs(
         &self,
-        up_to: &CardanoDbBeacon,
+        _up_to: &CardanoDbBeacon,
         transaction_hashes: &[TransactionHash],
     ) -> StdResult<Vec<CardanoTransactionsSetProof>> {
         // 1 - Compute the set of block ranges with transactions to prove
@@ -139,9 +136,9 @@ impl ProverService for MithrilProverService {
         let mk_trees = BTreeMap::from_iter(mk_trees?);
 
         // 3 - Compute block range roots Merkle map
-        self.compute_cache(up_to).await?;
-        let mut mk_map = self.mk_map_cache.lock().await;
-        let mk_map = mk_map.as_mut().unwrap();
+        let mut mk_map = self
+            .mk_map_cache
+            .acquire_resource(Duration::from_millis(1000))?;
 
         // 4 - Enrich the Merkle map with the block ranges Merkle trees
         for (block_range, mk_tree) in mk_trees {
@@ -150,6 +147,9 @@ impl ProverService for MithrilProverService {
 
         // 5 - Compute the proof for all transactions
         if let Ok(mk_proof) = mk_map.compute_proof(transaction_hashes) {
+            self.mk_map_cache
+                .return_resource(mk_map.into_inner(), mk_map.discriminant())?;
+
             let transaction_hashes_certified: Vec<TransactionHash> = transaction_hashes
                 .iter()
                 .filter(|hash| mk_proof.contains(&hash.as_str().into()).is_ok())
@@ -166,22 +166,30 @@ impl ProverService for MithrilProverService {
     }
 
     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");
+        if self.mk_map_cache.count()? == 0 {
+            println!("Computing Merkle map cache 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);
+            let discriminant_new = self.mk_map_cache.discriminant()? + 1;
+            self.mk_map_cache.set_discriminant(discriminant_new)?;
+            for i in 0..10 {
+                println!("Computing Merkle map cache from block range roots: {}", i);
+                self.mk_map_cache
+                    .return_resource(mk_map_cache.clone(), discriminant_new)?;
+            }
+            self.mk_map_cache
+                .return_resource(mk_map_cache, discriminant_new)?;
+            println!("Done computing Merkle map cache from block range roots");
         }
 
         Ok(())
     }
 
     async fn clear_cache(&self) -> StdResult<()> {
-        let mut mk_map = self.mk_map_cache.lock().await;
-        mk_map.take();
+        self.mk_map_cache.drain();
 
         Ok(())
     }
@@ -387,6 +395,7 @@ mod tests {
                     });
             },
         );
+        prover.compute_cache(&test_data.beacon).await.unwrap();
 
         let transactions_set_proof = prover
             .compute_transactions_proofs(&test_data.beacon, &test_data.transaction_hashes_to_prove)
@@ -440,6 +449,7 @@ mod tests {
                     });
             },
         );
+        prover.compute_cache(&test_data.beacon).await.unwrap();
 
         let transactions_set_proof = prover
             .compute_transactions_proofs(&test_data.beacon, &test_data.transaction_hashes_to_prove)
@@ -496,6 +506,7 @@ mod tests {
                     });
             },
         );
+        prover.compute_cache(&test_data.beacon).await.unwrap();
 
         let transactions_set_proof = prover
             .compute_transactions_proofs(&test_data.beacon, &test_data.transaction_hashes_to_prove)
@@ -533,6 +544,7 @@ mod tests {
                     .return_once(|_| MKMap::new(&[]));
             },
         );
+        prover.compute_cache(&test_data.beacon).await.unwrap();
 
         prover
             .compute_transactions_proofs(&test_data.beacon, &test_data.transaction_hashes_to_prove)

mithril-common/src/lib.rs

@@ -59,6 +59,7 @@ pub mod entities;
 pub mod era;
 pub mod messages;
 pub mod protocol;
+pub mod resource_pool;
 pub mod signable_builder;
 
 cfg_test_tools! {