Switch from using an Rc and forcing main thread access, to using an Arc

Mikayla Maki created

allowing this to be dropped from any thread.

Change summary

crates/gpui/src/app.rs                      |  24 +++--
crates/gpui/src/app/async_context.rs        |  15 ++-
crates/gpui/src/app/test_context.rs         |   1 
crates/gpui/src/executor.rs                 |  63 +++++--------
crates/gpui/src/platform.rs                 | 102 ++++++++++------------
crates/gpui/src/platform/mac/dispatcher.rs  |   5 
crates/gpui/src/platform/test/dispatcher.rs |   5 
7 files changed, 101 insertions(+), 114 deletions(-)

Detailed changes

crates/gpui/src/app.rs 🔗

@@ -34,15 +34,15 @@ use util::{ResultExt, debug_panic};
 #[cfg(any(feature = "inspector", debug_assertions))]
 use crate::InspectorElementRegistry;
 use crate::{
-    Action, ActionBuildError, ActionRegistry, Any, AnyView, AnyWindowHandle, AppContext, Asset,
-    AssetSource, BackgroundExecutor, Bounds, ClipboardItem, CursorStyle, DispatchPhase, DisplayId,
-    EventEmitter, FocusHandle, FocusMap, ForegroundExecutor, Global, KeyBinding, KeyContext,
-    Keymap, Keystroke, LayoutId, Menu, MenuItem, OwnedMenu, PathPromptOptions, Pixels, Platform,
-    PlatformDisplay, PlatformKeyboardLayout, PlatformKeyboardMapper, Point, Priority,
-    PromptBuilder, PromptButton, PromptHandle, PromptLevel, Render, RenderImage,
-    RenderablePromptHandle, Reservation, ScreenCaptureSource, SharedString, SubscriberSet,
-    Subscription, SvgRenderer, Task, TextSystem, Window, WindowAppearance, WindowHandle, WindowId,
-    WindowInvalidator,
+    Action, ActionBuildError, ActionRegistry, Any, AnyView, AnyWindowHandle, AppContext,
+    AppLiveness, Asset, AssetSource, BackgroundExecutor, Bounds, ClipboardItem, CursorStyle,
+    DispatchPhase, DisplayId, EventEmitter, FocusHandle, FocusMap, ForegroundExecutor, Global,
+    KeyBinding, KeyContext, Keymap, Keystroke, LayoutId, Menu, MenuItem, OwnedMenu,
+    PathPromptOptions, Pixels, Platform, PlatformDisplay, PlatformKeyboardLayout,
+    PlatformKeyboardMapper, Point, Priority, PromptBuilder, PromptButton, PromptHandle,
+    PromptLevel, Render, RenderImage, RenderablePromptHandle, Reservation, ScreenCaptureSource,
+    SharedString, SubscriberSet, Subscription, SvgRenderer, Task, TextSystem, Window,
+    WindowAppearance, WindowHandle, WindowId, WindowInvalidator,
     colors::{Colors, GlobalColors},
     current_platform, hash, init_app_menus,
 };
@@ -580,6 +580,9 @@ impl GpuiMode {
 /// You need a reference to an `App` to access the state of a [Entity].
 pub struct App {
     pub(crate) this: Weak<AppCell>,
+    /// Tracks whether this app is still alive. Used to cancel foreground tasks
+    /// when the app is dropped.
+    pub(crate) liveness: AppLiveness,
     pub(crate) platform: Rc<dyn Platform>,
     pub(crate) mode: GpuiMode,
     text_system: Arc<TextSystem>,
@@ -658,6 +661,7 @@ impl App {
         let app = Rc::new_cyclic(|this| AppCell {
             app: RefCell::new(App {
                 this: this.clone(),
+                liveness: AppLiveness::new(),
                 platform: platform.clone(),
                 text_system,
                 mode: GpuiMode::Production,
@@ -1474,8 +1478,10 @@ impl App {
     /// Creates an `AsyncApp`, which can be cloned and has a static lifetime
     /// so it can be held across `await` points.
     pub fn to_async(&self) -> AsyncApp {
+        let liveness_token = self.liveness.token();
         AsyncApp {
             app: self.this.clone(),
+            liveness_token,
             background_executor: self.background_executor.clone(),
             foreground_executor: self.foreground_executor.clone(),
         }

crates/gpui/src/app/async_context.rs 🔗

@@ -1,7 +1,8 @@
 use crate::{
-    AnyView, AnyWindowHandle, App, AppCell, AppContext, BackgroundExecutor, BorrowAppContext,
-    Entity, EventEmitter, Focusable, ForegroundExecutor, Global, PromptButton, PromptLevel, Render,
-    Reservation, Result, Subscription, Task, VisualContext, Window, WindowHandle,
+    AnyView, AnyWindowHandle, App, AppCell, AppContext, AppLivenessToken, BackgroundExecutor,
+    BorrowAppContext, Entity, EventEmitter, Focusable, ForegroundExecutor, Global, PromptButton,
+    PromptLevel, Render, Reservation, Result, Subscription, Task, VisualContext, Window,
+    WindowHandle,
 };
 use anyhow::{Context as _, anyhow};
 use derive_more::{Deref, DerefMut};
@@ -16,6 +17,7 @@ use super::{Context, WeakEntity};
 #[derive(Clone)]
 pub struct AsyncApp {
     pub(crate) app: Weak<AppCell>,
+    pub(crate) liveness_token: AppLivenessToken,
     pub(crate) background_executor: BackgroundExecutor,
     pub(crate) foreground_executor: ForegroundExecutor,
 }
@@ -185,7 +187,7 @@ impl AsyncApp {
     {
         let mut cx = self.clone();
         self.foreground_executor
-            .spawn_with_app(self.app.clone(), async move { f(&mut cx).await })
+            .spawn_context(self.liveness_token.clone(), async move { f(&mut cx).await })
     }
 
     /// Determine whether global state of the specified type has been assigned.
@@ -334,7 +336,10 @@ impl AsyncWindowContext {
     {
         let mut cx = self.clone();
         self.foreground_executor
-            .spawn_with_app(self.app.app.clone(), async move { f(&mut cx).await })
+            .spawn_context(
+                self.app.liveness_token.clone(),
+                async move { f(&mut cx).await },
+            )
     }
 
     /// Present a platform dialog.

crates/gpui/src/app/test_context.rs 🔗

@@ -405,6 +405,7 @@ impl TestAppContext {
     pub fn to_async(&self) -> AsyncApp {
         AsyncApp {
             app: Rc::downgrade(&self.app),
+            liveness_token: self.app.borrow().liveness.token(),
             background_executor: self.background_executor.clone(),
             foreground_executor: self.foreground_executor.clone(),
         }

crates/gpui/src/executor.rs 🔗

@@ -1,4 +1,6 @@
-use crate::{App, PlatformDispatcher, RunnableMeta, RunnableVariant, TaskTiming, profiler};
+use crate::{
+    App, AppLivenessToken, PlatformDispatcher, RunnableMeta, RunnableVariant, TaskTiming, profiler,
+};
 use async_task::Runnable;
 use futures::channel::mpsc;
 use parking_lot::{Condvar, Mutex};
@@ -766,7 +768,7 @@ impl ForegroundExecutor {
     where
         R: 'static,
     {
-        self.spawn_with_app_and_priority(None, Priority::default(), future)
+        self.inner_spawn(None, Priority::default(), future)
     }
 
     /// Enqueues the given Task to run on the main thread at some point in the future.
@@ -779,32 +781,25 @@ impl ForegroundExecutor {
     where
         R: 'static,
     {
-        self.spawn_with_app_and_priority(None, priority, future)
+        self.inner_spawn(None, priority, future)
     }
 
-    /// Enqueues the given Task to run on the main thread at some point in the future,
-    /// with a weak reference to the app for cancellation checking.
-    ///
-    /// When the app is dropped, pending tasks spawned with this method will be cancelled
-    /// before they run, rather than panicking when they try to access the dropped app.
     #[track_caller]
-    pub fn spawn_with_app<R>(
+    pub(crate) fn spawn_context<R>(
         &self,
-        app: std::rc::Weak<crate::AppCell>,
+        app: AppLivenessToken,
         future: impl Future<Output = R> + 'static,
     ) -> Task<R>
     where
         R: 'static,
     {
-        self.spawn_with_app_and_priority(Some(app), Priority::default(), future)
+        self.inner_spawn(Some(app), Priority::default(), future)
     }
 
-    /// Enqueues the given Task to run on the main thread at some point in the future,
-    /// with an optional weak reference to the app for cancellation checking and a specific priority.
     #[track_caller]
-    pub fn spawn_with_app_and_priority<R>(
+    pub(crate) fn inner_spawn<R>(
         &self,
-        app: Option<std::rc::Weak<crate::AppCell>>,
+        app: Option<AppLivenessToken>,
         priority: Priority,
         future: impl Future<Output = R> + 'static,
     ) -> Task<R>
@@ -819,21 +814,15 @@ impl ForegroundExecutor {
             dispatcher: Arc<dyn PlatformDispatcher>,
             future: AnyLocalFuture<R>,
             location: &'static core::panic::Location<'static>,
-            app: Option<std::rc::Weak<crate::AppCell>>,
+            app: Option<AppLivenessToken>,
             priority: Priority,
         ) -> Task<R> {
-            // SAFETY: We are on the main thread (ForegroundExecutor is !Send), and the
-            // MainThreadWeak will only be accessed on the main thread in the trampoline.
-            let app_weak = app.map(|weak| unsafe { crate::MainThreadWeak::new(weak) });
             let (runnable, task) = spawn_local_with_source_location(
                 future,
                 move |runnable| {
                     dispatcher.dispatch_on_main_thread(RunnableVariant::Meta(runnable), priority)
                 },
-                RunnableMeta {
-                    location,
-                    app: app_weak,
-                },
+                RunnableMeta { location, app },
             );
             runnable.schedule();
             Task(TaskState::Spawned(task))
@@ -977,7 +966,7 @@ mod test {
     use super::*;
     use crate::{App, TestDispatcher, TestPlatform};
     use rand::SeedableRng;
-    use std::{cell::RefCell, rc::Weak};
+    use std::cell::RefCell;
 
     #[test]
     fn sanity_test_tasks_run() {
@@ -991,11 +980,12 @@ mod test {
         let http_client = http_client::FakeHttpClient::with_404_response();
 
         let app = App::new_app(platform, asset_source, http_client);
+        let liveness_token = app.borrow().liveness.token();
 
         let task_ran = Rc::new(RefCell::new(false));
 
         foreground_executor
-            .spawn_with_app(Rc::downgrade(&app), {
+            .spawn_context(liveness_token, {
                 let task_ran = Rc::clone(&task_ran);
                 async move {
                     *task_ran.borrow_mut() = true;
@@ -1025,22 +1015,18 @@ mod test {
         let http_client = http_client::FakeHttpClient::with_404_response();
 
         let app = App::new_app(platform, asset_source, http_client);
+        let liveness_token = app.borrow().liveness.token();
         let app_weak = Rc::downgrade(&app);
 
         let task_ran = Rc::new(RefCell::new(false));
         let task_ran_clone = Rc::clone(&task_ran);
 
         foreground_executor
-            .spawn_with_app(Weak::clone(&app_weak), async move {
+            .spawn_context(liveness_token, async move {
                 *task_ran_clone.borrow_mut() = true;
             })
             .detach();
 
-        assert!(
-            Rc::weak_count(&app) > 0,
-            "Task should hold a weak reference"
-        );
-
         drop(app);
 
         assert!(app_weak.upgrade().is_none(), "App should have been dropped");
@@ -1066,6 +1052,7 @@ mod test {
         let http_client = http_client::FakeHttpClient::with_404_response();
 
         let app = App::new_app(platform, asset_source, http_client);
+        let liveness_token = app.borrow().liveness.token();
         let app_weak = Rc::downgrade(&app);
 
         let outer_completed = Rc::new(RefCell::new(false));
@@ -1079,17 +1066,17 @@ mod test {
         // Channel to block the inner task until we're ready
         let (tx, rx) = futures::channel::oneshot::channel::<()>();
 
-        // We need clones of executor and app_weak for the inner spawn
+        // We need clones of executor and liveness_token for the inner spawn
         let inner_executor = foreground_executor.clone();
-        let inner_app_weak = app_weak.clone();
+        let inner_liveness_token = liveness_token.clone();
 
         // Spawn outer task that will spawn and await an inner task
         foreground_executor
-            .spawn_with_app(Weak::clone(&app_weak), async move {
+            .spawn_context(liveness_token, async move {
                 let inner_flag_clone = Rc::clone(&inner_flag);
 
                 // Spawn inner task that blocks on a channel
-                let inner_task = inner_executor.spawn_with_app(inner_app_weak, async move {
+                let inner_task = inner_executor.spawn_context(inner_liveness_token, async move {
                     // Wait for signal (which will never come - we'll drop the app instead)
                     rx.await.ok();
                     *inner_flag_clone.borrow_mut() = true;
@@ -1191,9 +1178,10 @@ mod test {
         let http_client = http_client::FakeHttpClient::with_404_response();
 
         let app = App::new_app(platform, asset_source, http_client);
+        let liveness_token = app.borrow().liveness.token();
         let app_weak = Rc::downgrade(&app);
 
-        let task = foreground_executor.spawn_with_app(Weak::clone(&app_weak), async move { 42 });
+        let task = foreground_executor.spawn_context(liveness_token, async move { 42 });
 
         drop(app);
 
@@ -1216,10 +1204,11 @@ mod test {
         let http_client = http_client::FakeHttpClient::with_404_response();
 
         let app = App::new_app(platform, asset_source, http_client);
+        let liveness_token = app.borrow().liveness.token();
         let app_weak = Rc::downgrade(&app);
 
         let task = foreground_executor
-            .spawn_with_app(Weak::clone(&app_weak), async move { 42 })
+            .spawn_context(liveness_token, async move { 42 })
             .fallible();
 
         drop(app);

crates/gpui/src/platform.rs 🔗

@@ -572,75 +572,63 @@ pub(crate) trait PlatformWindow: HasWindowHandle + HasDisplayHandle {
     }
 }
 
-/// An rc::Weak<AppCell> that can cross thread boundaries but must only be accessed on the main thread.
+/// Tracks whether an App is still alive. This is used to cancel foreground tasks
+/// when the app is dropped.
 ///
-/// # Safety
-/// This type wraps a `Weak<AppCell>` (which is `!Send` and `!Sync`) and unsafely implements
-/// `Send` and `Sync`. The safety contract is:
-/// - Only create instances of this type on the main thread
-/// - Only access (upgrade) instances on the main thread
-/// - Only drop instances on the main thread
-///
-/// This is used to pass a weak reference to the app through the task scheduler, which
-/// requires `Send + Sync`, but the actual access only ever happens on the main thread
-/// in the trampoline function.
+#[derive(Clone)]
 #[doc(hidden)]
-pub struct MainThreadWeak {
-    weak: std::rc::Weak<crate::AppCell>,
-    #[cfg(debug_assertions)]
-    thread_id: std::thread::ThreadId,
+pub struct AppLiveness {
+    sentinel: std::sync::Arc<()>,
 }
 
-impl std::fmt::Debug for MainThreadWeak {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("MainThreadWeak").finish_non_exhaustive()
+impl AppLiveness {
+    /// Creates a new AppLiveness, initially alive.
+    pub fn new() -> Self {
+        Self {
+            sentinel: std::sync::Arc::new(()),
+        }
     }
-}
 
-impl MainThreadWeak {
-    /// Creates a new `MainThreadWeak` from a `Weak<AppCell>`.
-    ///
-    /// # Safety
-    /// Must only be called on the main thread.
-    pub unsafe fn new(weak: std::rc::Weak<crate::AppCell>) -> Self {
-        Self {
-            weak,
-            #[cfg(debug_assertions)]
-            thread_id: std::thread::current().id(),
+    /// Returns a token that can be stored in task metadata to check liveness.
+    pub fn token(&self) -> AppLivenessToken {
+        AppLivenessToken {
+            sentinel: std::sync::Arc::downgrade(&self.sentinel),
         }
     }
+}
 
-    /// Attempts to upgrade the weak reference to a strong reference.
-    ///
-    /// # Safety
-    /// Must only be called on the main thread.
-    pub unsafe fn upgrade(&self) -> Option<std::rc::Rc<crate::AppCell>> {
-        #[cfg(debug_assertions)]
-        debug_assert_eq!(
-            std::thread::current().id(),
-            self.thread_id,
-            "MainThreadWeak::upgrade called from a different thread than it was created on"
-        );
-        self.weak.upgrade()
+impl Default for AppLiveness {
+    fn default() -> Self {
+        Self::new()
     }
 }
 
-impl Drop for MainThreadWeak {
-    fn drop(&mut self) {
-        #[cfg(debug_assertions)]
-        debug_assert_eq!(
-            std::thread::current().id(),
-            self.thread_id,
-            "MainThreadWeak dropped on a different thread than it was created on"
-        );
+impl std::fmt::Debug for AppLiveness {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("AppLiveness").field("alive", &true).finish()
     }
 }
 
-// SAFETY: MainThreadWeak is only created, accessed, and dropped on the main thread.
-// The Send + Sync impls are needed because RunnableMeta must be Send + Sync for the
-// async-task crate, but we guarantee main-thread-only access.
-unsafe impl Send for MainThreadWeak {}
-unsafe impl Sync for MainThreadWeak {}
+/// A token that can be stored in task metadata to check if the app is still alive.
+#[derive(Clone)]
+pub struct AppLivenessToken {
+    sentinel: std::sync::Weak<()>,
+}
+
+impl AppLivenessToken {
+    /// Returns true if the app is still alive.
+    pub fn is_alive(&self) -> bool {
+        self.sentinel.strong_count() > 0
+    }
+}
+
+impl std::fmt::Debug for AppLivenessToken {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("AppLivenessToken")
+            .field("alive", &self.is_alive())
+            .finish()
+    }
+}
 
 /// This type is public so that our test macro can generate and use it, but it should not
 /// be considered part of our public API.
@@ -649,9 +637,9 @@ unsafe impl Sync for MainThreadWeak {}
 pub struct RunnableMeta {
     /// Location of the runnable
     pub location: &'static core::panic::Location<'static>,
-    /// Weak reference to the app, used to check if the app is still alive before running.
-    /// This must only be `Some()` on foreground tasks.
-    pub app: Option<MainThreadWeak>,
+    /// Token to check if the app is still alive before running.
+    /// This is `Some` for foreground tasks spawned with app tracking.
+    pub app: Option<AppLivenessToken>,
 }
 
 #[doc(hidden)]

crates/gpui/src/platform/mac/dispatcher.rs 🔗

@@ -254,9 +254,8 @@ extern "C" fn trampoline(runnable: *mut c_void) {
     let metadata = task.metadata();
     let location = metadata.location;
 
-    if let Some(ref app_weak) = metadata.app {
-        // SAFETY: App is only `Some()` when this trampoline is on the main thread.
-        if unsafe { app_weak.upgrade() }.is_none() {
+    if let Some(ref app_token) = metadata.app {
+        if !app_token.is_alive() {
             drop(task);
             return;
         }

crates/gpui/src/platform/test/dispatcher.rs 🔗

@@ -178,9 +178,8 @@ impl TestDispatcher {
         // todo(localcc): add timings to tests
         match runnable {
             RunnableVariant::Meta(runnable) => {
-                if let Some(ref app_weak) = runnable.metadata().app {
-                    // SAFETY: Test dispatcher should always run on the same thead as it's App
-                    if unsafe { app_weak.upgrade() }.is_none() {
+                if let Some(ref app_token) = runnable.metadata().app {
+                    if !app_token.is_alive() {
                         drop(runnable);
                         self.state.lock().is_main_thread = was_main_thread;
                         return true;