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

wsio_core/traits/task/
spawner.rs
