mithril_signer/services/signature_publisher/
delayer.rs

1use std::{sync::Arc, time::Duration};
2
3use mithril_common::{
4    StdResult,
5    entities::{ProtocolMessage, SignedEntityType, SingleSignature},
6    logging::LoggerExtensions,
7};
8use slog::{Logger, error};
9
10use super::SignaturePublisher;
11
12/// A decorator of [SignaturePublisher] that publishes right away on a first publisher
13/// and with a delay on the second publisher.
14pub struct SignaturePublisherDelayer {
15    first_publisher: Arc<dyn SignaturePublisher>,
16    second_publisher: Arc<dyn SignaturePublisher>,
17    delay_between_publish: Duration,
18    logger: Logger,
19}
20
21impl SignaturePublisherDelayer {
22    /// Creates a new [SignaturePublisherDelayer]
23    pub fn new(
24        first_publisher: Arc<dyn SignaturePublisher>,
25        second_publisher: Arc<dyn SignaturePublisher>,
26        delay_between_publish: Duration,
27        logger: Logger,
28    ) -> Self {
29        Self {
30            first_publisher,
31            second_publisher,
32            delay_between_publish,
33            logger: logger.new_with_component_name::<Self>(),
34        }
35    }
36}
37
38#[async_trait::async_trait]
39impl SignaturePublisher for SignaturePublisherDelayer {
40    async fn publish(
41        &self,
42        signed_entity_type: &SignedEntityType,
43        signature: &SingleSignature,
44        protocol_message: &ProtocolMessage,
45    ) -> StdResult<()> {
46        if let Err(e) = self
47            .first_publisher
48            .publish(signed_entity_type, signature, protocol_message)
49            .await
50        {
51            error!(
52                self.logger,
53                "Delayer failed to publish first signature";
54                "error" => ?e
55            );
56        }
57
58        tokio::time::sleep(self.delay_between_publish).await;
59
60        if let Err(e) = self
61            .second_publisher
62            .publish(signed_entity_type, signature, protocol_message)
63            .await
64        {
65            error!(
66                self.logger,
67                "Delayer failed to publish second signature";
68                "error" => ?e
69            );
70            return Err(e);
71        }
72
73        Ok(())
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use mithril_common::{entities::Epoch, test_utils::fake_data};
80
81    use crate::{services::MockSignaturePublisher, test_tools::TestLogger};
82
83    use super::*;
84
85    #[tokio::test]
86    async fn should_call_both_publishers_when_first_succeeds() {
87        let mut first_publisher = MockSignaturePublisher::new();
88        first_publisher.expect_publish().once().returning(|_, _, _| Ok(()));
89
90        let mut second_publisher = MockSignaturePublisher::new();
91        second_publisher.expect_publish().once().returning(|_, _, _| Ok(()));
92
93        let delayer = SignaturePublisherDelayer::new(
94            Arc::new(first_publisher),
95            Arc::new(second_publisher),
96            Duration::from_millis(0),
97            TestLogger::stdout(),
98        );
99
100        delayer
101            .publish(
102                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
103                &fake_data::single_signature(vec![1]),
104                &ProtocolMessage::default(),
105            )
106            .await
107            .unwrap();
108    }
109
110    #[tokio::test]
111    async fn should_call_second_publisher_even_if_first_fails_and_log_error() {
112        let (logger, log_inspector) = TestLogger::memory();
113        let mut first_publisher = MockSignaturePublisher::new();
114        first_publisher
115            .expect_publish()
116            .once()
117            .returning(|_, _, _| Err(anyhow::anyhow!("first publisher failure")));
118
119        let mut second_publisher = MockSignaturePublisher::new();
120        second_publisher.expect_publish().once().returning(|_, _, _| Ok(()));
121
122        let delayer = SignaturePublisherDelayer::new(
123            Arc::new(first_publisher),
124            Arc::new(second_publisher),
125            Duration::from_millis(0),
126            logger,
127        );
128
129        delayer
130            .publish(
131                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
132                &fake_data::single_signature(vec![1]),
133                &ProtocolMessage::default(),
134            )
135            .await
136            .unwrap();
137
138        assert!(log_inspector.contains_log("Delayer failed to publish first signature"));
139        assert!(log_inspector.contains_log("first publisher failure"));
140    }
141
142    #[tokio::test]
143    async fn should_return_and_log_error_if_second_publisher_fails() {
144        let (logger, log_inspector) = TestLogger::memory();
145        let mut first_publisher = MockSignaturePublisher::new();
146        first_publisher.expect_publish().once().returning(|_, _, _| Ok(()));
147
148        let mut second_publisher = MockSignaturePublisher::new();
149        second_publisher
150            .expect_publish()
151            .once()
152            .returning(|_, _, _| Err(anyhow::anyhow!("second publisher failure")));
153
154        let delayer = SignaturePublisherDelayer::new(
155            Arc::new(first_publisher),
156            Arc::new(second_publisher),
157            Duration::from_millis(0),
158            logger,
159        );
160
161        delayer
162            .publish(
163                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
164                &fake_data::single_signature(vec![1]),
165                &ProtocolMessage::default(),
166            )
167            .await
168            .expect_err("Expected error when delayed publisher failed");
169
170        assert!(log_inspector.contains_log("Delayer failed to publish second signature"));
171        assert!(log_inspector.contains_log("second publisher failure"));
172    }
173
174    #[tokio::test]
175    async fn should_wait_before_calling_second_publisher() {
176        let mut first_publisher = MockSignaturePublisher::new();
177        first_publisher.expect_publish().once().returning(|_, _, _| Ok(()));
178
179        let mut second_publisher = MockSignaturePublisher::new();
180        second_publisher.expect_publish().once().returning(|_, _, _| Ok(()));
181
182        let delay = Duration::from_millis(50);
183        let delayer = SignaturePublisherDelayer::new(
184            Arc::new(first_publisher),
185            Arc::new(second_publisher),
186            delay,
187            TestLogger::stdout(),
188        );
189
190        let start_time = std::time::Instant::now();
191        delayer
192            .publish(
193                &SignedEntityType::MithrilStakeDistribution(Epoch(1)),
194                &fake_data::single_signature(vec![1]),
195                &ProtocolMessage::default(),
196            )
197            .await
198            .unwrap();
199
200        let elapsed_time = start_time.elapsed();
201        assert!(
202            elapsed_time >= delay,
203            "Expected at least {delay:?} time elapsed, but got {elapsed_time:?}"
204        );
205        assert!(
206            elapsed_time < delay * 2,
207            "Expected less than {:?} time elapsed, but got {:?}",
208            delay * 2,
209            elapsed_time
210        );
211    }
212}