hotshot_testing/
completion_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::time::Duration;
8
9use async_broadcast::{Receiver, Sender};
10use hotshot_task_impls::helpers::broadcast_event;
11use tokio::{spawn, task::JoinHandle, time::timeout};
12
13use crate::test_task::TestEvent;
14
15/// Completion task state
16pub struct CompletionTask {
17    pub tx: Sender<TestEvent>,
18
19    pub rx: Receiver<TestEvent>,
20    /// Duration of the task.
21    pub duration: Duration,
22}
23
24impl CompletionTask {
25    pub fn run(mut self) -> JoinHandle<()> {
26        spawn(async move {
27            if timeout(self.duration, self.wait_for_shutdown())
28                .await
29                .is_err()
30            {
31                tracing::error!("Completion Task timed out");
32                broadcast_event(TestEvent::Shutdown, &self.tx).await;
33            }
34        })
35    }
36    async fn wait_for_shutdown(&mut self) {
37        while let Ok(event) = self.rx.recv_direct().await {
38            if matches!(event, TestEvent::Shutdown) {
39                tracing::error!("Completion Task shutting down");
40                return;
41            }
42        }
43    }
44}
45/// Description for a time-based completion task.
46#[derive(Clone, Debug)]
47pub struct TimeBasedCompletionTaskDescription {
48    /// Duration of the task.
49    pub duration: Duration,
50}
51
52/// Description for a completion task.
53#[derive(Clone, Debug)]
54pub enum CompletionTaskDescription {
55    /// Time-based completion task.
56    TimeBasedCompletionTaskBuilder(TimeBasedCompletionTaskDescription),
57}