Fixes Waker::will_wake not working (#1199)

gui/src/rt/mod.rs

@@ -4,13 +4,11 @@ pub use self::window::*;
 use self::context::Context;
 use self::event::WindowEvent;
 use self::task::TaskList;
-use self::waker::Waker;
 use std::cell::Cell;
 use std::collections::HashMap;
 use std::error::Error;
 use std::future::Future;
 use std::rc::{Rc, Weak};
-use std::sync::Arc;
 use thiserror::Error;
 use winit::application::ApplicationHandler;
 use winit::error::{EventLoopError, OsError};
@@ -22,7 +20,6 @@ mod context;
 mod event;
 mod hook;
 mod task;
-mod waker;
 mod window;
 
 /// Run the specified future to completion then return.
@@ -46,11 +43,12 @@ pub fn run<T: 'static>(main: impl Future<Output = T> + 'static) -> Result<T, Run
     };
 
     // Run event loop.
-    let mut tasks = TaskList::default();
-    let main: Box<dyn Future<Output = ()>> = Box::new(main);
-    let main = tasks.insert(None, Box::into_pin(main));
+    let proxy = el.create_proxy();
+    let mut tasks = TaskList::new(proxy.clone());
+    let main = tasks.create(main);
+    let main = tasks.insert(main);
     let mut rt = Runtime {
-        el: el.create_proxy(),
+        el: proxy,
         tasks,
         main,
         hooks: Vec::new(),
@@ -67,10 +65,9 @@ pub fn run<T: 'static>(main: impl Future<Output = T> + 'static) -> Result<T, Run
 /// # Panics
 /// If called from the other thread than main thread.
 pub fn spawn(task: impl Future<Output = ()> + 'static) {
-    let task: Box<dyn Future<Output = ()>> = Box::new(task);
-
     Context::with(move |cx| {
-        let id = cx.tasks.insert(None, Box::into_pin(task));
+        let task = cx.tasks.create(task);
+        let id = cx.tasks.insert(task);
 
         // We have a context so there is an event loop for sure.
         assert!(cx.proxy.send_event(Event::TaskReady(id)).is_ok());
@@ -135,7 +132,6 @@ impl<T> Runtime<T> {
         };
 
         // Setup context.
-        let waker = Arc::new(Waker::new(self.el.clone(), id));
         let mut cx = Context {
             el,
             proxy: &self.el,
@@ -151,14 +147,14 @@ impl<T> Runtime<T> {
 
         // Poll the task.
         let r = cx.run(|| {
-            let waker = std::task::Waker::from(waker);
+            let waker = std::task::Waker::from(task.waker().clone());
             let mut cx = std::task::Context::from_waker(&waker);
 
-            task.as_mut().poll(&mut cx)
+            task.future_mut().poll(&mut cx)
         });
 
         if r.is_pending() {
-            self.tasks.insert(Some(id), task);
+            self.tasks.insert(task);
         }
 
         true

gui/src/rt/task.rs

@@ -1,32 +1,82 @@
+use super::Event;
 use std::collections::HashMap;
 use std::future::Future;
 use std::pin::Pin;
+use std::sync::Arc;
+use std::task::Wake;
+use winit::event_loop::EventLoopProxy;
 
 /// List of pending tasks.
-#[derive(Default)]
 pub struct TaskList {
-    list: HashMap<u64, Pin<Box<dyn Future<Output = ()>>>>,
+    el: EventLoopProxy<Event>,
+    list: HashMap<u64, Task>,
     next: u64,
 }
 
 impl TaskList {
-    pub fn insert(&mut self, id: Option<u64>, task: Pin<Box<dyn Future<Output = ()>>>) -> u64 {
-        // Get ID.
-        let id = match id {
-            Some(v) => v,
-            None => {
-                let v = self.next;
-                self.next = self.next.checked_add(1).unwrap();
-                v
-            }
-        };
+    pub fn new(el: EventLoopProxy<Event>) -> Self {
+        Self {
+            el,
+            list: HashMap::default(),
+            next: 0,
+        }
+    }
 
-        assert!(self.list.insert(id, task).is_none());
+    pub fn create(&mut self, task: impl Future<Output = ()> + 'static) -> Task {
+        let id = self.next;
+
+        self.next = self.next.checked_add(1).unwrap(); // It should be impossible but just in case.
+
+        Task {
+            future: Box::pin(task),
+            waker: Arc::new(Waker {
+                el: self.el.clone(),
+                task: id,
+            }),
+        }
+    }
 
+    pub fn insert(&mut self, task: Task) -> u64 {
+        let id = task.waker.task;
+        assert!(self.list.insert(id, task).is_none());
         id
     }
 
-    pub fn remove(&mut self, id: u64) -> Option<Pin<Box<dyn Future<Output = ()>>>> {
+    pub fn remove(&mut self, id: u64) -> Option<Task> {
         self.list.remove(&id)
     }
 }
+
+/// Contains a [`Future`] and its waker for a pending task.
+///
+/// We need to use a single waker per [`Future`] to make [`std::task::Waker::will_wake()`] works.
+pub struct Task {
+    future: Pin<Box<dyn Future<Output = ()>>>,
+    waker: Arc<Waker>,
+}
+
+impl Task {
+    pub fn future_mut(&mut self) -> Pin<&mut dyn Future<Output = ()>> {
+        self.future.as_mut()
+    }
+
+    pub fn waker(&self) -> &Arc<impl Wake + Send + Sync + 'static> {
+        &self.waker
+    }
+}
+
+/// Implementation of [`Wake`].
+struct Waker {
+    el: EventLoopProxy<Event>,
+    task: u64,
+}
+
+impl Wake for Waker {
+    fn wake(self: Arc<Self>) {
+        self.wake_by_ref();
+    }
+
+    fn wake_by_ref(self: &Arc<Self>) {
+        drop(self.el.send_event(Event::TaskReady(self.task)));
+    }
+}