hotshot_testing/
test_task.rs

1// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
2// This file is part of the HotShot repository.
3
4// You should have received a copy of the MIT License
5// along with the HotShot repository. If not, see <https://mit-license.org/>.
6
7use std::{sync::Arc, time::Duration};
8
9use async_broadcast::{Receiver, Sender};
10use async_lock::RwLock;
11use async_trait::async_trait;
12use futures::future::select_all;
13use hotshot::{
14    traits::TestableNodeImplementation,
15    types::{Event, Message},
16};
17use hotshot_task_impls::{events::HotShotEvent, network::NetworkMessageTaskState};
18use hotshot_types::{
19    message::UpgradeLock,
20    traits::{network::ConnectedNetwork, node_implementation::NodeType},
21};
22use tokio::{
23    spawn,
24    task::JoinHandle,
25    time::{sleep, timeout},
26};
27use tracing::error;
28
29use crate::test_runner::Node;
30
31/// enum describing how the tasks completed
32pub enum TestResult {
33    /// the test task passed
34    Pass,
35    /// the test task failed with an error
36    Fail(Box<dyn std::fmt::Debug + Send + Sync>),
37}
38
39pub fn spawn_timeout_task(test_sender: Sender<TestEvent>, timeout: Duration) -> JoinHandle<()> {
40    tokio::spawn(async move {
41        sleep(timeout).await;
42
43        error!("Decide timeout triggered.");
44        let _ = test_sender.broadcast(TestEvent::Shutdown).await;
45    })
46}
47
48#[async_trait]
49/// Type for mutable task state that can be used as the state for a `Task`
50pub trait TestTaskState: Send {
51    /// Type of event sent and received by the task
52    type Event: Clone + Send + Sync;
53    /// Type of error produced by the task
54    type Error: std::fmt::Display;
55
56    /// Handles an event from one of multiple receivers.
57    async fn handle_event(
58        &mut self,
59        (event, id): (Self::Event, usize),
60    ) -> std::result::Result<(), Self::Error>;
61
62    /// Check the result of the test.
63    async fn check(&self) -> TestResult;
64}
65
66/// Type alias for type-erased [`TestTaskState`] to be used for
67/// dynamic dispatch
68pub type AnyTestTaskState<TYPES> = Box<
69    dyn TestTaskState<Event = hotshot_types::event::Event<TYPES>, Error = anyhow::Error>
70        + Send
71        + Sync,
72>;
73
74#[async_trait]
75impl<TYPES: NodeType> TestTaskState for AnyTestTaskState<TYPES> {
76    type Event = Event<TYPES>;
77    type Error = anyhow::Error;
78
79    async fn handle_event(
80        &mut self,
81        event: (Self::Event, usize),
82    ) -> std::result::Result<(), anyhow::Error> {
83        (**self).handle_event(event).await
84    }
85
86    async fn check(&self) -> TestResult {
87        (**self).check().await
88    }
89}
90
91#[async_trait]
92pub trait TestTaskStateSeed<TYPES, I>: Send
93where
94    TYPES: NodeType,
95    I: TestableNodeImplementation<TYPES>,
96{
97    async fn into_state(
98        self: Box<Self>,
99        handles: Arc<RwLock<Vec<Node<TYPES, I>>>>,
100    ) -> AnyTestTaskState<TYPES>;
101}
102
103/// A basic task which loops waiting for events to come from `event_receiver`
104/// and then handles them using it's state
105/// It sends events to other `Task`s through `event_sender`
106/// This should be used as the primary building block for long running
107/// or medium running tasks (i.e. anything that can't be described as a dependency task)
108pub struct TestTask<S: TestTaskState> {
109    /// The state of the task.  It is fed events from `event_sender`
110    /// and mutates it state ocordingly.  Also it signals the task
111    /// if it is complete/should shutdown
112    state: S,
113    /// Receives events that are broadcast from any task, including itself
114    receivers: Vec<Receiver<S::Event>>,
115    /// Receiver for test events, used for communication between test tasks.
116    test_receiver: Receiver<TestEvent>,
117}
118
119#[derive(Clone, Debug)]
120pub enum TestEvent {
121    Shutdown,
122}
123
124impl std::fmt::Display for TestEvent {
125    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126        match self {
127            TestEvent::Shutdown => write!(f, "Shutdown"),
128        }
129    }
130}
131
132impl<S: TestTaskState + Send + 'static> TestTask<S> {
133    /// Create a new task
134    pub fn new(
135        state: S,
136        receivers: Vec<Receiver<S::Event>>,
137        test_receiver: Receiver<TestEvent>,
138    ) -> Self {
139        TestTask {
140            state,
141            receivers,
142            test_receiver,
143        }
144    }
145
146    /// Spawn the task loop, consuming self.  Will continue until
147    /// the task reaches some shutdown condition
148    pub fn run(mut self) -> JoinHandle<TestResult> {
149        spawn(async move {
150            loop {
151                if let Ok(TestEvent::Shutdown) = self.test_receiver.try_recv() {
152                    break self.state.check().await;
153                }
154
155                self.receivers.retain(|receiver| !receiver.is_closed());
156
157                // All-node restart closes every receiver at once; `select_all` panics on
158                // an empty set. Receivers are never re-subscribed, so idle until shutdown.
159                if self.receivers.is_empty() {
160                    sleep(Duration::from_millis(500)).await;
161                    continue;
162                }
163
164                let mut messages = Vec::new();
165
166                for receiver in &mut self.receivers {
167                    messages.push(receiver.recv());
168                }
169
170                match timeout(Duration::from_millis(2500), select_all(messages)).await {
171                    Ok((Ok(input), id, _)) => {
172                        let _ = S::handle_event(&mut self.state, (input, id))
173                            .await
174                            .inspect_err(|e| tracing::error!("{e}"));
175                    },
176                    Ok((Err(e), _id, _)) => {
177                        error!("Error from one channel in test task {e:?}");
178                        sleep(Duration::from_millis(4000)).await;
179                    },
180                    _ => {},
181                };
182            }
183        })
184    }
185}
186
187/// Add the network task to handle messages and publish events.
188pub async fn add_network_message_test_task<
189    TYPES: NodeType,
190    NET: ConnectedNetwork<TYPES::SignatureKey>,
191>(
192    internal_event_stream: Sender<Arc<HotShotEvent<TYPES>>>,
193    external_event_stream: Sender<Event<TYPES>>,
194    upgrade_lock: UpgradeLock<TYPES>,
195    channel: Arc<NET>,
196    public_key: TYPES::SignatureKey,
197    id: u64,
198) -> JoinHandle<()> {
199    let net = Arc::clone(&channel);
200    let network_state: NetworkMessageTaskState<_> = NetworkMessageTaskState {
201        internal_event_stream: internal_event_stream.clone(),
202        external_event_stream: external_event_stream.clone(),
203        public_key,
204        upgrade_lock: upgrade_lock.clone(),
205        id,
206    };
207
208    let network = Arc::clone(&net);
209    let mut state = network_state.clone();
210
211    spawn(async move {
212        loop {
213            // Get the next message from the network
214            let message = match network.recv_message().await {
215                Ok(message) => message,
216                Err(e) => {
217                    error!("Failed to receive message: {e:?}");
218                    continue;
219                },
220            };
221
222            // Deserialize the message
223            let deserialized_message: Message<TYPES> = match upgrade_lock.deserialize(&message) {
224                Ok((message, _)) => message,
225                Err(e) => {
226                    tracing::error!("Failed to deserialize message: {e:?}");
227                    continue;
228                },
229            };
230
231            // Handle the message
232            state.handle_message(deserialized_message).await;
233        }
234    })
235}
hotshot_testing/test_task.rs

hotshot_testing/
test_task.rs
