Rust Callback Future

rust-cb-future.rs

use std::sync::{Arc, Mutex};
use futures::{
    task::{
        Poll,
        Context,
        Waker,
    },
    Future,
    executor::block_on,
};

struct CbFut<T: 'static + Send + Sync>(Arc<Mutex<(Option<T>, Option<Waker>)>>);

impl<T: 'static + Send + Sync> CbFut<T> {
    pub fn new() -> (Self, Box<dyn FnOnce(T) + 'static + Send + Sync>) {
        let inner: Arc<Mutex<(Option<T>, Option<Waker>)>> =
            Arc::new(Mutex::new((None, None)));
        let inner_clone = inner.clone();
        (
            Self(inner),
            Box::new(move |t: T| {
                let mut l = inner_clone.lock().unwrap();
                l.0 = Some(t);
                if let Some(w) = l.1.take() {
                    w.wake();
                }
            }),
        )
    }
}

impl<T: 'static + Send + Sync> Future for CbFut<T> {
    type Output = T;
    
    fn poll(
        self: std::pin::Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Self::Output> {
        let mut l = self.0.lock().unwrap();
        match &mut l.0 {
            s @ Some(_) => Poll::Ready(s.take().unwrap()),
            None => {
                l.1 = Some(cx.waker().clone());
                Poll::Pending
            }
        }
    }
}

async fn test() -> String {
    let (b, cb) = <CbFut<String>>::new();
    std::thread::spawn(move || {
        std::thread::sleep(std::time::Duration::from_millis(1000));
        cb("test-cb".to_string());
    });
    b.await
}

async fn async_main() {
    println!("{}", test().await);
}

fn main() {
    println!("{:?} b4", std::time::SystemTime::now());
    block_on(async_main());
    println!("{:?} done", std::time::SystemTime::now());
}