diff --git a/crates/gpui/src/executor.rs b/crates/gpui/src/executor.rs index 34132d72f841c20e3f97bf5cdf49200c0a56b568..86d4d0058f54290fae2ea67790493310e17fd9db 100644 --- a/crates/gpui/src/executor.rs +++ b/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(&self, future: impl Future + 'static) -> Task where R: 'static, { let dispatcher = self.dispatcher.clone(); + + #[track_caller] fn inner( dispatcher: Arc, future: AnyLocalFuture, ) -> Task { - 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( + future: Fut, + schedule: S, +) -> (Runnable<()>, async_task::Task) +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 { + id: ThreadId, + inner: ManuallyDrop, + location: &'static Location<'static>, + } + + impl Drop for Checked { + 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 Future for Checked { + type Output = F::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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,