Skip to main content

docbox_core/utils/
timing.rs

1use tokio::time::{Duration, sleep};
2
3/// Helper to run a callback when a future is taking longer than a `slow_duration`
4pub async fn handle_slow_future<Fut, T, F>(future: Fut, slow_duration: Duration, callback: F) -> T
5where
6    Fut: Future<Output = T>,
7    F: FnOnce(),
8{
9    tokio::pin!(future);
10
11    let mut slow_callback = Some(callback);
12
13    loop {
14        tokio::select! {
15            result = &mut future => return result,
16            _ = sleep(slow_duration), if slow_callback.is_some() => {
17
18                if let Some(slow_callback) = slow_callback.take() {
19                    slow_callback();
20                }
21            }
22        }
23    }
24}
25
26#[cfg(test)]
27mod test {
28    use crate::utils::timing::handle_slow_future;
29    use std::{sync::atomic::AtomicBool, time::Duration};
30    use tokio::time::sleep;
31
32    /// Tests that a slow future will trigger the callback
33    #[tokio::test]
34    async fn test_slow_future() {
35        let slow = AtomicBool::new(false);
36
37        let slow_future = sleep(Duration::from_secs(5));
38
39        handle_slow_future(slow_future, Duration::from_secs(2), || {
40            slow.store(true, std::sync::atomic::Ordering::SeqCst);
41        })
42        .await;
43
44        assert!(slow.load(std::sync::atomic::Ordering::SeqCst),);
45    }
46
47    /// Tests that a fast future will not trigger the callback
48    #[tokio::test]
49    async fn test_fast_future() {
50        let slow = AtomicBool::new(false);
51
52        let fast_future = sleep(Duration::from_millis(1));
53
54        handle_slow_future(fast_future, Duration::from_secs(2), || {
55            slow.store(true, std::sync::atomic::Ordering::SeqCst);
56        })
57        .await;
58
59        assert!(!slow.load(std::sync::atomic::Ordering::SeqCst));
60    }
61}