cuprate_helper/
asynch.rs

1//! `async` related
2//!
3//! `#[no_std]` compatible.
4
5//---------------------------------------------------------------------------------------------------- Use
6use core::{
7    future::Future,
8    pin::Pin,
9    task::{Context, Poll},
10};
11
12use futures::{channel::oneshot, FutureExt};
13
14//---------------------------------------------------------------------------------------------------- InfallibleOneshotReceiver
15/// A oneshot receiver channel that doesn't return an Error.
16///
17/// This requires the sender to always return a response.
18pub struct InfallibleOneshotReceiver<T>(oneshot::Receiver<T>);
19
20impl<T> From<oneshot::Receiver<T>> for InfallibleOneshotReceiver<T> {
21    fn from(value: oneshot::Receiver<T>) -> Self {
22        Self(value)
23    }
24}
25
26impl<T> Future for InfallibleOneshotReceiver<T> {
27    type Output = T;
28
29    #[inline]
30    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
31        self.0
32            .poll_unpin(ctx)
33            .map(|res| res.expect("Oneshot must not be cancelled before response!"))
34    }
35}
36
37//---------------------------------------------------------------------------------------------------- rayon_spawn_async
38/// Spawns a task for the rayon thread pool and awaits the result without blocking the async runtime.
39pub async fn rayon_spawn_async<F, R>(f: F) -> R
40where
41    F: FnOnce() -> R + Send + 'static,
42    R: Send + 'static,
43{
44    let (tx, rx) = oneshot::channel();
45    rayon::spawn(move || {
46        drop(tx.send(f()));
47    });
48    rx.await.expect("The sender must not be dropped")
49}
50
51//---------------------------------------------------------------------------------------------------- Tests
52#[cfg(test)]
53mod test {
54    use std::{
55        sync::{Arc, Barrier},
56        thread,
57        time::Duration,
58    };
59
60    use super::*;
61
62    #[tokio::test]
63    // Assert that basic channel operations work.
64    async fn infallible_oneshot_receiver() {
65        let (tx, rx) = oneshot::channel::<String>();
66        let msg = "hello world!".to_string();
67
68        tx.send(msg.clone()).unwrap();
69
70        let oneshot = InfallibleOneshotReceiver::from(rx);
71        assert_eq!(oneshot.await, msg);
72    }
73
74    #[test]
75    fn rayon_spawn_async_does_not_block() {
76        // There must be more than 1 rayon thread for this to work.
77        rayon::ThreadPoolBuilder::new()
78            .num_threads(2)
79            .build_global()
80            .unwrap();
81
82        // We use a barrier to make sure both tasks are executed together, we block the rayon thread
83        // until both rayon threads are blocked.
84        let barrier = Arc::new(Barrier::new(2));
85        let task = |barrier: &Barrier| barrier.wait();
86
87        let b_2 = Arc::clone(&barrier);
88
89        let (tx, rx) = std::sync::mpsc::channel();
90
91        thread::spawn(move || {
92            let runtime = tokio::runtime::Builder::new_current_thread()
93                .enable_all()
94                .build()
95                .unwrap();
96
97            runtime.block_on(async {
98                tokio::join!(
99                    // This polls them concurrently in the same task, so if the first one blocks the task then
100                    // the second wont run and if the second does not run the first does not unblock.
101                    rayon_spawn_async(move || task(&barrier)),
102                    rayon_spawn_async(move || task(&b_2)),
103                )
104            });
105
106            // if we managed to get here then rayon_spawn_async didn't block.
107            tx.send(()).unwrap();
108        });
109
110        rx.recv_timeout(Duration::from_secs(2))
111            .expect("rayon_spawn_async blocked");
112    }
113}
cuprate_helper/asynch.rs

cuprate_helper/
asynch.rs
