refactor: move transaction prover cache computation in artifact builder

Author: jpraynaud

mithril-aggregator/src/artifact_builder/cardano_transactions.rs

@@ -1,5 +1,9 @@
+use std::sync::Arc;
+
 use async_trait::async_trait;
 
+use crate::services::ProverService;
+
 use super::ArtifactBuilder;
 use anyhow::{anyhow, Context};
 use mithril_common::{
@@ -11,12 +15,14 @@ use mithril_common::{
 };
 
 /// A [CardanoTransactionsArtifact] builder
-pub struct CardanoTransactionsArtifactBuilder {}
+pub struct CardanoTransactionsArtifactBuilder {
+    prover_service: Arc<dyn ProverService>,
+}
 
 impl CardanoTransactionsArtifactBuilder {
     /// CardanoTransactions artifact builder factory
-    pub fn new() -> Self {
-        Self {}
+    pub fn new(prover_service: Arc<dyn ProverService>) -> Self {
+        Self { prover_service }
     }
 }
 
@@ -41,6 +47,7 @@ impl ArtifactBuilder<CardanoDbBeacon, CardanoTransactionsSnapshot>
                     SignedEntityType::CardanoTransactions(beacon.clone())
                 )
             })?;
+        self.prover_service.compute_cache(&beacon).await?;
 
         Ok(CardanoTransactionsSnapshot::new(
             merkle_root.to_string(),
@@ -51,10 +58,9 @@ impl ArtifactBuilder<CardanoDbBeacon, CardanoTransactionsSnapshot>
 
 #[cfg(test)]
 mod tests {
-    use mithril_common::{
-        entities::ProtocolMessage,
-        test_utils::fake_data::{self},
-    };
+    use mithril_common::{entities::ProtocolMessage, test_utils::fake_data};
+
+    use crate::services::MockProverService;
 
     use super::*;
 
@@ -72,7 +78,10 @@ mod tests {
         };
 
         let beacon = certificate.as_cardano_db_beacon();
-        let cardano_transaction_artifact_builder = CardanoTransactionsArtifactBuilder::new();
+        let mut mock_prover = MockProverService::new();
+        mock_prover.expect_compute_cache().returning(|_| Ok(()));
+        let cardano_transaction_artifact_builder =
+            CardanoTransactionsArtifactBuilder::new(Arc::new(mock_prover));
         let artifact = cardano_transaction_artifact_builder
             .compute_artifact(beacon.clone(), &certificate)
             .await
@@ -90,7 +99,10 @@ mod tests {
             certificate
         };
 
-        let cardano_transaction_artifact_builder = CardanoTransactionsArtifactBuilder::new();
+        let mut mock_prover = MockProverService::new();
+        mock_prover.expect_compute_cache().returning(|_| Ok(()));
+        let cardano_transaction_artifact_builder =
+            CardanoTransactionsArtifactBuilder::new(Arc::new(mock_prover));
         cardano_transaction_artifact_builder
             .compute_artifact(CardanoDbBeacon::default(), &certificate)
             .await

mithril-aggregator/src/dependency_injection/builder.rs

@@ -1104,15 +1104,29 @@ impl DependenciesBuilder {
                 snapshot_uploader,
                 self.configuration.snapshot_compression_algorithm,
             ));
-        let cardano_transactions_artifact_builder =
-            Arc::new(CardanoTransactionsArtifactBuilder::new());
+        let prover_service = self.get_prover_service().await?;
+        let cardano_transactions_artifact_builder = Arc::new(
+            CardanoTransactionsArtifactBuilder::new(prover_service.clone()),
+        );
         let signed_entity_service = Arc::new(MithrilSignedEntityService::new(
             signed_entity_storer,
             mithril_stake_distribution_artifact_builder,
             cardano_immutable_files_full_artifact_builder,
             cardano_transactions_artifact_builder,
         ));
 
+        // Compute the cache pool for prover service
+        // This is done here to avoid circular dependencies between the prover service and the signed entity service
+        // TODO: Make this part of a warmup phase of the aggregator?
+        if let Some(signed_entity) = signed_entity_service
+            .get_last_cardano_transaction_snapshot()
+            .await?
+        {
+            prover_service
+                .compute_cache(&signed_entity.artifact.beacon)
+                .await?;
+        }
+
         Ok(signed_entity_service)
     }
 
@@ -1378,17 +1392,6 @@ impl DependenciesBuilder {
             logger,
         );
 
-        // Compute the cache pool for prover service
-        let signed_entity_service = self.get_signed_entity_service().await?;
-        if let Some(signed_entity) = signed_entity_service
-            .get_last_cardano_transaction_snapshot()
-            .await?
-        {
-            prover_service
-                .compute_cache(&signed_entity.artifact.beacon)
-                .await?;
-        }
-
         Ok(Arc::new(prover_service))
     }
 

mithril-aggregator/src/runtime/runner.rs

@@ -399,13 +399,6 @@ 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(())
     }