Use futures_util mutex instead of flume for sending flush lock wakeups

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.