hotshot_testing/
completion_task.rs1use 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
15pub struct CompletionTask {
17 pub tx: Sender<TestEvent>,
18
19 pub rx: Receiver<TestEvent>,
20 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#[derive(Clone, Debug)]
47pub struct TimeBasedCompletionTaskDescription {
48 pub duration: Duration,
50}
51
52#[derive(Clone, Debug)]
54pub enum CompletionTaskDescription {
55 TimeBasedCompletionTaskBuilder(TimeBasedCompletionTaskDescription),
57}