mithril_signer/services/signature_publisher/
retrier.rs

1use std::{sync::Arc, time::Duration};
2
3use mithril_common::{
4    entities::{ProtocolMessage, SignedEntityType, SingleSignature},
5    StdResult,
6};
7
8use super::SignaturePublisher;
9
10/// Policy for retrying signature publishing
11#[derive(Debug, PartialEq, Clone)]
12pub struct SignaturePublishRetryPolicy {
13    /// Number of attempts to publish a signature
14    pub attempts: u8,
15    /// Delay between two attempts
16    pub delay_between_attempts: Duration,
17}
18
19impl SignaturePublishRetryPolicy {
20    /// Create a policy that never retries
21    pub fn never() -> Self {
22        Self {
23            attempts: 1,
24            delay_between_attempts: Duration::from_secs(0),
25        }
26    }
27}
28
29impl Default for SignaturePublishRetryPolicy {
30    /// Create a default retry policy
31    fn default() -> Self {
32        Self {
33            attempts: 3,
34            delay_between_attempts: Duration::from_secs(5),
35        }
36    }
37}
38
39/// A decorator of [SignaturePublisher] that retries the publishing of signatures in case of failure
40pub struct SignaturePublisherRetrier {
41    publisher: Arc<dyn SignaturePublisher>,
42    retry_policy: SignaturePublishRetryPolicy,
43}
44
45impl SignaturePublisherRetrier {
46    /// Creates a new [SignaturePublisherRetrier]
47    pub fn new(
48        publisher: Arc<dyn SignaturePublisher>,
49        retry_policy: SignaturePublishRetryPolicy,
50    ) -> Self {
51        Self {
52            publisher,
53            retry_policy,
54        }
55    }
56}
57
58#[async_trait::async_trait]
59impl SignaturePublisher for SignaturePublisherRetrier {
60    async fn publish(
61        &self,
62        signed_entity_type: &SignedEntityType,
63        signature: &SingleSignature,
64        protocol_message: &ProtocolMessage,
65    ) -> StdResult<()> {
66        let mut nb_attempts = 0;
67        loop {
68            nb_attempts += 1;
69
70            match self
71                .publisher
72                .publish(signed_entity_type, signature, protocol_message)
73                .await
74            {
75                Ok(_) => return Ok(()),
76                Err(e) if nb_attempts >= self.retry_policy.attempts => {
77                    return Err(anyhow::anyhow!(e)
78                        .context(format!("Publish failed after {nb_attempts} attempts")));
79                }
80                _ => tokio::time::sleep(self.retry_policy.delay_between_attempts).await,
81            }
82        }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use mithril_common::{entities::Epoch, test_utils::fake_data};
89
90    use crate::services::MockSignaturePublisher;
91
92    use super::*;
93
94    #[tokio::test]
95    async fn should_call_publish_once_when_no_retry_policy() {
96        let retry_policy = SignaturePublishRetryPolicy::never();
97
98        let mut publisher = MockSignaturePublisher::new();
99        publisher
100            .expect_publish()
101            .once()
102            .returning(|_, _, _| Ok(()));
103
104        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
105
106        retrier
107            .publish(
108                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
109                &fake_data::single_signature(vec![1]),
110                &ProtocolMessage::default(),
111            )
112            .await
113            .unwrap();
114    }
115
116    #[tokio::test]
117    async fn should_not_retry_when_publish_fails_and_retry_policy_is_never() {
118        let retry_policy = SignaturePublishRetryPolicy::never();
119
120        let mut publisher = MockSignaturePublisher::new();
121        publisher
122            .expect_publish()
123            .once()
124            .returning(|_, _, _| Err(anyhow::anyhow!("error while publishing")));
125
126        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
127
128        retrier
129            .publish(
130                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
131                &fake_data::single_signature(vec![1]),
132                &ProtocolMessage::default(),
133            )
134            .await
135            .expect_err("An error should be returned");
136    }
137
138    #[tokio::test]
139    async fn should_retry_once_and_succeed_on_second_attempt() {
140        let retry_policy = SignaturePublishRetryPolicy {
141            attempts: 2,
142            delay_between_attempts: Duration::from_secs(0),
143        };
144
145        let mut publisher = MockSignaturePublisher::new();
146        publisher
147            .expect_publish()
148            .times(1)
149            .return_once(|_, _, _| Err(anyhow::anyhow!("error")));
150        publisher
151            .expect_publish()
152            .times(1)
153            .return_once(|_, _, _| Ok(()));
154
155        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
156
157        retrier
158            .publish(
159                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
160                &fake_data::single_signature(vec![1]),
161                &ProtocolMessage::default(),
162            )
163            .await
164            .unwrap();
165    }
166
167    #[tokio::test]
168    async fn should_retry_and_return_error_after_max_attempts() {
169        let retry_policy = SignaturePublishRetryPolicy {
170            attempts: 2,
171            delay_between_attempts: Duration::from_secs(0),
172        };
173
174        let mut publisher = MockSignaturePublisher::new();
175        publisher
176            .expect_publish()
177            .times(2)
178            .returning(|_, _, _| Err(anyhow::anyhow!("error")));
179
180        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
181
182        retrier
183            .publish(
184                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
185                &fake_data::single_signature(vec![1]),
186                &ProtocolMessage::default(),
187            )
188            .await
189            .expect_err("An error should be returned after max attempts");
190    }
191
192    #[tokio::test]
193    async fn should_wait_between_retries_according_to_policy() {
194        let delay_between_attempts = Duration::from_millis(50);
195        let retry_policy = SignaturePublishRetryPolicy {
196            attempts: 2,
197            delay_between_attempts,
198        };
199
200        let mut publisher = MockSignaturePublisher::new();
201        publisher
202            .expect_publish()
203            .once()
204            .return_once(|_, _, _| Err(anyhow::anyhow!("error")));
205        publisher
206            .expect_publish()
207            .once()
208            .return_once(|_, _, _| Ok(()));
209
210        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
211
212        let start_time = std::time::Instant::now();
213        retrier
214            .publish(
215                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
216                &fake_data::single_signature(vec![1]),
217                &ProtocolMessage::default(),
218            )
219            .await
220            .unwrap();
221
222        let elapsed_time = start_time.elapsed();
223        assert!(
224            elapsed_time >= delay_between_attempts,
225            "Expected at least {delay_between_attempts:?} time elapsed, but got {elapsed_time:?}"
226        );
227        assert!(
228            elapsed_time < delay_between_attempts * 2,
229            "Expected less than {:?} time elapsed, but got {:?}",
230            delay_between_attempts * 2,
231            elapsed_time
232        );
233    }
234}
mithril_signer/services/signature_publisher/retrier.rs

mithril_signer/services/signature_publisher/
retrier.rs
