Add `async_task::spawn_local` variant that includes caller in panics (#21758)

Michael Sloan created

For debugging #21020. Copy-modified [from async_task
here](https://github.com/smol-rs/async-task/blob/ca9dbe1db9c422fd765847fa91306e30a6bb58a9/src/runnable.rs#L432)

Release Notes:

- N/A

Change summary

crates/gpui/src/executor.rs | 74 ++++++++++++++++++++++++++++++++++++++
1 file changed, 73 insertions(+), 1 deletion(-)

Detailed changes

crates/gpui/src/executor.rs 🔗

@@ -1,6 +1,10 @@
 use crate::{AppContext, PlatformDispatcher};
+use async_task::Runnable;
 use futures::channel::mpsc;
 use smol::prelude::*;
+use std::mem::ManuallyDrop;
+use std::panic::Location;
+use std::thread::{self, ThreadId};
 use std::{
     fmt::Debug,
     marker::PhantomData,
@@ -440,16 +444,19 @@ impl ForegroundExecutor {
     }
 
     /// Enqueues the given Task to run on the main thread at some point in the future.
+    #[track_caller]
     pub fn spawn<R>(&self, future: impl Future<Output = R> + 'static) -> Task<R>
     where
         R: 'static,
     {
         let dispatcher = self.dispatcher.clone();
+
+        #[track_caller]
         fn inner<R: 'static>(
             dispatcher: Arc<dyn PlatformDispatcher>,
             future: AnyLocalFuture<R>,
         ) -> Task<R> {
-            let (runnable, task) = async_task::spawn_local(future, move |runnable| {
+            let (runnable, task) = spawn_local_with_source_location(future, move |runnable| {
                 dispatcher.dispatch_on_main_thread(runnable)
             });
             runnable.schedule();
@@ -459,6 +466,71 @@ impl ForegroundExecutor {
     }
 }
 
+/// Variant of `async_task::spawn_local` that includes the source location of the spawn in panics.
+///
+/// Copy-modified from:
+/// https://github.com/smol-rs/async-task/blob/ca9dbe1db9c422fd765847fa91306e30a6bb58a9/src/runnable.rs#L405
+#[track_caller]
+fn spawn_local_with_source_location<Fut, S>(
+    future: Fut,
+    schedule: S,
+) -> (Runnable<()>, async_task::Task<Fut::Output, ()>)
+where
+    Fut: Future + 'static,
+    Fut::Output: 'static,
+    S: async_task::Schedule<()> + Send + Sync + 'static,
+{
+    #[inline]
+    fn thread_id() -> ThreadId {
+        std::thread_local! {
+            static ID: ThreadId = thread::current().id();
+        }
+        ID.try_with(|id| *id)
+            .unwrap_or_else(|_| thread::current().id())
+    }
+
+    struct Checked<F> {
+        id: ThreadId,
+        inner: ManuallyDrop<F>,
+        location: &'static Location<'static>,
+    }
+
+    impl<F> Drop for Checked<F> {
+        fn drop(&mut self) {
+            assert!(
+                self.id == thread_id(),
+                "local task dropped by a thread that didn't spawn it. Task spawned at {}",
+                self.location
+            );
+            unsafe {
+                ManuallyDrop::drop(&mut self.inner);
+            }
+        }
+    }
+
+    impl<F: Future> Future for Checked<F> {
+        type Output = F::Output;
+
+        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            assert!(
+                self.id == thread_id(),
+                "local task polled by a thread that didn't spawn it. Task spawned at {}",
+                self.location
+            );
+            unsafe { self.map_unchecked_mut(|c| &mut *c.inner).poll(cx) }
+        }
+    }
+
+    // Wrap the future into one that checks which thread it's on.
+    let future = Checked {
+        id: thread_id(),
+        inner: ManuallyDrop::new(future),
+        location: Location::caller(),
+    };
+
+    unsafe { async_task::spawn_unchecked(future, schedule) }
+}
+
 /// Scope manages a set of tasks that are enqueued and waited on together. See [`BackgroundExecutor::scoped`].
 pub struct Scope<'a> {
     executor: BackgroundExecutor,