mithril_aggregator/database/repository/
epoch_settings_store.rs

1use std::sync::Arc;
2
3use anyhow::Context;
4use async_trait::async_trait;
5
6use mithril_common::entities::{Epoch, ProtocolParameters};
7use mithril_common::StdResult;
8use mithril_persistence::sqlite::{ConnectionExtensions, SqliteConnection};
9
10use crate::database::query::{
11    DeleteEpochSettingsQuery, GetEpochSettingsQuery, UpdateEpochSettingsQuery,
12};
13use crate::entities::AggregatorEpochSettings;
14use crate::services::EpochPruningTask;
15use crate::{EpochSettingsStorer, ProtocolParametersRetriever};
16
17/// Service to deal with epoch settings (read & write).
18pub struct EpochSettingsStore {
19    connection: Arc<SqliteConnection>,
20
21    /// Number of epochs before previous records will be pruned at the next call to
22    /// [save_protocol_parameters][EpochSettingStore::save_protocol_parameters].
23    retention_limit: Option<u64>,
24}
25
26impl EpochSettingsStore {
27    /// Create a new EpochSettings store
28    pub fn new(connection: Arc<SqliteConnection>, retention_limit: Option<u64>) -> Self {
29        Self {
30            connection,
31            retention_limit,
32        }
33    }
34}
35
36#[async_trait]
37impl ProtocolParametersRetriever for EpochSettingsStore {
38    async fn get_protocol_parameters(&self, epoch: Epoch) -> StdResult<Option<ProtocolParameters>> {
39        Ok(self
40            .get_epoch_settings(epoch)
41            .await?
42            .map(|epoch_settings| epoch_settings.protocol_parameters))
43    }
44}
45
46#[async_trait]
47impl EpochSettingsStorer for EpochSettingsStore {
48    async fn save_epoch_settings(
49        &self,
50        epoch: Epoch,
51        epoch_settings: AggregatorEpochSettings,
52    ) -> StdResult<Option<AggregatorEpochSettings>> {
53        let epoch_settings_record = self
54            .connection
55            .fetch_first(UpdateEpochSettingsQuery::one(epoch, epoch_settings))
56            .with_context(|| format!("persist epoch settings failure for epoch {epoch:?}"))?
57            .unwrap_or_else(|| panic!("No entity returned by the persister, epoch = {epoch:?}"));
58
59        Ok(Some(epoch_settings_record.into()))
60    }
61
62    async fn get_epoch_settings(&self, epoch: Epoch) -> StdResult<Option<AggregatorEpochSettings>> {
63        let mut cursor = self
64            .connection
65            .fetch(GetEpochSettingsQuery::by_epoch(epoch)?)
66            .with_context(|| format!("Could not get epoch settings: epoch = {epoch:?}"))?;
67
68        if let Some(epoch_settings_record) = cursor.next() {
69            return Ok(Some(epoch_settings_record.into()));
70        }
71        Ok(None)
72    }
73}
74
75#[async_trait]
76impl EpochPruningTask for EpochSettingsStore {
77    fn pruned_data(&self) -> &'static str {
78        "Epoch settings"
79    }
80
81    /// Prune useless old epoch settings.
82    async fn prune(&self, epoch: Epoch) -> StdResult<()> {
83        if let Some(threshold) = self.retention_limit {
84            self.connection
85                .apply(DeleteEpochSettingsQuery::below_epoch_threshold(
86                    epoch - threshold,
87                ))?;
88        }
89        Ok(())
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use crate::database::test_helper::{insert_epoch_settings, main_db_connection};
96
97    use super::*;
98
99    #[tokio::test]
100    async fn prune_epoch_settings_older_than_threshold() {
101        const EPOCH_SETTINGS_PRUNE_EPOCH_THRESHOLD: u64 = 5;
102
103        let connection = main_db_connection().unwrap();
104        insert_epoch_settings(&connection, &[1, 2]).unwrap();
105        let store = EpochSettingsStore::new(
106            Arc::new(connection),
107            Some(EPOCH_SETTINGS_PRUNE_EPOCH_THRESHOLD),
108        );
109
110        store
111            .prune(Epoch(2) + EPOCH_SETTINGS_PRUNE_EPOCH_THRESHOLD)
112            .await
113            .unwrap();
114
115        let epoch1_params = store.get_epoch_settings(Epoch(1)).await.unwrap();
116        let epoch2_params = store.get_epoch_settings(Epoch(2)).await.unwrap();
117
118        assert!(
119            epoch1_params.is_none(),
120            "Epoch settings at epoch 1 should have been pruned",
121        );
122        assert!(
123            epoch2_params.is_some(),
124            "Epoch settings at epoch 2 should still exist",
125        );
126    }
127
128    #[tokio::test]
129    async fn without_threshold_nothing_is_pruned() {
130        let connection = main_db_connection().unwrap();
131        insert_epoch_settings(&connection, &[1, 2]).unwrap();
132        let store = EpochSettingsStore::new(Arc::new(connection), None);
133
134        store.prune(Epoch(100)).await.unwrap();
135
136        let epoch1_params = store.get_epoch_settings(Epoch(1)).await.unwrap();
137        let epoch2_params = store.get_epoch_settings(Epoch(2)).await.unwrap();
138
139        assert!(
140            epoch1_params.is_some(),
141            "Epoch settings at epoch 1 should have been pruned",
142        );
143        assert!(
144            epoch2_params.is_some(),
145            "Epoch settings at epoch 2 should still exist",
146        );
147    }
148
149    #[tokio::test]
150    async fn save_epoch_settings_stores_in_database() {
151        let connection = main_db_connection().unwrap();
152
153        let store = EpochSettingsStore::new(Arc::new(connection), None);
154
155        store
156            .save_epoch_settings(Epoch(2), AggregatorEpochSettings::dummy())
157            .await
158            .expect("saving epoch settings should not fails");
159        {
160            let epoch_settings = store.get_epoch_settings(Epoch(1)).await.unwrap();
161            assert_eq!(None, epoch_settings);
162        }
163        {
164            let epoch_settings = store.get_epoch_settings(Epoch(2)).await.unwrap().unwrap();
165            assert_eq!(AggregatorEpochSettings::dummy(), epoch_settings);
166        }
167        {
168            let epoch_settings = store.get_epoch_settings(Epoch(3)).await.unwrap();
169            assert_eq!(None, epoch_settings);
170        }
171    }
172}
mithril_aggregator/database/repository/epoch_settings_store.rs

mithril_aggregator/database/repository/
epoch_settings_store.rs
