wsio_core/traits/task/
spawner.rs

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

wsio_core/traits/task/
spawner.rs
