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

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);
         }
     }