Merge pull request #2108 from input-output-hk/sfa/2075/refactor_signer_pruning_with_upkeep_service

Refactor signer pruning with upkeep service

Author: sfauvel

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-persistence"
-version = "0.2.32"
+version = "0.2.33"
 description = "Common types, interfaces, and utilities to persist data for Mithril nodes."
 authors = { workspace = true }
 edition = { workspace = true }

internal/mithril-persistence/Cargo.toml

@@ -5,5 +5,5 @@ pub mod adapter;
 mod stake_store;
 mod store_pruner;
 
-pub use stake_store::{StakeStore, StakeStorer};
+pub use stake_store::StakeStorer;
 pub use store_pruner::StorePruner;

internal/mithril-persistence/src/store/mod.rs

@@ -1,12 +1,8 @@
 use async_trait::async_trait;
-use mithril_common::entities::{Epoch, StakeDistribution};
-use mithril_common::signable_builder::StakeDistributionRetriever;
-use mithril_common::StdResult;
-use tokio::sync::RwLock;
-
-use super::{adapter::StoreAdapter, StorePruner};
-
-type Adapter = Box<dyn StoreAdapter<Key = Epoch, Record = StakeDistribution>>;
+use mithril_common::{
+    entities::{Epoch, StakeDistribution},
+    StdResult,
+};
 
 /// Represent a way to store the stake of mithril party members.
 #[async_trait]
@@ -21,184 +17,3 @@ pub trait StakeStorer: Sync + Send {
     /// Get the stakes of all party at a given `epoch`.
     async fn get_stakes(&self, epoch: Epoch) -> StdResult<Option<StakeDistribution>>;
 }
-
-/// A [StakeStorer] that use a [StoreAdapter] to store data.
-pub struct StakeStore {
-    adapter: RwLock<Adapter>,
-    retention_limit: Option<usize>,
-}
-
-impl StakeStore {
-    /// StakeStore factory
-    pub fn new(adapter: Adapter, retention_limit: Option<usize>) -> Self {
-        Self {
-            adapter: RwLock::new(adapter),
-            retention_limit,
-        }
-    }
-}
-
-#[async_trait]
-impl StorePruner for StakeStore {
-    type Key = Epoch;
-    type Record = StakeDistribution;
-
-    fn get_adapter(
-        &self,
-    ) -> &RwLock<Box<dyn StoreAdapter<Key = Self::Key, Record = Self::Record>>> {
-        &self.adapter
-    }
-
-    fn get_max_records(&self) -> Option<usize> {
-        self.retention_limit
-    }
-}
-
-#[async_trait]
-impl StakeStorer for StakeStore {
-    async fn save_stakes(
-        &self,
-        epoch: Epoch,
-        stakes: StakeDistribution,
-    ) -> StdResult<Option<StakeDistribution>> {
-        let signers = {
-            let mut adapter = self.adapter.write().await;
-            let signers = adapter.get_record(&epoch).await?;
-            adapter.store_record(&epoch, &stakes).await?;
-
-            signers
-        };
-        // it is important the adapter gets out of the scope to free the write lock it holds.
-        // Otherwise the method below will hang forever waiting for the lock.
-        self.prune().await?;
-
-        Ok(signers)
-    }
-
-    async fn get_stakes(&self, epoch: Epoch) -> StdResult<Option<StakeDistribution>> {
-        Ok(self.adapter.read().await.get_record(&epoch).await?)
-    }
-}
-
-#[async_trait]
-impl StakeDistributionRetriever for StakeStore {
-    async fn retrieve(&self, epoch: Epoch) -> StdResult<Option<StakeDistribution>> {
-        let stake_distribution = self.get_stakes(epoch).await?;
-
-        Ok(stake_distribution)
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::super::adapter::MemoryAdapter;
-    use super::*;
-
-    fn init_store(
-        nb_epoch: u64,
-        signers_per_epoch: u64,
-        retention_limit: Option<usize>,
-    ) -> StakeStore {
-        let mut values: Vec<(Epoch, StakeDistribution)> = Vec::new();
-
-        for epoch in 1..=nb_epoch {
-            let mut signers: StakeDistribution = StakeDistribution::new();
-
-            for party_idx in 1..=signers_per_epoch {
-                let party_id = format!("{party_idx}");
-                signers.insert(party_id.clone(), 100 * party_idx + 1);
-            }
-            values.push((Epoch(epoch), signers));
-        }
-
-        let values = if !values.is_empty() {
-            Some(values)
-        } else {
-            None
-        };
-        let adapter: MemoryAdapter<Epoch, StakeDistribution> = MemoryAdapter::new(values).unwrap();
-        StakeStore::new(Box::new(adapter), retention_limit)
-    }
-
-    #[tokio::test]
-    async fn save_key_in_empty_store() {
-        let store = init_store(0, 0, None);
-        let res = store
-            .save_stakes(Epoch(1), StakeDistribution::from([("1".to_string(), 123)]))
-            .await
-            .expect("Test adapter should not fail.");
-
-        assert!(res.is_none());
-    }
-
-    #[tokio::test]
-    async fn update_signer_in_store() {
-        let store = init_store(1, 1, None);
-        let res = store
-            .save_stakes(Epoch(1), StakeDistribution::from([("1".to_string(), 123)]))
-            .await
-            .expect("Test adapter should not fail.");
-
-        assert_eq!(
-            StakeDistribution::from([("1".to_string(), 101)]),
-            res.expect("the result should not be empty"),
-        );
-    }
-
-    #[tokio::test]
-    async fn get_stakes_for_empty_epoch() {
-        let store = init_store(2, 1, None);
-        let res = store
-            .get_stakes(Epoch(0))
-            .await
-            .expect("Test adapter should not fail.");
-
-        assert!(res.is_none());
-    }
-
-    #[tokio::test]
-    async fn get_stakes_for_existing_epoch() {
-        let store = init_store(2, 2, None);
-        let res = store
-            .get_stakes(Epoch(1))
-            .await
-            .expect("Test adapter should not fail.");
-
-        assert!(res.is_some());
-        assert_eq!(2, res.expect("Query result should not be empty.").len());
-    }
-
-    #[tokio::test]
-    async fn check_retention_limit() {
-        let store = init_store(2, 2, Some(2));
-        let _res = store
-            .save_stakes(Epoch(3), StakeDistribution::from([("1".to_string(), 123)]))
-            .await
-            .unwrap();
-        assert!(store.get_stakes(Epoch(1)).await.unwrap().is_none());
-    }
-
-    #[tokio::test]
-    async fn retrieve_with_no_stakes_returns_none() {
-        let store = init_store(0, 0, None);
-
-        let result = store.retrieve(Epoch(1)).await.unwrap();
-
-        assert!(result.is_none());
-    }
-
-    #[tokio::test]
-    async fn retrieve_returns_stake_distribution() {
-        let stake_distribution_to_retrieve =
-            StakeDistribution::from([("pool-123".to_string(), 123)]);
-        let store = init_store(0, 0, None);
-        store
-            .save_stakes(Epoch(1), stake_distribution_to_retrieve.clone())
-            .await
-            .unwrap();
-
-        let stake_distribution = store.retrieve(Epoch(1)).await.unwrap();
-
-        assert_eq!(stake_distribution, Some(stake_distribution_to_retrieve));
-    }
-}

internal/mithril-persistence/src/store/stake_store.rs

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-signer"
-version = "0.2.211"
+version = "0.2.212"
 description = "A Mithril Signer"
 authors = { workspace = true }
 edition = { workspace = true }

mithril-signer/Cargo.toml

@@ -33,7 +33,6 @@ use mithril_persistence::database::repository::CardanoTransactionRepository;
 use mithril_persistence::database::{ApplicationNodeType, SqlMigration};
 use mithril_persistence::sqlite::{ConnectionBuilder, SqliteConnection, SqliteConnectionPool};
 use mithril_persistence::store::adapter::SQLiteAdapter;
-use mithril_persistence::store::StakeStore;
 
 use crate::database::repository::SignedBeaconRepository;
 use crate::dependency_injection::SignerDependencyContainer;
@@ -44,7 +43,7 @@ use crate::services::{
     SignerUpkeepService, TransactionsImporterByChunk, TransactionsImporterWithPruner,
     TransactionsImporterWithVacuum,
 };
-use crate::store::{MKTreeStoreSqlite, ProtocolInitializerStore};
+use crate::store::{MKTreeStoreSqlite, ProtocolInitializerStore, StakeStore};
 use crate::{
     Configuration, MetricsService, HTTP_REQUEST_TIMEOUT_DURATION, SQLITE_FILE,
     SQLITE_FILE_CARDANO_TRANSACTION,
@@ -371,7 +370,11 @@ impl<'a> DependenciesBuilder<'a> {
             sqlite_connection.clone(),
             sqlite_connection_cardano_transaction_pool,
             signed_entity_type_lock.clone(),
-            vec![signed_beacon_repository.clone()],
+            vec![
+                signed_beacon_repository.clone(),
+                stake_store.clone(),
+                protocol_initializer_store.clone(),
+            ],
             self.root_logger(),
         ));
         let certifier = Arc::new(SignerCertifierService::new(

mithril-signer/src/dependency_injection/builder.rs

@@ -8,13 +8,12 @@ use mithril_common::era::{EraChecker, EraReader};
 use mithril_common::signable_builder::SignableBuilderService;
 use mithril_common::signed_entity_type_lock::SignedEntityTypeLock;
 use mithril_common::TickerService;
-use mithril_persistence::store::StakeStore;
 use tokio::sync::RwLock;
 
 use crate::services::{
     AggregatorClient, CertifierService, EpochService, SingleSigner, UpkeepService,
 };
-use crate::store::ProtocolInitializerStorer;
+use crate::store::{ProtocolInitializerStorer, StakeStore};
 use crate::MetricsService;
 
 type StakeStoreService = Arc<StakeStore>;

mithril-signer/src/dependency_injection/containers.rs

@@ -358,7 +358,7 @@ mod tests {
         MithrilTickerService, TickerService,
     };
     use mithril_persistence::store::adapter::{DumbStoreAdapter, MemoryAdapter};
-    use mithril_persistence::store::{StakeStore, StakeStorer};
+    use mithril_persistence::store::StakeStorer;
 
     use crate::database::repository::SignedBeaconRepository;
     use crate::database::test_helper::main_db_connection;
@@ -368,7 +368,7 @@ mod tests {
         MithrilSingleSigner, MockTransactionStore, MockUpkeepService, SignerCertifierService,
         SignerSignableSeedBuilder, SignerSignedEntityConfigProvider,
     };
-    use crate::store::ProtocolInitializerStore;
+    use crate::store::{ProtocolInitializerStore, StakeStore};
     use crate::test_tools::TestLogger;
 
     use super::*;

mithril-signer/src/runtime/runner.rs

@@ -5,6 +5,11 @@ use std::collections::BTreeSet;
 use std::sync::Arc;
 use thiserror::Error;
 
+use crate::dependency_injection::EpochServiceWrapper;
+use crate::entities::SignerEpochSettings;
+use crate::services::SignedEntityConfigProvider;
+use crate::store::ProtocolInitializerStorer;
+use crate::RunnerError;
 use mithril_common::crypto_helper::ProtocolInitializer;
 use mithril_common::entities::{
     CardanoTransactionsSigningConfig, Epoch, PartyId, ProtocolParameters, SignedEntityConfig,
@@ -14,12 +19,6 @@ use mithril_common::logging::LoggerExtensions;
 use mithril_common::{CardanoNetwork, StdResult};
 use mithril_persistence::store::StakeStorer;
 
-use crate::dependency_injection::EpochServiceWrapper;
-use crate::entities::SignerEpochSettings;
-use crate::services::SignedEntityConfigProvider;
-use crate::store::ProtocolInitializerStorer;
-use crate::RunnerError;
-
 /// Errors dedicated to the EpochService.
 #[derive(Debug, Error)]
 pub enum EpochServiceError {
@@ -327,9 +326,9 @@ impl MithrilEpochService {
     /// `TEST ONLY` - Create a new instance of the service using dumb dependencies.
     pub fn new_with_dumb_dependencies() -> Self {
         use crate::store::ProtocolInitializerStore;
+        use crate::store::StakeStore;
         use crate::test_tools::TestLogger;
         use mithril_persistence::store::adapter::DumbStoreAdapter;
-        use mithril_persistence::store::StakeStore;
 
         let stake_store = Arc::new(StakeStore::new(Box::new(DumbStoreAdapter::new()), None));
         let protocol_initializer_store = Arc::new(ProtocolInitializerStore::new(
@@ -437,11 +436,10 @@ mod tests {
     use mithril_common::entities::{Epoch, StakeDistribution};
     use mithril_common::test_utils::{fake_data, MithrilFixtureBuilder};
     use mithril_persistence::store::adapter::{DumbStoreAdapter, MemoryAdapter};
-    use mithril_persistence::store::{StakeStore, StakeStorer};
 
     use crate::entities::SignerEpochSettings;
     use crate::services::MithrilProtocolInitializerBuilder;
-    use crate::store::ProtocolInitializerStore;
+    use crate::store::{ProtocolInitializerStore, StakeStore};
     use crate::test_tools::TestLogger;
 
     use super::*;

mithril-signer/src/services/epoch_service.rs

@@ -176,10 +176,9 @@ mod tests {
     use mithril_common::entities::{Epoch, ProtocolMessagePartKey};
     use mithril_common::test_utils::MithrilFixtureBuilder;
     use mithril_persistence::store::adapter::{DumbStoreAdapter, MemoryAdapter};
-    use mithril_persistence::store::StakeStore;
 
     use crate::services::MithrilEpochService;
-    use crate::store::ProtocolInitializerStore;
+    use crate::store::{ProtocolInitializerStore, StakeStore};
     use crate::test_tools::TestLogger;
 
     use super::*;