async_utility/
time.rs

1// Copyright (c) 2022-2023 Yuki Kishimoto
2// Distributed under the MIT software license
3
4//! Time module
5
6use core::future::Future;
7use core::time::Duration;
8
9use futures_util::future::{AbortHandle, Abortable};
10#[cfg(target_arch = "wasm32")]
11use wasm_bindgen_futures::spawn_local;
12
13#[cfg(not(target_arch = "wasm32"))]
14use crate::runtime;
15
16/// Sleep
17pub async fn sleep(duration: Duration) {
18    #[cfg(not(target_arch = "wasm32"))]
19    if runtime::is_tokio_context() {
20        tokio::time::sleep(duration).await;
21    } else {
22        // No need to propagate error
23        let _ = runtime::handle()
24            .spawn(async move {
25                tokio::time::sleep(duration).await;
26            })
27            .await;
28    }
29
30    #[cfg(target_arch = "wasm32")]
31    gloo_timers::future::sleep(duration).await;
32}
33
34/// Timeout
35pub async fn timeout<F>(timeout: Option<Duration>, future: F) -> Option<F::Output>
36where
37    F: Future,
38{
39    #[cfg(not(target_arch = "wasm32"))]
40    if let Some(timeout) = timeout {
41        if runtime::is_tokio_context() {
42            tokio::time::timeout(timeout, future).await.ok()
43        } else {
44            let (abort_handle, abort_registration) = AbortHandle::new_pair();
45            let future = Abortable::new(future, abort_registration);
46            tokio::select! {
47                res = future => {
48                    res.ok()
49                }
50                _ = sleep(timeout) => {
51                    abort_handle.abort();
52                    None
53                }
54            }
55        }
56    } else {
57        Some(future.await)
58    }
59
60    #[cfg(target_arch = "wasm32")]
61    {
62        if let Some(timeout) = timeout {
63            let (abort_handle, abort_registration) = AbortHandle::new_pair();
64            let future = Abortable::new(future, abort_registration);
65            spawn_local(async move {
66                gloo_timers::callback::Timeout::new(timeout.as_millis() as u32, move || {
67                    abort_handle.abort();
68                })
69                .forget();
70            });
71            future.await.ok()
72        } else {
73            Some(future.await)
74        }
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    // TODO: test also wasm
83
84    #[tokio::test]
85    #[cfg(not(target_arch = "wasm32"))]
86    async fn test_sleep_in_tokio() {
87        sleep(Duration::from_secs(5)).await;
88    }
89
90    #[async_std::test]
91    #[cfg(not(target_arch = "wasm32"))]
92    async fn test_sleep_in_async_std() {
93        sleep(Duration::from_secs(5)).await;
94    }
95
96    #[test]
97    #[cfg(not(target_arch = "wasm32"))]
98    fn test_sleep_in_smol() {
99        smol::block_on(async {
100            sleep(Duration::from_secs(5)).await;
101        });
102    }
103
104    #[tokio::test]
105    #[cfg(not(target_arch = "wasm32"))]
106    async fn test_timeout_tokio() {
107        // Timeout
108        let result = timeout(Some(Duration::from_secs(1)), async {
109            sleep(Duration::from_secs(2)).await;
110        })
111        .await;
112        assert!(result.is_none());
113
114        // Not timeout
115        let result = timeout(Some(Duration::from_secs(10)), async {
116            sleep(Duration::from_secs(1)).await;
117        })
118        .await;
119        assert!(result.is_some());
120    }
121
122    #[async_std::test]
123    #[cfg(not(target_arch = "wasm32"))]
124    async fn test_timeout_async_std() {
125        // Timeout
126        let result = timeout(Some(Duration::from_secs(1)), async {
127            sleep(Duration::from_secs(2)).await;
128        })
129        .await;
130        assert!(result.is_none());
131
132        // Not timeout
133        let result = timeout(Some(Duration::from_secs(10)), async {
134            sleep(Duration::from_secs(1)).await;
135        })
136        .await;
137        assert!(result.is_some());
138    }
139
140    #[test]
141    #[cfg(not(target_arch = "wasm32"))]
142    fn test_timeout_smol() {
143        smol::block_on(async {
144            // Timeout
145            let result = timeout(Some(Duration::from_secs(1)), async {
146                sleep(Duration::from_secs(2)).await;
147            })
148            .await;
149            assert!(result.is_none());
150
151            // Not timeout
152            let result = timeout(Some(Duration::from_secs(10)), async {
153                sleep(Duration::from_secs(1)).await;
154            })
155            .await;
156            assert!(result.is_some());
157        });
158    }
159}
async_utility/time.rs

async_utility/
time.rs
