wsio_core/traits/task/
spawner.rs

1use std::sync::Arc;
2
3use anyhow::Result;
4use tokio::{
5    select,
6    spawn,
7};
8use tokio_util::sync::CancellationToken;
9
10pub trait TaskSpawner: Send + Sync + 'static {
11    fn cancel_token(&self) -> Arc<CancellationToken>;
12
13    #[inline]
14    fn spawn_task<F: Future<Output = Result<()>> + Send + 'static>(&self, future: F) {
15        let cancel_token = self.cancel_token();
16        spawn(async move {
17            select! {
18                _ = cancel_token.cancelled() => {},
19                _ = future => {},
20            }
21        });
22    }
23}
24
25#[cfg(test)]
26mod tests {
27    use std::{
28        future::pending,
29        time::Duration,
30    };
31
32    use tokio::{
33        sync::oneshot::channel,
34        time::timeout,
35    };
36
37    use super::*;
38
39    struct TestSpawner {
40        cancel_token: Arc<CancellationToken>,
41    }
42
43    impl TaskSpawner for TestSpawner {
44        fn cancel_token(&self) -> Arc<CancellationToken> {
45            self.cancel_token.clone()
46        }
47    }
48
49    #[tokio::test]
50    async fn test_spawn_task_runs_to_completion() {
51        let spawner = TestSpawner {
52            cancel_token: Arc::new(CancellationToken::new()),
53        };
54
55        let (completed_tx, completed_rx) = channel();
56
57        spawner.spawn_task(async move {
58            let _ = completed_tx.send(());
59            Ok(())
60        });
61
62        timeout(Duration::from_secs(1), completed_rx)
63            .await
64            .expect("task should complete before timeout")
65            .expect("task should run to completion");
66    }
67
68    #[tokio::test]
69    async fn test_spawn_task_is_cancelled() {
70        let cancel_token = Arc::new(CancellationToken::new());
71        let spawner = TestSpawner {
72            cancel_token: cancel_token.clone(),
73        };
74
75        let (started_tx, started_rx) = channel();
76        let (dropped_tx, dropped_rx) = channel::<()>();
77
78        spawner.spawn_task(async move {
79            let _drop_signal = dropped_tx;
80            let _ = started_tx.send(());
81            pending::<()>().await;
82            Ok(())
83        });
84
85        timeout(Duration::from_secs(1), started_rx)
86            .await
87            .expect("task should start before timeout")
88            .expect("task should start before cancellation");
89
90        cancel_token.cancel();
91
92        timeout(Duration::from_secs(1), dropped_rx)
93            .await
94            .expect("task should be dropped before timeout")
95            .expect_err("cancellation should drop the task future");
96    }
97}
wsio_core/traits/task/spawner.rs

wsio_core/traits/task/
spawner.rs
