mithril_signer/services/signature_publisher/
retrier.rs

1use std::{sync::Arc, time::Duration};
2
3use mithril_common::{
4    StdResult,
5    entities::{ProtocolMessage, SignedEntityType, SingleSignature},
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.expect_publish().once().returning(|_, _, _| Ok(()));
100
101        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
102
103        retrier
104            .publish(
105                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
106                &fake_data::single_signature(vec![1]),
107                &ProtocolMessage::default(),
108            )
109            .await
110            .unwrap();
111    }
112
113    #[tokio::test]
114    async fn should_not_retry_when_publish_fails_and_retry_policy_is_never() {
115        let retry_policy = SignaturePublishRetryPolicy::never();
116
117        let mut publisher = MockSignaturePublisher::new();
118        publisher
119            .expect_publish()
120            .once()
121            .returning(|_, _, _| Err(anyhow::anyhow!("error while publishing")));
122
123        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
124
125        retrier
126            .publish(
127                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
128                &fake_data::single_signature(vec![1]),
129                &ProtocolMessage::default(),
130            )
131            .await
132            .expect_err("An error should be returned");
133    }
134
135    #[tokio::test]
136    async fn should_retry_once_and_succeed_on_second_attempt() {
137        let retry_policy = SignaturePublishRetryPolicy {
138            attempts: 2,
139            delay_between_attempts: Duration::from_secs(0),
140        };
141
142        let mut publisher = MockSignaturePublisher::new();
143        publisher
144            .expect_publish()
145            .times(1)
146            .return_once(|_, _, _| Err(anyhow::anyhow!("error")));
147        publisher.expect_publish().times(1).return_once(|_, _, _| Ok(()));
148
149        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
150
151        retrier
152            .publish(
153                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
154                &fake_data::single_signature(vec![1]),
155                &ProtocolMessage::default(),
156            )
157            .await
158            .unwrap();
159    }
160
161    #[tokio::test]
162    async fn should_retry_and_return_error_after_max_attempts() {
163        let retry_policy = SignaturePublishRetryPolicy {
164            attempts: 2,
165            delay_between_attempts: Duration::from_secs(0),
166        };
167
168        let mut publisher = MockSignaturePublisher::new();
169        publisher
170            .expect_publish()
171            .times(2)
172            .returning(|_, _, _| Err(anyhow::anyhow!("error")));
173
174        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
175
176        retrier
177            .publish(
178                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
179                &fake_data::single_signature(vec![1]),
180                &ProtocolMessage::default(),
181            )
182            .await
183            .expect_err("An error should be returned after max attempts");
184    }
185
186    #[tokio::test]
187    async fn should_wait_between_retries_according_to_policy() {
188        let delay_between_attempts = Duration::from_millis(50);
189        let retry_policy = SignaturePublishRetryPolicy {
190            attempts: 2,
191            delay_between_attempts,
192        };
193
194        let mut publisher = MockSignaturePublisher::new();
195        publisher
196            .expect_publish()
197            .once()
198            .return_once(|_, _, _| Err(anyhow::anyhow!("error")));
199        publisher.expect_publish().once().return_once(|_, _, _| Ok(()));
200
201        let retrier = SignaturePublisherRetrier::new(Arc::new(publisher), retry_policy);
202
203        let start_time = std::time::Instant::now();
204        retrier
205            .publish(
206                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
207                &fake_data::single_signature(vec![1]),
208                &ProtocolMessage::default(),
209            )
210            .await
211            .unwrap();
212
213        let elapsed_time = start_time.elapsed();
214        assert!(
215            elapsed_time >= delay_between_attempts,
216            "Expected at least {delay_between_attempts:?} time elapsed, but got {elapsed_time:?}"
217        );
218        assert!(
219            elapsed_time < delay_between_attempts * 2,
220            "Expected less than {:?} time elapsed, but got {:?}",
221            delay_between_attempts * 2,
222            elapsed_time
223        );
224    }
225}
mithril_signer/services/signature_publisher/retrier.rs

mithril_signer/services/signature_publisher/
retrier.rs
