feat: implement a resource pool for transaction prover

jpraynaud · jpraynaud · commit 81f3bc58c05c · 2024-05-23T11:52:25.000+02:00
diff --git a/mithril-aggregator/src/dependency_injection/builder.rs b/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,
@@ -1365,9 +1365,24 @@ impl DependenciesBuilder {
 
     /// build Prover service
     pub async fn build_prover_service(&mut self) -> Result<Arc<dyn ProverService>> {
+        let mk_map_pool_size = 30; // TODO: make this configurable
         let transaction_retriever = self.get_transaction_repository().await?;
         let block_range_root_retriever = self.get_transaction_repository().await?;
-        let service = MithrilProverService::new(transaction_retriever, block_range_root_retriever);
+        let logger = self.get_logger().await?;
+        let service = MithrilProverService::new(
+            transaction_retriever,
+            block_range_root_retriever,
+            mk_map_pool_size,
+            logger,
+        );
+
+        service
+            .compute_cache(&CardanoDbBeacon::new(
+                self.configuration.get_network()?.to_string(),
+                0,
+                u64::MAX,
+            ))
+            .await?;
 
         Ok(Arc::new(service))
     }
diff --git a/mithril-aggregator/src/runtime/runner.rs b/mithril-aggregator/src/runtime/runner.rs
@@ -399,6 +399,13 @@ impl AggregatorRunnerTrait for AggregatorRunner {
                 )
             })?;
 
+        if let SignedEntityType::CardanoTransactions(up_to_beacon) = signed_entity_type {
+            self.dependencies
+                .prover_service
+                .compute_cache(up_to_beacon)
+                .await?;
+        }
+
         Ok(())
     }
 
diff --git a/mithril-aggregator/src/services/prover.rs b/mithril-aggregator/src/services/prover.rs
@@ -1,16 +1,19 @@
+use async_trait::async_trait;
+use rayon::prelude::*;
+use slog::{debug, info, Logger};
 use std::{
     collections::{BTreeMap, BTreeSet, HashMap},
     sync::Arc,
+    time::Duration,
 };
 
-use async_trait::async_trait;
-
 use mithril_common::{
-    crypto_helper::MKTree,
+    crypto_helper::{MKMap, MKMapNode, MKTree},
     entities::{
         BlockRange, CardanoDbBeacon, CardanoTransaction, CardanoTransactionsSetProof,
         TransactionHash,
     },
+    resource_pool::ResourcePool,
     signable_builder::BlockRangeRootRetriever,
     StdResult,
 };
@@ -25,6 +28,9 @@ 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<()>;
 }
 
 /// Transactions retriever
@@ -51,17 +57,23 @@ pub trait TransactionsRetriever: Sync + Send {
 pub struct MithrilProverService {
     transaction_retriever: Arc<dyn TransactionsRetriever>,
     block_range_root_retriever: Arc<dyn BlockRangeRootRetriever>,
+    mk_map_pool: ResourcePool<MKMap<BlockRange, MKMapNode<BlockRange>>>,
+    logger: Logger,
 }
 
 impl MithrilProverService {
     /// Create a new Mithril prover
     pub fn new(
         transaction_retriever: Arc<dyn TransactionsRetriever>,
         block_range_root_retriever: Arc<dyn BlockRangeRootRetriever>,
+        mk_map_pool_size: usize,
+        logger: Logger,
     ) -> Self {
         Self {
             transaction_retriever,
             block_range_root_retriever,
+            mk_map_pool: ResourcePool::new(mk_map_pool_size, vec![]),
+            logger,
         }
     }
 
@@ -106,7 +118,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
@@ -124,9 +136,8 @@ 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?;
+            .mk_map_pool
+            .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 {
@@ -135,6 +146,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_pool
+                .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())
@@ -149,6 +163,35 @@ impl ProverService for MithrilProverService {
             Ok(vec![])
         }
     }
+
+    async fn compute_cache(&self, up_to: &CardanoDbBeacon) -> StdResult<()> {
+        let pool_size = self.mk_map_pool.size();
+        info!(
+            self.logger,
+            "Prover starts computing the Merkle map pool pool resource of size {pool_size}"
+        );
+        self.mk_map_pool.drain();
+        let mk_map_cache = self
+            .block_range_root_retriever
+            .compute_merkle_map_from_block_range_roots(up_to.immutable_file_number)
+            .await?;
+        let discriminant_new = self.mk_map_pool.discriminant()? + 1;
+        self.mk_map_pool.set_discriminant(discriminant_new)?;
+        for i in 1..=pool_size {
+            debug!(
+                self.logger,
+                "Prover is computing the Merkle map pool pool resource {i}/{pool_size}"
+            );
+            self.mk_map_pool
+                .return_resource(mk_map_cache.clone(), discriminant_new)?;
+        }
+        info!(
+            self.logger,
+            "Prover completed computing the Merkle map pool pool resource of size {pool_size}"
+        );
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -305,10 +348,14 @@ mod tests {
         transaction_retriever_mock_config(&mut transaction_retriever);
         let mut block_range_root_retriever = MockBlockRangeRootRetrieverImpl::new();
         block_range_root_retriever_mock_config(&mut block_range_root_retriever);
+        let mk_map_pool_size = 1;
+        let logger = slog_scope::logger();
 
         MithrilProverService::new(
             Arc::new(transaction_retriever),
             Arc::new(block_range_root_retriever),
+            mk_map_pool_size,
+            logger,
         )
     }
 
@@ -351,6 +398,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)
@@ -404,6 +452,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)
@@ -460,6 +509,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)
@@ -497,6 +547,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)
diff --git a/mithril-common/src/lib.rs b/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! {
diff --git a/mithril-common/src/resource_pool.rs b/mithril-common/src/resource_pool.rs

Original file line number	Diff line number	Diff line change
`@@ -399,6 +399,13 @@ impl AggregatorRunnerTrait for AggregatorRunner {`
`399`	`399`	`)`
`400`	`400`	`})?;`
`401`	`401`
	`402`	`+ if let SignedEntityType::CardanoTransactions(up_to_beacon) = signed_entity_type {`
	`403`	`+ self.dependencies`
	`404`	`+ .prover_service`
	`405`	`+ .compute_cache(up_to_beacon)`
	`406`	`+ .await?;`
	`407`	`+ }`
	`408`	`+`
`402`	`409`	`Ok(())`
`403`	`410`	`}`
`404`	`411`