From 3295935de7c501e8abf3198fb2c52f2de9405964 Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Fri, 2 Feb 2024 14:33:02 -0800
Subject: [PATCH] Fix panic on pausing dead tasks

---
 examples/control_focus.rs  | 23 ++++++++++++++---------
 packages/core/src/tasks.rs | 30 +++++++++++++++++++-----------
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/examples/control_focus.rs b/examples/control_focus.rs
index 1ca3f7aa8..1172327ea 100644
--- a/examples/control_focus.rs
+++ b/examples/control_focus.rs
@@ -12,16 +12,21 @@ fn app() -> Element {
 
     use_future(move || async move {
         let mut focused = 0;
-        if running() {
-            loop {
-                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
-                if let Some(element) = elements.with(|f| f.get(focused).cloned()) {
-                    _ = element.set_focus(true).await;
-                } else {
-                    focused = 0;
-                }
-                focused += 1;
+
+        loop {
+            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+
+            if !running() {
+                continue;
             }
+
+            if let Some(element) = elements.with(|f| f.get(focused).cloned()) {
+                _ = element.set_focus(true).await;
+            } else {
+                focused = 0;
+            }
+
+            focused += 1;
         }
     });
 
diff --git a/packages/core/src/tasks.rs b/packages/core/src/tasks.rs
index 1e0d60162..85ecb72c5 100644
--- a/packages/core/src/tasks.rs
+++ b/packages/core/src/tasks.rs
@@ -37,12 +37,24 @@ impl Task {
 
     /// Pause the task.
     pub fn pause(&self) {
-        Runtime::with(|rt| rt.tasks.borrow()[self.0].active.set(false));
+        self.set_active(false);
+    }
+
+    /// Resume the task.
+    pub fn resume(&self) {
+        self.set_active(true);
     }
 
     /// Check if the task is paused.
     pub fn paused(&self) -> bool {
-        Runtime::with(|rt| !rt.tasks.borrow()[self.0].active.get()).unwrap_or_default()
+        Runtime::with(|rt| {
+            if let Some(task) = rt.tasks.borrow().get(self.0) {
+                !task.active.get()
+            } else {
+                false
+            }
+        })
+        .unwrap_or_default()
     }
 
     /// Wake the task.
@@ -57,16 +69,12 @@ impl Task {
 
     /// Set the task as active or paused.
     pub fn set_active(&self, active: bool) {
-        Runtime::with(|rt| rt.tasks.borrow()[self.0].active.set(active));
-    }
-
-    /// Resume the task.
-    pub fn resume(&self) {
         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));
+            if let Some(task) = rt.tasks.borrow().get(self.0) {
+                let was_active = task.active.replace(active);
+                if !was_active && active {
+                    _ = rt.sender.unbounded_send(SchedulerMsg::TaskNotified(*self));
+                }
             }
         });
     }