mithril_aggregator/services/
prover.rs1use async_trait::async_trait;
2use rayon::prelude::*;
3use slog::{Logger, debug, info};
4use std::sync::Arc;
5
6use mithril_common::{
7 StdResult,
8 crypto_helper::{MKMap, MKMapNode, MKTreeStorer},
9 entities::{BlockNumber, BlockRange},
10 logging::LoggerExtensions,
11 signable_builder::BlockRangeRootRetriever,
12};
13use mithril_resource_pool::ResourcePool;
14
15use crate::services::TransactionsRetriever;
16
17#[cfg_attr(test, mockall::automock)]
20#[async_trait]
21pub trait ProverService: Sync + Send {
22 async fn compute_cache(&self, up_to: BlockNumber) -> StdResult<()>;
24}
25
26pub struct MithrilProverService<S: MKTreeStorer> {
28 _transaction_retriever: Arc<dyn TransactionsRetriever>,
29 block_range_root_retriever: Arc<dyn BlockRangeRootRetriever<S>>,
30 mk_map_pool: ResourcePool<MKMap<BlockRange, MKMapNode<BlockRange, S>, S>>,
31 logger: Logger,
32}
33
34impl<S: MKTreeStorer> MithrilProverService<S> {
35 pub fn new(
37 _transaction_retriever: Arc<dyn TransactionsRetriever>,
38 block_range_root_retriever: Arc<dyn BlockRangeRootRetriever<S>>,
39 mk_map_pool_size: usize,
40 logger: Logger,
41 ) -> Self {
42 Self {
43 _transaction_retriever,
44 block_range_root_retriever,
45 mk_map_pool: ResourcePool::new(mk_map_pool_size, vec![]),
46 logger: logger.new_with_component_name::<Self>(),
47 }
48 }
49}
50
51#[async_trait]
52impl<S: MKTreeStorer> ProverService for MithrilProverService<S> {
53 async fn compute_cache(&self, up_to: BlockNumber) -> StdResult<()> {
54 let pool_size = self.mk_map_pool.size();
55 info!(
56 self.logger, "Starts computing the Merkle map pool resource of size {pool_size}";
57 "up_to_block_number" => *up_to,
58 );
59 let mk_map_cache = self
60 .block_range_root_retriever
61 .compute_merkle_map_from_block_range_roots(up_to)
62 .await?;
63 let mk_maps_new = (1..=pool_size)
64 .into_par_iter()
65 .map(|i| {
66 debug!(
67 self.logger,
68 "Computing the Merkle map pool resource {i}/{pool_size}"
69 );
70 mk_map_cache.clone()
71 })
72 .collect::<Vec<MKMap<_, _, _>>>();
73 debug!(self.logger, "Draining the Merkle map pool");
74 let discriminant_new = self.mk_map_pool.discriminant()? + 1;
75 self.mk_map_pool.set_discriminant(discriminant_new)?;
76 self.mk_map_pool.clear();
77 debug!(
78 self.logger,
79 "Giving back new resources to the Merkle map pool"
80 );
81 mk_maps_new
82 .into_iter()
83 .map(|mk_map| self.mk_map_pool.give_back_resource(mk_map, discriminant_new))
84 .collect::<StdResult<Vec<_>>>()?;
85 info!(
86 self.logger,
87 "Completed computing the Merkle map pool resource of size {pool_size}"
88 );
89
90 Ok(())
91 }
92}