Implement lock/unlock mechanism for signed entity types during the create artifact process

Author: dlachaume

mithril-aggregator/src/dependency_injection/builder.rs

@@ -1135,6 +1135,7 @@ impl DependenciesBuilder {
             mithril_stake_distribution_artifact_builder,
             cardano_immutable_files_full_artifact_builder,
             cardano_transactions_artifact_builder,
+            self.get_signed_entity_lock().await?,
         ));
 
         // Compute the cache pool for prover service

mithril-aggregator/src/services/signed_entity.rs

@@ -2,7 +2,7 @@
 //!
 //! This service is responsible for dealing with [SignedEntity] type.
 //! It creates [Artifact] that can be accessed by clients.
-use anyhow::Context;
+use anyhow::{anyhow, Context};
 use async_trait::async_trait;
 use chrono::Utc;
 use slog_scope::info;
@@ -16,6 +16,7 @@ use mithril_common::{
         Snapshot,
     },
     signable_builder::Artifact,
+    signed_entity_type_lock::SignedEntityTypeLock,
     StdResult,
 };
 
@@ -79,6 +80,7 @@ pub struct MithrilSignedEntityService {
         Arc<dyn ArtifactBuilder<CardanoDbBeacon, Snapshot>>,
     cardano_transactions_artifact_builder:
         Arc<dyn ArtifactBuilder<BlockNumber, CardanoTransactionsSnapshot>>,
+    signed_entity_type_lock: Arc<SignedEntityTypeLock>,
 }
 
 impl MithrilSignedEntityService {
@@ -94,29 +96,56 @@ impl MithrilSignedEntityService {
         cardano_transactions_artifact_builder: Arc<
             dyn ArtifactBuilder<BlockNumber, CardanoTransactionsSnapshot>,
         >,
+        signed_entity_type_lock: Arc<SignedEntityTypeLock>,
     ) -> Self {
         Self {
             signed_entity_storer,
             mithril_stake_distribution_artifact_builder,
             cardano_immutable_files_full_artifact_builder,
             cardano_transactions_artifact_builder,
+            signed_entity_type_lock,
         }
     }
 
-    fn create_artifact_return_join_handle(
+    async fn create_artifact_return_join_handle(
         &self,
         signed_entity_type: SignedEntityType,
         certificate: &Certificate,
-    ) -> JoinHandle<StdResult<()>> {
+    ) -> StdResult<JoinHandle<StdResult<()>>> {
+        if self
+            .signed_entity_type_lock
+            .is_locked(&signed_entity_type)
+            .await
+        {
+            return Err(anyhow!(
+                "Signed entity type '{:?}' is already locked",
+                signed_entity_type
+            ));
+        }
+
         let service = self.clone();
         let certificate_cloned = certificate.clone();
-        tokio::task::spawn_blocking(|| {
-            Handle::current().block_on(async move {
-                service
-                    .create_artifact(signed_entity_type, &certificate_cloned)
+        service
+            .signed_entity_type_lock
+            .lock(&signed_entity_type)
+            .await;
+
+        Ok(tokio::task::spawn(async move {
+            let signed_entity_type_clone = signed_entity_type.clone();
+            let service_clone = service.clone();
+            let result = tokio::task::spawn(async move {
+                service_clone
+                    .create_artifact(signed_entity_type_clone, &certificate_cloned)
                     .await
             })
-        })
+            .await;
+            service
+                .signed_entity_type_lock
+                .release(signed_entity_type)
+                .await;
+
+            result.unwrap()
+        }))
     }
 
     /// Compute artifact from signed entity type
@@ -190,6 +219,11 @@ impl SignedEntityService for MithrilSignedEntityService {
             "certificate_hash" => &certificate.hash
         );
 
+        println!(
+            "MithrilSignedEntityService::create_artifact: signed_entity_type: {:?}, certificate_hash: {}",
+            signed_entity_type, certificate.hash
+        );
+
         let mut remaining_retries = 2;
         let artifact = loop {
             remaining_retries -= 1;
@@ -317,14 +351,15 @@ impl SignedEntityService for MithrilSignedEntityService {
 
 #[cfg(test)]
 mod tests {
-    use std::time::Duration;
+    use std::{sync::atomic::Ordering, time::Duration};
 
     use mithril_common::{
         entities::{CardanoTransactionsSnapshot, Epoch},
         signable_builder,
         test_utils::fake_data,
     };
     use serde::{de::DeserializeOwned, Serialize};
+    use std::sync::atomic::AtomicBool;
 
     use crate::artifact_builder::MockArtifactBuilder;
     use crate::database::repository::MockSignedEntityStorer;
@@ -386,6 +421,58 @@ mod tests {
                 Arc::new(self.mock_mithril_stake_distribution_artifact_builder),
                 Arc::new(self.mock_cardano_immutable_files_full_artifact_builder),
                 Arc::new(self.mock_cardano_transactions_artifact_builder),
+                Arc::new(SignedEntityTypeLock::default()),
+            )
+        }
+
+        fn build_artifact_builder_service_with_time_consuming_process(
+            mut self,
+            atomic_stop: Arc<AtomicBool>,
+        ) -> MithrilSignedEntityService {
+            struct LongArtifactBuilder {
+                atomic_stop: Arc<AtomicBool>,
+                snapshot: Snapshot,
+            }
+            impl LongArtifactBuilder {}
+
+            let snapshot = fake_data::snapshots(1).first().unwrap().to_owned();
+
+            #[async_trait]
+            impl ArtifactBuilder<CardanoDbBeacon, Snapshot> for LongArtifactBuilder {
+                async fn compute_artifact(
+                    &self,
+                    _beacon: CardanoDbBeacon,
+                    _certificate: &Certificate,
+                ) -> StdResult<Snapshot> {
+                    let mut max_iteration = 100;
+                    while !self.atomic_stop.load(Ordering::Relaxed) {
+                        max_iteration -= 1;
+                        if max_iteration <= 0 {
+                            return Err(anyhow!("Test should handle the stop"));
+                        }
+                        tokio::time::sleep(Duration::from_millis(10)).await;
+                    }
+                    Ok(self.snapshot.clone())
+                }
+            }
+            let cardano_immutable_files_full_long_artifact_builder = LongArtifactBuilder {
+                atomic_stop: atomic_stop.clone(),
+                snapshot: snapshot.clone(),
+            };
+
+            let artifact_clone: Arc<dyn Artifact> = Arc::new(snapshot);
+            let signed_entity_artifact = serde_json::to_string(&artifact_clone).unwrap();
+            self.mock_signed_entity_storer
+                .expect_store_signed_entity()
+                .withf(move |signed_entity| signed_entity.artifact == signed_entity_artifact)
+                .return_once(|_| Ok(()));
+
+            MithrilSignedEntityService::new(
+                Arc::new(self.mock_signed_entity_storer),
+                Arc::new(self.mock_mithril_stake_distribution_artifact_builder),
+                Arc::new(cardano_immutable_files_full_long_artifact_builder),
+                Arc::new(self.mock_cardano_transactions_artifact_builder),
+                Arc::new(SignedEntityTypeLock::default()),
             )
         }
 
@@ -585,7 +672,7 @@ mod tests {
     // TODO: Verify the relevance of this test
     #[tokio::test]
     async fn create_artifact_for_two_signed_entity_types_in_sequence() {
-        let artifact_builder_service = {
+        let signed_entity_type_service = {
             let mut mock_container = MockDependencyInjector::new();
             mock_container.mock_immutable_files_processing(
                 fake_data::snapshots(1).first().unwrap().to_owned(),
@@ -599,49 +686,167 @@ mod tests {
 
         let signed_entity_type_immutable =
             SignedEntityType::CardanoImmutableFilesFull(CardanoDbBeacon::default());
-        artifact_builder_service
+        signed_entity_type_service
             .create_artifact(signed_entity_type_immutable, &certificate)
             .await
             .unwrap();
 
         let signed_entity_type_msd = SignedEntityType::MithrilStakeDistribution(Epoch(1));
-        artifact_builder_service
+        signed_entity_type_service
             .create_artifact(signed_entity_type_msd, &certificate)
             .await
             .unwrap();
     }
 
     #[tokio::test]
     async fn create_artifact_for_two_signed_entity_types_in_sequence_not_blocking() {
-        let artifact_builder_service = {
+        let atomic_stop = Arc::new(AtomicBool::new(false));
+        let signed_entity_type_service = {
             let mut mock_container = MockDependencyInjector::new();
-            let snapshot = fake_data::snapshots(1).first().unwrap().to_owned();
-            mock_container
-                .mock_cardano_immutable_files_full_artifact_builder
-                .expect_compute_artifact()
-                .times(1)
-                .return_once(|_, _| {
-                    std::thread::sleep(Duration::from_millis(1000));
-                    Ok(snapshot)
-                });
 
             let msd = create_stake_distribution(Epoch(1), 5);
             mock_container.mock_stake_distribution_processing(msd);
 
-            Arc::new(mock_container.build_artifact_builder_service())
+            mock_container
+                .build_artifact_builder_service_with_time_consuming_process(atomic_stop.clone())
         };
         let certificate = fake_data::certificate("hash".to_string());
 
         let signed_entity_type_immutable =
             SignedEntityType::CardanoImmutableFilesFull(CardanoDbBeacon::default());
-        let first_task_that_never_finished = artifact_builder_service
-            .create_artifact_return_join_handle(signed_entity_type_immutable, &certificate);
+        let first_task_that_never_finished = signed_entity_type_service
+            .create_artifact_return_join_handle(signed_entity_type_immutable, &certificate)
+            .await
+            .unwrap();
 
         let signed_entity_type_msd = SignedEntityType::MithrilStakeDistribution(Epoch(1));
-        let second_task_that_finish_first = artifact_builder_service
-            .create_artifact_return_join_handle(signed_entity_type_msd, &certificate);
+        let second_task_that_finish_first = signed_entity_type_service
+            .create_artifact_return_join_handle(signed_entity_type_msd, &certificate)
+            .await
+            .unwrap();
 
-        let _ = second_task_that_finish_first.await.unwrap();
+        second_task_that_finish_first.await.unwrap().unwrap();
         assert!(!first_task_that_never_finished.is_finished());
+
+        atomic_stop.swap(true, Ordering::Relaxed);
+    }
+
+    #[tokio::test]
+    async fn create_artifact_lock_unlock_signed_entity_type_while_processing() {
+        let atomic_stop = Arc::new(AtomicBool::new(false));
+        let signed_entity_type_service = MockDependencyInjector::new()
+            .build_artifact_builder_service_with_time_consuming_process(atomic_stop.clone());
+        let certificate = fake_data::certificate("hash".to_string());
+
+        let signed_entity_type_immutable =
+            SignedEntityType::CardanoImmutableFilesFull(CardanoDbBeacon::default());
+        assert!(
+            !signed_entity_type_service
+                .signed_entity_type_lock
+                .is_locked(&signed_entity_type_immutable)
+                .await
+        );
+        let join_handle = signed_entity_type_service
+            .create_artifact_return_join_handle(signed_entity_type_immutable.clone(), &certificate)
+            .await
+            .unwrap();
+
+        // Results are stored to finalize the task before assertions,
+        // ensuring 'atomic_stop' is always assigned a new value.
+        let is_locked = signed_entity_type_service
+            .signed_entity_type_lock
+            .is_locked(&signed_entity_type_immutable)
+            .await;
+        let is_finished = join_handle.is_finished();
+
+        atomic_stop.swap(true, Ordering::Relaxed);
+        join_handle.await.unwrap().unwrap();
+
+        assert!(is_locked);
+        assert!(!is_finished);
+
+        assert!(
+            !signed_entity_type_service
+                .signed_entity_type_lock
+                .is_locked(&signed_entity_type_immutable)
+                .await
+        );
+    }
+
+    #[tokio::test]
+    async fn create_artifact_unlock_signed_entity_type_when_error() {
+        let signed_entity_type_service = {
+            let mut mock_container = MockDependencyInjector::new();
+            mock_container
+                .mock_cardano_immutable_files_full_artifact_builder
+                .expect_compute_artifact()
+                .returning(|_, _| Err(anyhow::anyhow!("Error while computing artifact")));
+
+            mock_container.build_artifact_builder_service()
+        };
+        let certificate = fake_data::certificate("hash".to_string());
+
+        let signed_entity_type_immutable =
+            SignedEntityType::CardanoImmutableFilesFull(CardanoDbBeacon::default());
+
+        let join_handle = signed_entity_type_service
+            .create_artifact_return_join_handle(signed_entity_type_immutable.clone(), &certificate)
+            .await
+            .unwrap();
+
+        join_handle.await.unwrap().unwrap_err();
+
+        assert!(
+            !signed_entity_type_service
+                .signed_entity_type_lock
+                .is_locked(&signed_entity_type_immutable)
+                .await
+        );
+    }
+
+    #[tokio::test]
+    async fn create_artifact_unlock_signed_entity_type_when_panic() {
+        let signed_entity_type_service =
+            MockDependencyInjector::new().build_artifact_builder_service();
+        let certificate = fake_data::certificate("hash".to_string());
+
+        let signed_entity_type_immutable =
+            SignedEntityType::CardanoImmutableFilesFull(CardanoDbBeacon::default());
+
+        let join_handle = signed_entity_type_service
+            .create_artifact_return_join_handle(signed_entity_type_immutable.clone(), &certificate)
+            .await
+            .unwrap();
+
+        join_handle.await.unwrap_err();
+
+        assert!(
+            !signed_entity_type_service
+                .signed_entity_type_lock
+                .is_locked(&signed_entity_type_immutable)
+                .await
+        );
+    }
+
+    #[tokio::test]
+    async fn create_artifact_for_a_signed_entity_type_already_lock_return_error() {
+        let atomic_stop = Arc::new(AtomicBool::new(false));
+        let signed_entity_service = MockDependencyInjector::new()
+            .build_artifact_builder_service_with_time_consuming_process(atomic_stop.clone());
+        let certificate = fake_data::certificate("hash".to_string());
+        let signed_entity_type_immutable =
+            SignedEntityType::CardanoImmutableFilesFull(CardanoDbBeacon::default());
+
+        signed_entity_service
+            .create_artifact_return_join_handle(signed_entity_type_immutable.clone(), &certificate)
+            .await
+            .unwrap();
+
+        signed_entity_service
+            .create_artifact_return_join_handle(signed_entity_type_immutable, &certificate)
+            .await
+            .expect_err("Should return error when signed entity type is already locked");
+
+        atomic_stop.swap(true, Ordering::Relaxed);
     }
 }