Properly implement wait_for_suspense

Wait_for_suspense now does not call flush_sync, making it
appropriate for deferring effects on the server. Futures will
still run on the server (as needed to progress suspense)
but use_effect / and memo watchers will not run on the server

examples/signals.rs

@@ -17,6 +17,17 @@ fn app() -> Element {
     // effects will always run after first mount and then whenever the signal values change
     use_effect(move || println!("Count changed to {}", count()));
 
+    // We can do early returns and conditional rendering which will pause all futures that haven't been polled
+    if count() > 30 {
+        return rsx! {
+            h1 { "Count is too high!" }
+            button {
+                onclick: move |_| count.set(0),
+                "Press to reset"
+            }
+        };
+    }
+
     // use_future will spawn an infinitely running future that can be started and stopped
     use_future(|| async move {
         loop {

packages/core/src/global_context.rs

@@ -216,17 +216,33 @@ pub fn use_drop<D: FnOnce() + 'static>(destroy: D) {
     });
 }
 
+pub fn use_before_render(f: impl FnMut() + 'static) {
+    use_hook(|| before_render(f));
+}
+
+pub fn use_after_render(f: impl FnMut() + 'static) {
+    use_hook(|| after_render(f));
+}
+
 /// Push a function to be run before the next render
 /// This is a hook and will always run, so you can't unschedule it
 /// Will run for every progression of suspense, though this might change in the future
-pub fn use_before_render(f: impl FnMut() + 'static) {
+pub fn before_render(f: impl FnMut() + 'static) {
     Runtime::with_current_scope(|cx| cx.push_before_render(f));
 }
 
+/// Push a function to be run after the render is complete, even if it didn't complete successfully
+pub fn after_render(f: impl FnMut() + 'static) {
+    Runtime::with_current_scope(|cx| cx.push_after_render(f));
+}
+
 /// Wait for the virtualdom to finish its sync work before proceeding
 ///
 /// This is useful if you've just triggered an update and want to wait for it to finish before proceeding with valid
 /// DOM nodes.
+///
+/// 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() {
     let mut polled = false;
 
@@ -246,7 +262,8 @@ pub async fn flush_sync() {
     .await;
 
     // If the the future got polled, then we don't need to prevent it from being dropped
-    // This would all be solved with generational indicies on tasks
+    // If we had generational indicies on tasks we could simply let the task remain in the queue and just be a no-op
+    // when it's run
     std::mem::forget(_task);
 
     struct FlushKey(Task);

packages/core/src/lib.rs

@@ -90,9 +90,9 @@ pub mod prelude {
         consume_context, consume_context_from_scope, current_scope_id, fc_to_builder, flush_sync,
         generation, has_context, needs_update, parent_scope, provide_context, provide_root_context,
         remove_future, schedule_update, schedule_update_any, spawn, spawn_forever, suspend,
-        try_consume_context, use_before_render, use_drop, use_error_boundary, use_hook,
-        use_hook_with_cleanup, AnyValue, Attribute, Component, ComponentFunction, Element,
-        ErrorBoundary, Event, EventHandler, Fragment, HasAttributes, IntoAttributeValue,
+        try_consume_context, use_after_render, use_before_render, use_drop, use_error_boundary,
+        use_hook, use_hook_with_cleanup, AnyValue, Attribute, Component, ComponentFunction,
+        Element, ErrorBoundary, Event, EventHandler, Fragment, HasAttributes, IntoAttributeValue,
         IntoDynNode, OptionStringFromMarker, Properties, Runtime, RuntimeGuard, ScopeId,
         ScopeState, SuperFrom, SuperInto, Task, Template, TemplateAttribute, TemplateNode, Throw,
         VNode, VNodeInner, VirtualDom,

packages/core/src/scope_arena.rs

@@ -57,6 +57,11 @@ impl VirtualDom {
 
         let context = scope.state();
 
+        // Run all post-render hooks
+        for post_run in context.after_render.borrow_mut().iter_mut() {
+            post_run();
+        }
+
         // And move the render generation forward by one
         context.render_count.set(context.render_count.get() + 1);
 

packages/core/src/scope_context.rs

@@ -24,6 +24,7 @@ pub(crate) struct Scope {
     pub(crate) shared_contexts: RefCell<Vec<Box<dyn Any>>>,
     pub(crate) spawned_tasks: RefCell<FxHashSet<Task>>,
     pub(crate) before_render: RefCell<Vec<Box<dyn FnMut()>>>,
+    pub(crate) after_render: RefCell<Vec<Box<dyn FnMut()>>>,
 }
 
 impl Scope {
@@ -45,6 +46,7 @@ impl Scope {
             hooks: RefCell::new(vec![]),
             hook_index: Cell::new(0),
             before_render: RefCell::new(vec![]),
+            after_render: RefCell::new(vec![]),
         }
     }
 
@@ -285,6 +287,10 @@ impl Scope {
         self.before_render.borrow_mut().push(Box::new(f));
     }
 
+    pub fn push_after_render(&self, f: impl FnMut() + 'static) {
+        self.after_render.borrow_mut().push(Box::new(f));
+    }
+
     /// Get the current render since the inception of this component
     ///
     /// This can be used as a helpful diagnostic when debugging hooks/renders, etc

packages/core/src/tasks.rs

@@ -58,6 +58,9 @@ impl Task {
         Runtime::with(|rt| {
             // set the active flag, and then ping the scheduler to ensure the task gets queued
             let was_active = rt.tasks.borrow()[self.0].active.replace(true);
+            if !was_active {
+                _ = rt.sender.unbounded_send(SchedulerMsg::TaskNotified(*self));
+            }
         });
     }
 }

packages/core/src/virtual_dom.rs

@@ -422,6 +422,12 @@ impl VirtualDom {
         // Ping tasks waiting on the flush table - they're waiting for sync stuff to be done before progressing
         self.clear_flush_table();
 
+        // And then poll the futures
+        self.poll_tasks().await;
+    }
+
+    /// Poll futures without progressing any futures from the flush table
+    async fn poll_tasks(&mut self) {
         loop {
             // Process all events - Scopes are marked dirty, etc
             self.process_events();
@@ -542,6 +548,7 @@ impl VirtualDom {
         self.flush_templates(to);
 
         // Process any events that might be pending in the queue
+        // Signals marked with .write() need a chance to be handled by the effect driver
         self.process_events();
 
         // Next, diff any dirty scopes
@@ -573,19 +580,18 @@ impl VirtualDom {
     ///
     /// The mutations will be thrown out, so it's best to use this method for things like SSR that have async content
     ///
-    /// Tasks waiting to be flushed are *cleared* here *without running them*
-    /// This behavior is subject to change, but configured this way so use_future/use_memo/use_future won't react on the server
+    /// We don't call "flush_sync" here since there's no sync work to be done. Futures will be progressed like usual,
+    /// however any futures wating on flush_sync will remain pending
     pub async fn wait_for_suspense(&mut self) {
         loop {
             if self.suspended_scopes.is_empty() {
-                return;
+                break;
             }
 
-            // not sure if we should be doing this?
-            self.runtime.flush_table.borrow_mut().clear();
-
-            self.wait_for_work().await;
+            // Wait for a work to be ready (IE new suspense leaves to pop up)
+            self.poll_tasks().await;
 
+            // Render whatever work needs to be rendered, unlocking new futures and suspense leaves
             self.render_immediate(&mut NoOpMutations);
         }
     }

packages/hooks/src/lib.rs

@@ -81,3 +81,6 @@ pub use use_resource::*;
 
 mod use_root_context;
 pub use use_root_context::*;
+
+mod use_hook_did_run;
+pub use use_hook_did_run::*;

packages/hooks/src/use_future.rs

@@ -1,14 +1,12 @@
 #![allow(missing_docs)]
+use crate::use_hook_did_run;
 use dioxus_core::{
-    prelude::{spawn, use_before_render, use_drop, use_hook},
-    ScopeState, Task,
+    prelude::{spawn, use_drop, use_hook},
+    Task,
 };
 use dioxus_signals::*;
 use dioxus_signals::{Readable, Writable};
-use futures_util::{future, pin_mut, FutureExt};
-use std::{any::Any, cell::Cell, future::Future, pin::Pin, rc::Rc, sync::Arc, task::Poll};
-
-use crate::use_callback;
+use std::future::Future;
 
 /// A hook that allows you to spawn a future
 ///
@@ -17,44 +15,28 @@ pub fn use_future<F>(mut future: impl FnMut() -> F) -> UseFuture
 where
     F: Future + 'static,
 {
-    let state = use_signal(|| UseFutureState::Pending);
+    let mut state = use_signal(|| UseFutureState::Pending);
 
     // Create the task inside a copyvalue so we can reset it in-place later
     let task = use_hook(|| {
         let fut = future();
         CopyValue::new(spawn(async move {
             fut.await;
+            state.set(UseFutureState::Complete);
         }))
     });
 
-    /*
-    Early returns in dioxus have consequences for use_memo, use_resource, and use_future, etc
-    We *don't* want futures to be running if the component early returns. It's a rather weird behavior to have
-    use_memo running in the background even if the component isn't hitting those hooks anymore.
-
-    React solves this by simply not having early returns interleave with hooks.
-    However, since dioxus allows early returns (since we use them for suspense), we need to solve this problem.
-
-
-     */
-    // Track if this *current* render is the same
-    let gen = use_hook(|| CopyValue::new((0, 0)));
-
-    // Early returns will pause this task, effectively
-    use_before_render(move || {
-        gen.write().0 += 1;
-        task.peek().set_active(false);
+    // Early returns in dioxus have consequences for use_memo, use_resource, and use_future, etc
+    // We *don't* want futures to be running if the component early returns. It's a rather weird behavior to have
+    // use_memo running in the background even if the component isn't hitting those hooks anymore.
+    //
+    // React solves this by simply not having early returns interleave with hooks.
+    // However, since dioxus allows early returns (since we use them for suspense), we need to solve this problem
+    use_hook_did_run(move |did_run| match did_run {
+        true => task.peek().resume(),
+        false => task.peek().pause(),
     });
 
-    // However when we actually run this component, we want to resume the task
-    task.peek().set_active(true);
-    gen.write().1 += 1;
-
-    // if the gens are different, we need to wake the task
-    if gen().0 != gen().1 {
-        task.peek().wake();
-    }
-
     use_drop(move || task.peek().stop());
 
     UseFuture { task, state }

packages/hooks/src/use_hook_did_run.rs

@@ -0,0 +1,16 @@
+use dioxus_core::prelude::*;
+use dioxus_signals::{CopyValue, Writable};
+
+/// A hook that uses before/after lifecycle hooks to determine if the hook was run
+pub fn use_hook_did_run(mut handler: impl FnMut(bool) + 'static) {
+    let mut did_run_ = use_hook(|| CopyValue::new(false));
+
+    // Before render always set the value to false
+    use_before_render(move || did_run_.set(false));
+
+    // Only when this hook is hit do we want to set the value to true
+    did_run_.set(true);
+
+    // After render, we can check if the hook was run
+    use_after_render(move || handler(did_run_()));
+}

packages/signals/src/effect.rs

@@ -47,53 +47,74 @@ impl EffectStackRef {
 }
 
 pub(crate) fn get_effect_ref() -> EffectStackRef {
-    match try_consume_context() {
-        Some(rt) => rt,
-        None => {
-            let (sender, mut receiver) = futures_channel::mpsc::unbounded();
-            spawn_forever(async move {
-                let mut queued_memos = Vec::new();
+    if let Some(rt) = try_consume_context() {
+        return rt;
+    }
 
-                loop {
-                    // Wait for a flush
-                    // This gives a chance for effects to be updated in place and memos to compute their values
-                    let flush_await = flush_sync();
-                    pin_mut!(flush_await);
+    let (sender, receiver) = futures_channel::mpsc::unbounded();
 
-                    loop {
-                        let res =
-                            futures_util::future::select(&mut flush_await, receiver.next()).await;
+    spawn_forever(async move { effect_driver(receiver).await });
 
-                        match res {
-                            Either::Right((_queued, _)) => {
-                                if let Some(task) = _queued {
-                                    queued_memos.push(task);
-                                }
-                                continue;
-                            }
-                            Either::Left(_flushed) => break,
-                        }
+    let stack_ref = EffectStackRef {
+        rerun_effect: sender,
+    };
+
+    provide_root_context(stack_ref.clone());
+
+    stack_ref
+}
+
+/// The primary top-level driver of all effects
+///
+/// In Dioxus, effects are neither react effects nor solidjs effects. They are a hybrid of the two, making our model
+/// more complex but also more powerful.
+///
+/// In react, when a component renders, it can queue up effects to be run after the component is done rendering.
+/// This is done *only during render* and determined by the dependency array attached to the effect. In Dioxus,
+/// we track effects using signals, so these effects can actually run multiple times after the component has rendered.
+///
+///
+async fn effect_driver(
+    mut receiver: futures_channel::mpsc::UnboundedReceiver<GenerationalBoxId>,
+) -> ! {
+    let mut queued_memos = Vec::new();
+
+    loop {
+        // Wait for a flush
+        // This gives a chance for effects to be updated in place and memos to compute their values
+        let flush_await = flush_sync();
+        pin_mut!(flush_await);
+
+        // Until the flush is ready, wait for a new effect to be queued
+        // We don't run the effects immediately because we want to batch them on the next call to flush
+        // todo: the queued memos should be unqueued when components are dropped
+        loop {
+            match futures_util::future::select(&mut flush_await, receiver.next()).await {
+                // VDOM is flushed and we can run the queued effects
+                Either::Left(_flushed) => break,
+
+                // A new effect was queued to be run after the next flush
+                // Marking components as dirty is handled syncrhonously on write, though we could try
+                // batching them here too
+                Either::Right((_queued, _)) => {
+                    if let Some(task) = _queued {
+                        queued_memos.push(task);
                     }
-
-                    EFFECT_STACK.with(|stack| {
-                        for id in queued_memos.drain(..) {
-                            let effect_mapping = stack.effect_mapping.read();
-                            if let Some(mut effect) = effect_mapping.get(&id).copied() {
-                                tracing::trace!("Rerunning effect: {:?}", id);
-                                effect.try_run();
-                            } else {
-                                tracing::trace!("Effect not found: {:?}", id);
-                            }
-                        }
-                    });
                 }
-            });
-            let stack_ref = EffectStackRef {
-                rerun_effect: sender,
-            };
-            provide_root_context(stack_ref.clone());
-            stack_ref
+            }
         }
+
+        EFFECT_STACK.with(|stack| {
+            for id in queued_memos.drain(..) {
+                let effect_mapping = stack.effect_mapping.read();
+                if let Some(mut effect) = effect_mapping.get(&id).copied() {
+                    tracing::trace!("Rerunning effect: {:?}", id);
+                    effect.try_run();
+                } else {
+                    tracing::trace!("Effect not found: {:?}", id);
+                }
+            }
+        });
     }
 }