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

mithril_signer/services/signature_publisher/
retrier.rs
