Web Async Rework (#3082)

src/platform_impl/web/async/channel.rs

@@ -0,0 +1,115 @@
+use atomic_waker::AtomicWaker;
+use std::future;
+use std::rc::Rc;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::mpsc::{self, Receiver, RecvError, SendError, Sender, TryRecvError};
+use std::sync::{Arc, Mutex};
+use std::task::Poll;
+
+// NOTE: This channel doesn't wake up when all senders or receivers are
+// dropped. This is acceptable as long as it's only used in `Dispatcher`, which
+// has it's own `Drop` behavior.
+
+pub fn channel<T>() -> (AsyncSender<T>, AsyncReceiver<T>) {
+    let (sender, receiver) = mpsc::channel();
+    let sender = Arc::new(Mutex::new(sender));
+    let inner = Arc::new(Inner {
+        closed: AtomicBool::new(false),
+        waker: AtomicWaker::new(),
+    });
+
+    let sender = AsyncSender {
+        sender,
+        inner: Arc::clone(&inner),
+    };
+    let receiver = AsyncReceiver {
+        receiver: Rc::new(receiver),
+        inner,
+    };
+
+    (sender, receiver)
+}
+
+pub struct AsyncSender<T> {
+    // We need to wrap it into a `Mutex` to make it `Sync`. So the sender can't
+    // be accessed on the main thread, as it could block. Additionally we need
+    // to wrap it in an `Arc` to make it clonable on the main thread without
+    // having to block.
+    sender: Arc<Mutex<Sender<T>>>,
+    inner: Arc<Inner>,
+}
+
+impl<T> AsyncSender<T> {
+    pub fn send(&self, event: T) -> Result<(), SendError<T>> {
+        self.sender.lock().unwrap().send(event)?;
+        self.inner.waker.wake();
+
+        Ok(())
+    }
+
+    pub fn close(&self) {
+        self.inner.closed.store(true, Ordering::Relaxed);
+        self.inner.waker.wake()
+    }
+}
+
+impl<T> Clone for AsyncSender<T> {
+    fn clone(&self) -> Self {
+        Self {
+            sender: Arc::clone(&self.sender),
+            inner: Arc::clone(&self.inner),
+        }
+    }
+}
+
+pub struct AsyncReceiver<T> {
+    receiver: Rc<Receiver<T>>,
+    inner: Arc<Inner>,
+}
+
+impl<T> AsyncReceiver<T> {
+    pub async fn next(&self) -> Result<T, RecvError> {
+        future::poll_fn(|cx| match self.receiver.try_recv() {
+            Ok(event) => Poll::Ready(Ok(event)),
+            Err(TryRecvError::Empty) => {
+                self.inner.waker.register(cx.waker());
+
+                match self.receiver.try_recv() {
+                    Ok(event) => Poll::Ready(Ok(event)),
+                    Err(TryRecvError::Empty) => {
+                        if self.inner.closed.load(Ordering::Relaxed) {
+                            Poll::Ready(Err(RecvError))
+                        } else {
+                            Poll::Pending
+                        }
+                    }
+                    Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError)),
+                }
+            }
+            Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError)),
+        })
+        .await
+    }
+
+    pub fn try_recv(&self) -> Result<Option<T>, RecvError> {
+        match self.receiver.try_recv() {
+            Ok(value) => Ok(Some(value)),
+            Err(TryRecvError::Empty) => Ok(None),
+            Err(TryRecvError::Disconnected) => Err(RecvError),
+        }
+    }
+}
+
+impl<T> Clone for AsyncReceiver<T> {
+    fn clone(&self) -> Self {
+        Self {
+            receiver: Rc::clone(&self.receiver),
+            inner: Arc::clone(&self.inner),
+        }
+    }
+}
+
+struct Inner {
+    closed: AtomicBool,
+    waker: AtomicWaker,
+}

src/platform_impl/web/async/dispatcher.rs

@@ -0,0 +1,113 @@
+use super::{channel, AsyncReceiver, AsyncSender, Wrapper};
+use std::{
+    cell::Ref,
+    sync::{Arc, Condvar, Mutex},
+};
+
+pub struct Dispatcher<T: 'static>(Wrapper<true, T, AsyncSender<Closure<T>>, Closure<T>>);
+
+struct Closure<T>(Box<dyn FnOnce(&T) + Send>);
+
+impl<T> Dispatcher<T> {
+    #[track_caller]
+    pub fn new(value: T) -> Option<(Self, DispatchRunner<T>)> {
+        let (sender, receiver) = channel::<Closure<T>>();
+
+        Wrapper::new(
+            value,
+            |value, Closure(closure)| {
+                // SAFETY: The given `Closure` here isn't really `'static`, so we shouldn't do anything
+                // funny with it here. See `Self::queue()`.
+                closure(value.borrow().as_ref().unwrap())
+            },
+            {
+                let receiver = receiver.clone();
+                move |value| async move {
+                    while let Ok(Closure(closure)) = receiver.next().await {
+                        // SAFETY: The given `Closure` here isn't really `'static`, so we shouldn't do anything
+                        // funny with it here. See `Self::queue()`.
+                        closure(value.borrow().as_ref().unwrap())
+                    }
+                }
+            },
+            sender,
+            |sender, closure| {
+                // SAFETY: The given `Closure` here isn't really `'static`, so we shouldn't do anything
+                // funny with it here. See `Self::queue()`.
+                sender.send(closure).unwrap()
+            },
+        )
+        .map(|wrapper| (Self(wrapper.clone()), DispatchRunner { wrapper, receiver }))
+    }
+
+    pub fn value(&self) -> Option<Ref<'_, T>> {
+        self.0.value()
+    }
+
+    pub fn dispatch(&self, f: impl 'static + FnOnce(&T) + Send) {
+        if let Some(value) = self.0.value() {
+            f(&value)
+        } else {
+            self.0.send(Closure(Box::new(f)))
+        }
+    }
+
+    pub fn queue<R: Send>(&self, f: impl FnOnce(&T) -> R + Send) -> R {
+        if let Some(value) = self.0.value() {
+            f(&value)
+        } else {
+            let pair = Arc::new((Mutex::new(None), Condvar::new()));
+            let closure = Box::new({
+                let pair = pair.clone();
+                move |value: &T| {
+                    *pair.0.lock().unwrap() = Some(f(value));
+                    pair.1.notify_one();
+                }
+            }) as Box<dyn FnOnce(&T) + Send>;
+            // SAFETY: The `transmute` is necessary because `Closure` requires `'static`. This is
+            // safe because this function won't return until `f` has finished executing. See
+            // `Self::new()`.
+            let closure = Closure(unsafe { std::mem::transmute(closure) });
+
+            self.0.send(closure);
+
+            let mut started = pair.0.lock().unwrap();
+
+            while started.is_none() {
+                started = pair.1.wait(started).unwrap();
+            }
+
+            started.take().unwrap()
+        }
+    }
+}
+
+impl<T> Drop for Dispatcher<T> {
+    fn drop(&mut self) {
+        self.0.with_sender_data(|sender| sender.close())
+    }
+}
+
+pub struct DispatchRunner<T: 'static> {
+    wrapper: Wrapper<true, T, AsyncSender<Closure<T>>, Closure<T>>,
+    receiver: AsyncReceiver<Closure<T>>,
+}
+
+impl<T> DispatchRunner<T> {
+    pub fn run(&self) {
+        while let Some(Closure(closure)) = self
+            .receiver
+            .try_recv()
+            .expect("should only be closed when `Dispatcher` is dropped")
+        {
+            // SAFETY: The given `Closure` here isn't really `'static`, so we shouldn't do anything
+            // funny with it here. See `Self::queue()`.
+            closure(
+                &self
+                    .wrapper
+                    .value()
+                    .expect("don't call this outside the main thread"),
+            )
+        }
+    }
+}

src/platform_impl/web/async/mod.rs

@@ -0,0 +1,9 @@
+mod channel;
+mod dispatcher;
+mod waker;
+mod wrapper;
+
+use self::channel::{channel, AsyncReceiver, AsyncSender};
+pub use self::dispatcher::{DispatchRunner, Dispatcher};
+pub use self::waker::{Waker, WakerSpawner};
+use self::wrapper::Wrapper;

src/platform_impl/web/async/waker.rs

@@ -0,0 +1,123 @@
+use super::Wrapper;
+use atomic_waker::AtomicWaker;
+use std::future;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::task::Poll;
+
+pub struct WakerSpawner<T: 'static>(Wrapper<false, Handler<T>, Sender, usize>);
+
+pub struct Waker<T: 'static>(Wrapper<false, Handler<T>, Sender, usize>);
+
+struct Handler<T> {
+    value: T,
+    handler: fn(&T, usize),
+}
+
+#[derive(Clone)]
+struct Sender(Arc<Inner>);
+
+impl<T> WakerSpawner<T> {
+    #[track_caller]
+    pub fn new(value: T, handler: fn(&T, usize)) -> Option<Self> {
+        let inner = Arc::new(Inner {
+            counter: AtomicUsize::new(0),
+            waker: AtomicWaker::new(),
+            closed: AtomicBool::new(false),
+        });
+
+        let handler = Handler { value, handler };
+
+        let sender = Sender(Arc::clone(&inner));
+
+        let wrapper = Wrapper::new(
+            handler,
+            |handler, count| {
+                let handler = handler.borrow();
+                let handler = handler.as_ref().unwrap();
+                (handler.handler)(&handler.value, count);
+            },
+            {
+                let inner = Arc::clone(&inner);
+
+                move |handler| async move {
+                    while let Some(count) = future::poll_fn(|cx| {
+                        let count = inner.counter.swap(0, Ordering::Relaxed);
+
+                        if count > 0 {
+                            Poll::Ready(Some(count))
+                        } else {
+                            inner.waker.register(cx.waker());
+
+                            let count = inner.counter.swap(0, Ordering::Relaxed);
+
+                            if count > 0 {
+                                Poll::Ready(Some(count))
+                            } else {
+                                if inner.closed.load(Ordering::Relaxed) {
+                                    return Poll::Ready(None);
+                                }
+
+                                Poll::Pending
+                            }
+                        }
+                    })
+                    .await
+                    {
+                        let handler = handler.borrow();
+                        let handler = handler.as_ref().unwrap();
+                        (handler.handler)(&handler.value, count);
+                    }
+                }
+            },
+            sender,
+            |inner, _| {
+                inner.0.counter.fetch_add(1, Ordering::Relaxed);
+                inner.0.waker.wake();
+            },
+        )?;
+
+        Some(Self(wrapper))
+    }
+
+    pub fn waker(&self) -> Waker<T> {
+        Waker(self.0.clone())
+    }
+
+    pub fn fetch(&self) -> usize {
+        debug_assert!(
+            self.0.is_main_thread(),
+            "this should only be called from the main thread"
+        );
+
+        self.0
+            .with_sender_data(|inner| inner.0.counter.swap(0, Ordering::Relaxed))
+    }
+}
+
+impl<T> Drop for WakerSpawner<T> {
+    fn drop(&mut self) {
+        self.0.with_sender_data(|inner| {
+            inner.0.closed.store(true, Ordering::Relaxed);
+            inner.0.waker.wake();
+        });
+    }
+}
+
+impl<T> Waker<T> {
+    pub fn wake(&self) {
+        self.0.send(1)
+    }
+}
+
+impl<T> Clone for Waker<T> {
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+struct Inner {
+    counter: AtomicUsize,
+    waker: AtomicWaker,
+    closed: AtomicBool,
+}

src/platform_impl/web/async/wrapper.rs

@@ -0,0 +1,131 @@
+use std::cell::{Ref, RefCell};
+use std::future::Future;
+use std::marker::PhantomData;
+use std::sync::Arc;
+use wasm_bindgen::prelude::wasm_bindgen;
+use wasm_bindgen::{JsCast, JsValue};
+
+// Unsafe wrapper type that allows us to use `T` when it's not `Send` from other threads.
+// `value` **must** only be accessed on the main thread.
+pub struct Wrapper<const SYNC: bool, V: 'static, S: Clone + Send, E> {
+    value: Value<SYNC, V>,
+    handler: fn(&RefCell<Option<V>>, E),
+    sender_data: S,
+    sender_handler: fn(&S, E),
+}
+
+struct Value<const SYNC: bool, V> {
+    // SAFETY:
+    // This value must not be accessed if not on the main thread.
+    //
+    // - We wrap this in an `Arc` to allow it to be safely cloned without
+    //   accessing the value.
+    // - The `RefCell` lets us mutably access in the main thread but is safe to
+    //   drop in any thread because it has no `Drop` behavior.
+    // - The `Option` lets us safely drop `T` only in the main thread.
+    value: Arc<RefCell<Option<V>>>,
+    // Prevent's `Send` or `Sync` to be automatically implemented.
+    local: PhantomData<*const ()>,
+}
+
+// SAFETY: See `Self::value`.
+unsafe impl<const SYNC: bool, V> Send for Value<SYNC, V> {}
+// SAFETY: See `Self::value`.
+unsafe impl<V> Sync for Value<true, V> {}
+
+impl<const SYNC: bool, V, S: Clone + Send, E> Wrapper<SYNC, V, S, E> {
+    thread_local! {
+        static MAIN_THREAD: bool = {
+            #[wasm_bindgen]
+            extern "C" {
+                #[derive(Clone)]
+                type Global;
+
+                #[wasm_bindgen(method, getter, js_name = Window)]
+                fn window(this: &Global) -> JsValue;
+            }
+
+            let global: Global = js_sys::global().unchecked_into();
+            !global.window().is_undefined()
+        };
+    }
+
+    #[track_caller]
+    pub fn new<R: Future<Output = ()>>(
+        value: V,
+        handler: fn(&RefCell<Option<V>>, E),
+        receiver: impl 'static + FnOnce(Arc<RefCell<Option<V>>>) -> R,
+        sender_data: S,
+        sender_handler: fn(&S, E),
+    ) -> Option<Self> {
+        Self::MAIN_THREAD.with(|safe| {
+            if !safe {
+                panic!("only callable from inside the `Window`")
+            }
+        });
+
+        let value = Arc::new(RefCell::new(Some(value)));
+
+        wasm_bindgen_futures::spawn_local({
+            let value = Arc::clone(&value);
+            async move {
+                receiver(Arc::clone(&value)).await;
+                drop(value.borrow_mut().take().unwrap());
+            }
+        });
+
+        Some(Self {
+            value: Value {
+                value,
+                local: PhantomData,
+            },
+            handler,
+            sender_data,
+            sender_handler,
+        })
+    }
+
+    pub fn send(&self, event: E) {
+        Self::MAIN_THREAD.with(|is_main_thread| {
+            if *is_main_thread {
+                (self.handler)(&self.value.value, event)
+            } else {
+                (self.sender_handler)(&self.sender_data, event)
+            }
+        })
+    }
+
+    pub fn is_main_thread(&self) -> bool {
+        Self::MAIN_THREAD.with(|is_main_thread| *is_main_thread)
+    }
+
+    pub fn value(&self) -> Option<Ref<'_, V>> {
+        Self::MAIN_THREAD.with(|is_main_thread| {
+            if *is_main_thread {
+                Some(Ref::map(self.value.value.borrow(), |value| {
+                    value.as_ref().unwrap()
+                }))
+            } else {
+                None
+            }
+        })
+    }
+
+    pub fn with_sender_data<T>(&self, f: impl FnOnce(&S) -> T) -> T {
+        f(&self.sender_data)
+    }
+}
+
+impl<const SYNC: bool, V, S: Clone + Send, E> Clone for Wrapper<SYNC, V, S, E> {
+    fn clone(&self) -> Self {
+        Self {
+            value: Value {
+                value: self.value.value.clone(),
+                local: PhantomData,
+            },
+            handler: self.handler,
+            sender_data: self.sender_data.clone(),
+            sender_handler: self.sender_handler,
+        }
+    }
+}