Skip to content

Commit

Permalink
feat: implement a resource pool
Browse files Browse the repository at this point in the history
  • Loading branch information
jpraynaud committed May 23, 2024
1 parent 5c750d9 commit 259f3f9
Show file tree
Hide file tree
Showing 4 changed files with 357 additions and 20 deletions.
10 changes: 9 additions & 1 deletion mithril-aggregator/src/dependency_injection/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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))
}

Expand Down
50 changes: 31 additions & 19 deletions mithril-aggregator/src/services/prover.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand All @@ -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
Expand All @@ -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 {
Expand All @@ -74,7 +71,7 @@ impl MithrilProverService {
Self {
transaction_retriever,
block_range_root_retriever,
mk_map_cache: Mutex::new(None),
mk_map_cache: ResourcePool::default(),
}
}

Expand Down Expand Up @@ -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
Expand All @@ -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 {
Expand All @@ -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())
Expand All @@ -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(())
}
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions mithril-common/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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! {
Expand Down
Loading

0 comments on commit 259f3f9

Please sign in to comment.