From e4764f2bba4f6a13c17016428ecb003bd3d5f76e Mon Sep 17 00:00:00 2001
From: Evan Almloff <evanalmloff@gmail.com>
Date: Fri, 7 Jun 2024 23:36:08 +0200
Subject: [PATCH] switch from slab to slotmap for tasks to fix that ABA problem
 (#2488)

---
 Cargo.lock                   | 11 ++++++++++
 Cargo.toml                   |  1 +
 packages/core/Cargo.toml     |  1 +
 packages/core/src/runtime.rs |  4 +++-
 packages/core/src/tasks.rs   | 40 +++++++++++++++++++++---------------
 5 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 18f647dd1..e9721cf49 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2253,6 +2253,7 @@ dependencies = [
  "rustversion",
  "serde",
  "slab",
+ "slotmap",
  "tokio",
  "tracing",
  "tracing-fluent-assertions",
@@ -8343,6 +8344,16 @@ dependencies = [
  "rustc-hash",
 ]
 
+[[package]]
+name = "slotmap"
+version = "1.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dbff4acf519f630b3a3ddcfaea6c06b42174d9a44bc70c620e9ed1649d58b82a"
+dependencies = [
+ "serde",
+ "version_check",
+]
+
 [[package]]
 name = "smallvec"
 version = "1.13.2"
diff --git a/Cargo.toml b/Cargo.toml
index 74f39c4d8..ab586d3f6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -86,6 +86,7 @@ tracing-futures = "0.2.5"
 toml = "0.8"
 tokio = "1.28"
 slab = "0.4.2"
+slotmap = { version = "1.0.7", features = ["serde"] }
 futures-channel = "0.3.21"
 futures-util = { version = "0.3", default-features = false }
 rustc-hash = "1.1.0"
diff --git a/packages/core/Cargo.toml b/packages/core/Cargo.toml
index 47f835b34..f41c4aff3 100644
--- a/packages/core/Cargo.toml
+++ b/packages/core/Cargo.toml
@@ -17,6 +17,7 @@ futures-util = { workspace = true, default-features = false, features = [
     "std",
 ] }
 slab = { workspace = true }
+slotmap = { workspace = true }
 futures-channel = { workspace = true }
 tracing = { workspace = true }
 serde = { version = "1", features = ["derive"], optional = true }
diff --git a/packages/core/src/runtime.rs b/packages/core/src/runtime.rs
index ee727f348..95fe44855 100644
--- a/packages/core/src/runtime.rs
+++ b/packages/core/src/runtime.rs
@@ -1,3 +1,5 @@
+use slotmap::DefaultKey;
+
 use crate::innerlude::Effect;
 use crate::{
     innerlude::{LocalTask, SchedulerMsg},
@@ -27,7 +29,7 @@ pub struct Runtime {
     pub(crate) current_task: Cell<Option<Task>>,
 
     /// Tasks created with cx.spawn
-    pub(crate) tasks: RefCell<slab::Slab<Rc<LocalTask>>>,
+    pub(crate) tasks: RefCell<slotmap::SlotMap<DefaultKey, Rc<LocalTask>>>,
 
     // Currently suspended tasks
     pub(crate) suspended_tasks: Cell<usize>,
diff --git a/packages/core/src/tasks.rs b/packages/core/src/tasks.rs
index cd57de861..d81c309b5 100644
--- a/packages/core/src/tasks.rs
+++ b/packages/core/src/tasks.rs
@@ -3,6 +3,7 @@ use crate::innerlude::ScopeOrder;
 use crate::innerlude::{remove_future, spawn, Runtime};
 use crate::ScopeId;
 use futures_util::task::ArcWake;
+use slotmap::DefaultKey;
 use std::sync::Arc;
 use std::task::Waker;
 use std::{cell::Cell, future::Future};
@@ -14,7 +15,7 @@ use std::{pin::Pin, task::Poll};
 /// `Task` is a unique identifier for a task that has been spawned onto the runtime. It can be used to cancel the task
 #[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
-pub struct Task(pub(crate) usize);
+pub struct Task(pub(crate) slotmap::DefaultKey);
 
 impl Task {
     /// Start a new future on the same thread as the rest of the VirtualDom.
@@ -139,24 +140,29 @@ impl Runtime {
         let (task, task_id) = {
             let mut tasks = self.tasks.borrow_mut();
 
-            let entry = tasks.vacant_entry();
-            let task_id = Task(entry.key());
+            let mut task_id = Task(DefaultKey::default());
+            let mut local_task = None;
+            tasks.insert_with_key(|key| {
+                task_id = Task(key);
 
-            let task = Rc::new(LocalTask {
-                scope,
-                active: Cell::new(true),
-                parent: self.current_task(),
-                task: RefCell::new(Box::pin(task)),
-                waker: futures_util::task::waker(Arc::new(LocalTaskHandle {
-                    id: task_id,
-                    tx: self.sender.clone(),
-                })),
-                ty: Cell::new(ty),
+                let new_task = Rc::new(LocalTask {
+                    scope,
+                    active: Cell::new(true),
+                    parent: self.current_task(),
+                    task: RefCell::new(Box::pin(task)),
+                    waker: futures_util::task::waker(Arc::new(LocalTaskHandle {
+                        id: task_id,
+                        tx: self.sender.clone(),
+                    })),
+                    ty: Cell::new(ty),
+                });
+
+                local_task = Some(new_task.clone());
+
+                new_task
             });
 
-            entry.insert(task.clone());
-
-            (task, task_id)
+            (local_task.unwrap(), task_id)
         };
 
         // Get a borrow on the task, holding no borrows on the tasks map
@@ -244,7 +250,7 @@ impl Runtime {
     ///
     /// This does not abort the task, so you'll want to wrap it in an abort handle if that's important to you
     pub(crate) fn remove_task(&self, id: Task) -> Option<Rc<LocalTask>> {
-        let task = self.tasks.borrow_mut().try_remove(id.0);
+        let task = self.tasks.borrow_mut().remove(id.0);
         if let Some(task) = &task {
             if task.suspended() {
                 self.suspended_tasks.set(self.suspended_tasks.get() - 1);