gpui_tokio: Fix `child task panicked: JoinError::Cancelled` panics (#51995) (cherry-pick to preview) (#51997)

zed-zippy[bot] and Lukas Wirth created

Cherry-pick of #51995 to preview

----
This panic stemmed from us dropping the tokio runtime before dropping
wasm tasks which then could attempt to spawn_blocking, immediately
failing and panicking on unwrap inside wasmtime.

Closes ZED-5DE

Release Notes:

- N/A or Added/Fixed/Improved ...

Co-authored-by: Lukas Wirth <lukas@zed.dev>

Change summary

crates/gpui/src/app.rs            | 59 ++++++++++++++++++++++----------
crates/gpui/src/app/entity_map.rs |  5 +-
2 files changed, 44 insertions(+), 20 deletions(-)

Detailed changes

crates/gpui/src/app.rs 🔗

@@ -579,21 +579,13 @@ impl GpuiMode {
 pub struct App {
     pub(crate) this: Weak<AppCell>,
     pub(crate) platform: Rc<dyn Platform>,
-    pub(crate) mode: GpuiMode,
     text_system: Arc<TextSystem>,
-    flushing_effects: bool,
-    pending_updates: usize,
+
     pub(crate) actions: Rc<ActionRegistry>,
     pub(crate) active_drag: Option<AnyDrag>,
     pub(crate) background_executor: BackgroundExecutor,
     pub(crate) foreground_executor: ForegroundExecutor,
-    pub(crate) loading_assets: FxHashMap<(TypeId, u64), Box<dyn Any>>,
-    asset_source: Arc<dyn AssetSource>,
-    pub(crate) svg_renderer: SvgRenderer,
-    http_client: Arc<dyn HttpClient>,
-    pub(crate) globals_by_type: FxHashMap<TypeId, Box<dyn Any>>,
     pub(crate) entities: EntityMap,
-    pub(crate) window_update_stack: Vec<WindowId>,
     pub(crate) new_entity_observers: SubscriberSet<TypeId, NewEntityListener>,
     pub(crate) windows: SlotMap<WindowId, Option<Box<Window>>>,
     pub(crate) window_handles: FxHashMap<WindowId, AnyWindowHandle>,
@@ -604,10 +596,8 @@ pub struct App {
     pub(crate) global_action_listeners:
         FxHashMap<TypeId, Vec<Rc<dyn Fn(&dyn Any, DispatchPhase, &mut Self)>>>,
     pending_effects: VecDeque<Effect>,
-    pub(crate) pending_notifications: FxHashSet<EntityId>,
-    pub(crate) pending_global_notifications: FxHashSet<TypeId>,
+
     pub(crate) observers: SubscriberSet<EntityId, Handler>,
-    // TypeId is the type of the event that the listener callback expects
     pub(crate) event_listeners: SubscriberSet<EntityId, (TypeId, Listener)>,
     pub(crate) keystroke_observers: SubscriberSet<(), KeystrokeObserver>,
     pub(crate) keystroke_interceptors: SubscriberSet<(), KeystrokeObserver>,
@@ -617,8 +607,30 @@ pub struct App {
     pub(crate) global_observers: SubscriberSet<TypeId, Handler>,
     pub(crate) quit_observers: SubscriberSet<(), QuitHandler>,
     pub(crate) restart_observers: SubscriberSet<(), Handler>,
-    pub(crate) restart_path: Option<PathBuf>,
     pub(crate) window_closed_observers: SubscriberSet<(), WindowClosedHandler>,
+
+    /// Per-App element arena. This isolates element allocations between different
+    /// App instances (important for tests where multiple Apps run concurrently).
+    pub(crate) element_arena: RefCell<Arena>,
+    /// Per-App event arena.
+    pub(crate) event_arena: Arena,
+
+    // Drop globals last. We need to ensure all tasks owned by entities and
+    // callbacks are marked cancelled at this point as this will also shutdown
+    // the tokio runtime. As any task attempting to spawn a blocking tokio task,
+    // might panic.
+    pub(crate) globals_by_type: FxHashMap<TypeId, Box<dyn Any>>,
+
+    // assets
+    pub(crate) loading_assets: FxHashMap<(TypeId, u64), Box<dyn Any>>,
+    asset_source: Arc<dyn AssetSource>,
+    pub(crate) svg_renderer: SvgRenderer,
+    http_client: Arc<dyn HttpClient>,
+
+    // below is plain data, the drop order is insignificant here
+    pub(crate) pending_notifications: FxHashSet<EntityId>,
+    pub(crate) pending_global_notifications: FxHashSet<TypeId>,
+    pub(crate) restart_path: Option<PathBuf>,
     pub(crate) layout_id_buffer: Vec<LayoutId>, // We recycle this memory across layout requests.
     pub(crate) propagate_event: bool,
     pub(crate) prompt_builder: Option<PromptBuilder>,
@@ -632,13 +644,18 @@ pub struct App {
     #[cfg(any(test, feature = "test-support", debug_assertions))]
     pub(crate) name: Option<&'static str>,
     pub(crate) text_rendering_mode: Rc<Cell<TextRenderingMode>>,
+
+    pub(crate) window_update_stack: Vec<WindowId>,
+    pub(crate) mode: GpuiMode,
+    flushing_effects: bool,
+    pending_updates: usize,
     quit_mode: QuitMode,
     quitting: bool,
-    /// Per-App element arena. This isolates element allocations between different
-    /// App instances (important for tests where multiple Apps run concurrently).
-    pub(crate) element_arena: RefCell<Arena>,
-    /// Per-App event arena.
-    pub(crate) event_arena: Arena,
+
+    // We need to ensure the leak detector drops last, after all tasks, callbacks and things have been dropped.
+    // Otherwise it may report false positives.
+    #[cfg(any(test, feature = "leak-detection"))]
+    _ref_counts: Arc<RwLock<EntityRefCounts>>,
 }
 
 impl App {
@@ -660,6 +677,9 @@ impl App {
         let keyboard_layout = platform.keyboard_layout();
         let keyboard_mapper = platform.keyboard_mapper();
 
+        #[cfg(any(test, feature = "leak-detection"))]
+        let _ref_counts = entities.ref_counts_drop_handle();
+
         let app = Rc::new_cyclic(|this| AppCell {
             app: RefCell::new(App {
                 this: this.clone(),
@@ -719,6 +739,9 @@ impl App {
                 name: None,
                 element_arena: RefCell::new(Arena::new(1024 * 1024)),
                 event_arena: Arena::new(1024 * 1024),
+
+                #[cfg(any(test, feature = "leak-detection"))]
+                _ref_counts,
             }),
         });
 

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

@@ -59,7 +59,8 @@ pub(crate) struct EntityMap {
     ref_counts: Arc<RwLock<EntityRefCounts>>,
 }
 
-struct EntityRefCounts {
+#[doc(hidden)]
+pub(crate) struct EntityRefCounts {
     counts: SlotMap<EntityId, AtomicUsize>,
     dropped_entity_ids: Vec<EntityId>,
     #[cfg(any(test, feature = "leak-detection"))]
@@ -84,7 +85,7 @@ impl EntityMap {
     }
 
     #[doc(hidden)]
-    pub fn ref_counts_drop_handle(&self) -> impl Sized + use<> {
+    pub fn ref_counts_drop_handle(&self) -> Arc<RwLock<EntityRefCounts>> {
         self.ref_counts.clone()
     }