Web Async Rework (#3082)

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,
+}