Use futures_util mutex instead of flume for sending flush lock wakeups

packages/core/src/global_context.rs

@@ -1,5 +1,5 @@
 use crate::{runtime::Runtime, Element, ScopeId, Task};
-use futures_util::{future::poll_fn, Future};
+use futures_util::Future;
 use std::sync::Arc;
 
 /// Get the current scope id
@@ -257,13 +257,11 @@ pub fn after_render(f: impl FnMut() + 'static) {
 /// Effects rely on this to ensure that they only run effects after the DOM has been updated. Without flush_sync effects
 /// are run immediately before diffing the DOM, which causes all sorts of out-of-sync weirdness.
 pub async fn flush_sync() {
-    // if flushing() {
-    //      return;
-    // }
-
-    _ = Runtime::with(|rt| rt.flush.clone())
+    // Wait for the flush lock to be available
+    // We release it immediately, so it's impossible for the lock to be held longer than this function
+    Runtime::with(|rt| rt.flush_mutex.clone())
         .unwrap()
-        .recv_async()
+        .lock()
         .await;
 }
 

packages/core/src/runtime.rs

@@ -1,5 +1,3 @@
-use rustc_hash::FxHashSet;
-
 use crate::{
     innerlude::{LocalTask, SchedulerMsg},
     scope_context::Scope,
@@ -9,6 +7,7 @@ use crate::{
 use std::{
     cell::{Cell, Ref, RefCell},
     rc::Rc,
+    sync::Arc,
 };
 
 thread_local! {
@@ -32,18 +31,18 @@ pub struct Runtime {
 
     pub(crate) sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
 
-    // Tasks waiting to be manually resumed when we call wait_for_work
-    pub(crate) flush: flume::Receiver<()>,
+    // the virtualdom will hold this lock while it's doing syncronous work
+    // when the lock is lifted, tasks waiting for the lock will be able to run
+    pub(crate) flush_mutex: Arc<futures_util::lock::Mutex<()>>,
+    pub(crate) flush_lock: Cell<Option<futures_util::lock::OwnedMutexGuard<()>>>,
 }
 
 impl Runtime {
-    pub(crate) fn new(
-        sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
-        flush: flume::Receiver<()>,
-    ) -> Rc<Self> {
+    pub(crate) fn new(sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>) -> Rc<Self> {
         Rc::new(Self {
             sender,
-            flush,
+            flush_mutex: Default::default(),
+            flush_lock: Default::default(),
             rendering: Cell::new(true),
             scope_states: Default::default(),
             scope_stack: Default::default(),
@@ -150,6 +149,23 @@ impl Runtime {
     pub(crate) fn with_scope<R>(scope: ScopeId, f: impl FnOnce(&Scope) -> R) -> Option<R> {
         Self::with(|rt| rt.get_state(scope).map(|sc| f(&sc))).flatten()
     }
+
+    /// Acquire the flush lock and store it interally
+    ///
+    /// This means the virtual dom is currently doing syncronous work
+    /// The lock will be held until `release_flush_lock` is called - and then the OwnedLock will be dropped
+    pub(crate) fn acquire_flush_lock(&self) {
+        self.flush_lock
+            .set(Some(self.flush_mutex.try_lock_owned().unwrap()));
+    }
+
+    /// Release the flush lock
+    ///
+    /// On the drop of the flush lock, all tasks waiting on `flush_sync` will spring to life via their wakers.
+    /// You can now freely poll those tasks and they can progress
+    pub(crate) fn release_flush_lock(&self) {
+        self.flush_lock.take();
+    }
 }
 
 /// A guard for a new runtime. This must be used to override the current runtime when importing components from a dynamic library that has it's own runtime.

packages/core/src/virtual_dom.rs

@@ -202,7 +202,6 @@ pub struct VirtualDom {
     pub(crate) suspended_scopes: FxHashSet<ScopeId>,
 
     rx: futures_channel::mpsc::UnboundedReceiver<SchedulerMsg>,
-    flush_tx: flume::Sender<()>,
 }
 
 impl VirtualDom {
@@ -306,12 +305,10 @@ impl VirtualDom {
     /// ```
     pub(crate) fn new_with_component(root: impl AnyProps + 'static) -> Self {
         let (tx, rx) = futures_channel::mpsc::unbounded();
-        let (flush_tx, flush_rx) = flume::unbounded(); // I don't think this needs to be unbounded
 
         let mut dom = Self {
             rx,
-            flush_tx,
-            runtime: Runtime::new(tx, flush_rx),
+            runtime: Runtime::new(tx),
             scopes: Default::default(),
             dirty_scopes: Default::default(),
             templates: Default::default(),
@@ -430,12 +427,13 @@ impl VirtualDom {
         self.poll_tasks().await;
     }
 
-    /// Poll futures without progressing any futures from the flush table
+    ///
     async fn poll_tasks(&mut self) {
-        loop {
-            // Send the flush signal to the runtime, waking up tasks
-            _ = self.flush_tx.try_send(());
+        // Release the flush lock
+        // This will cause all the flush wakers to immediately spring to life, which we will off with process_events
+        self.runtime.release_flush_lock();
 
+        loop {
             // Process all events - Scopes are marked dirty, etc
             // Sometimes when wakers fire we get a slew of updates at once, so its important that we drain this completely
             self.process_events();
@@ -448,6 +446,10 @@ impl VirtualDom {
             // Make sure we set the runtime since we're running user code
             let _runtime = RuntimeGuard::new(self.runtime.clone());
 
+            // Hold a lock to the flush sync to prevent tasks from running in the event we get an immediate
+            // When we're doing awaiting the rx, the lock will be dropped and tasks waiting on the lock will get waked
+            self.runtime.acquire_flush_lock();
+
             match self.rx.next().await.expect("channel should never close") {
                 SchedulerMsg::Immediate(id) => self.mark_dirty(id),
                 SchedulerMsg::TaskNotified(id) => self.runtime.handle_task_wakeup(id),
@@ -546,6 +548,7 @@ impl VirtualDom {
 
         // Process any events that might be pending in the queue
         // Signals marked with .write() need a chance to be handled by the effect driver
+        // This also processes futures which might progress into immediates
         self.process_events();
 
         // Next, diff any dirty scopes