fix desktop to poll without budget

packages/core/src/virtual_dom.rs

@@ -13,21 +13,10 @@ use crate::{
     scopes::{ScopeId, ScopeState},
     AttributeValue, Element, Event, Scope, SuspenseContext,
 };
-use futures_util::{
-    pin_mut,
-    stream::{futures_unordered, FuturesUnordered},
-    StreamExt,
-};
+use futures_util::{pin_mut, StreamExt};
 use rustc_hash::FxHashMap;
 use slab::Slab;
-use std::{
-    any::Any,
-    borrow::BorrowMut,
-    cell::Cell,
-    collections::{BTreeSet, HashSet},
-    future::Future,
-    rc::Rc,
-};
+use std::{any::Any, borrow::BorrowMut, cell::Cell, collections::BTreeSet, future::Future, rc::Rc};
 
 /// A virtual node system that progresses user events and diffs UI trees.
 ///
@@ -431,7 +420,7 @@ impl VirtualDom {
         loop {
             match some_msg.take() {
                 // If a bunch of messages are ready in a sequence, try to pop them off synchronously
-                Some(msg) => match dbg!(msg) {
+                Some(msg) => match msg {
                     SchedulerMsg::Immediate(id) => self.mark_dirty(id),
                     SchedulerMsg::TaskNotified(task) => self.handle_task_wakeup(task),
                     SchedulerMsg::SuspenseNotified(id) => self.handle_suspense_wakeup(id),
@@ -441,20 +430,23 @@ impl VirtualDom {
                 None => {
                     match self.rx.try_next() {
                         Ok(Some(val)) => some_msg = Some(val),
-                        Ok(None) => panic!("Scheduler channel closed"),
+                        Ok(None) => return,
                         Err(_) => {
                             // If we have any dirty scopes, or finished fiber trees then we should exit
                             if !self.dirty_scopes.is_empty() || !self.finished_fibers.is_empty() {
                                 return;
                             }
 
-                            println!("Waiting for next message...");
-                            let mut tasks = self.scheduler.tasks.borrow_mut();
+                            let res = (&mut *self.scheduler.tasks.borrow_mut()).next().await;
 
-                            let takss = &mut *tasks;
-
-                            use futures_util::StreamExt;
-                            takss.next().await;
+                            if res.is_none() {
+                                // If we have no tasks, then we should wait for a message
+                                if let Some(msg) = self.rx.next().await {
+                                    some_msg = Some(msg);
+                                } else {
+                                    return;
+                                }
+                            }
                         }
                     }
                 }