chore: Refractor out common logic

Signed-off-by: John Nunley <[email protected]>

Author: notgull

src/lib.rs

@@ -150,24 +150,8 @@ impl<'a> Executor<'a> {
     pub fn spawn<T: Send + 'a>(&self, future: impl Future<Output = T> + Send + 'a) -> Task<T> {
         let mut active = self.state().active.lock().unwrap();
 
-        // Remove the task from the set of active tasks when the future finishes.
-        let index = active.vacant_entry().key();
-        let state = self.state().clone();
-        let future = async move {
-            let _guard = CallOnDrop(move || drop(state.active.lock().unwrap().try_remove(index)));
-            future.await
-        };
-
-        // Create the task and register it in the set of active tasks.
-        let (runnable, task) = unsafe {
-            Builder::new()
-                .propagate_panic(true)
-                .spawn_unchecked(|()| future, self.schedule())
-        };
-        active.insert(runnable.waker());
-
-        runnable.schedule();
-        task
+        // SAFETY: `T` and the future are `Send`.
+        unsafe { self.spawn_inner(future, &mut active) }
     }
 
     /// Spawns many tasks onto the executor.
@@ -218,25 +202,8 @@ impl<'a> Executor<'a> {
         let mut active = self.state().active.lock().unwrap();
 
         for (i, future) in futures.into_iter().enumerate() {
-            // Remove the task from the set of active tasks when the future finishes.
-            let index = active.vacant_entry().key();
-            let state = self.state().clone();
-            let future = async move {
-                let _guard =
-                    CallOnDrop(move || drop(state.active.lock().unwrap().try_remove(index)));
-                future.await
-            };
-
-            // Create the task and register it in the set of active tasks.
-            let (runnable, task) = unsafe {
-                Builder::new()
-                    .propagate_panic(true)
-                    .spawn_unchecked(|()| future, self.schedule())
-            };
-            active.insert(runnable.waker());
-
-            runnable.schedule();
-            handles.extend(Some(task));
+            // SAFETY: `T` and the future are `Send`.
+            handles.extend(Some(unsafe { self.spawn_inner(future, &mut active) }));
 
             // Yield the lock every once in a while to ease contention.
             if i.wrapping_sub(1) % 500 == 0 {
@@ -246,6 +213,36 @@ impl<'a> Executor<'a> {
         }
     }
 
+    /// Spawn a future using the inner lock.
+    ///
+    /// # Safety
+    ///
+    /// If this is an `Executor`, `F` and `T` must be `Send`.
+    unsafe fn spawn_inner<T: 'a>(
+        &self,
+        future: impl Future<Output = T> + 'a,
+        active: &mut Slab<Waker>,
+    ) -> Task<T> {
+        // Remove the task from the set of active tasks when the future finishes.
+        let index = active.vacant_entry().key();
+        let state = self.state().clone();
+        let future = async move {
+            let _guard = CallOnDrop(move || drop(state.active.lock().unwrap().try_remove(index)));
+            future.await
+        };
+
+        // Create the task and register it in the set of active tasks.
+        let (runnable, task) = unsafe {
+            Builder::new()
+                .propagate_panic(true)
+                .spawn_unchecked(|()| future, self.schedule())
+        };
+        active.insert(runnable.waker());
+
+        runnable.schedule();
+        task
+    }
+
     /// Attempts to run a task if at least one is scheduled.
     ///
     /// Running a scheduled task means simply polling its future once.
@@ -473,24 +470,8 @@ impl<'a> LocalExecutor<'a> {
     pub fn spawn<T: 'a>(&self, future: impl Future<Output = T> + 'a) -> Task<T> {
         let mut active = self.inner().state().active.lock().unwrap();
 
-        // Remove the task from the set of active tasks when the future finishes.
-        let index = active.vacant_entry().key();
-        let state = self.inner().state().clone();
-        let future = async move {
-            let _guard = CallOnDrop(move || drop(state.active.lock().unwrap().try_remove(index)));
-            future.await
-        };
-
-        // Create the task and register it in the set of active tasks.
-        let (runnable, task) = unsafe {
-            Builder::new()
-                .propagate_panic(true)
-                .spawn_unchecked(|()| future, self.schedule())
-        };
-        active.insert(runnable.waker());
-
-        runnable.schedule();
-        task
+        // SAFETY: This future is not thread safe.
+        unsafe { self.inner().spawn_inner(future, &mut active) }
     }
 
     /// Spawns many tasks onto the executor.
@@ -541,25 +522,10 @@ impl<'a> LocalExecutor<'a> {
         let mut active = self.inner().state().active.lock().unwrap();
 
         for future in futures {
-            // Remove the task from the set of active tasks when the future finishes.
-            let index = active.vacant_entry().key();
-            let state = self.inner().state().clone();
-            let future = async move {
-                let _guard =
-                    CallOnDrop(move || drop(state.active.lock().unwrap().try_remove(index)));
-                future.await
-            };
-
-            // Create the task and register it in the set of active tasks.
-            let (runnable, task) = unsafe {
-                Builder::new()
-                    .propagate_panic(true)
-                    .spawn_unchecked(|()| future, self.schedule())
-            };
-            active.insert(runnable.waker());
-
-            runnable.schedule();
-            handles.extend(Some(task));
+            // SAFETY: F and T are both `Send`.
+            handles.extend(Some(unsafe {
+                self.inner().spawn_inner(future, &mut active)
+            }));
         }
     }