Fix leaked editor (#25530)

Closes #ISSUE

Release Notes:

- Fixed a bug that would prevent rejoining projects sometimes

crates/gpui/Cargo.toml

@@ -13,7 +13,7 @@ workspace = true
 [features]
 default = ["http_client", "font-kit", "wayland", "x11"]
 test-support = [
-    "backtrace",
+    "leak-detection",
     "collections/test-support",
     "rand",
     "util/test-support",
@@ -21,6 +21,7 @@ test-support = [
     "wayland",
     "x11",
 ]
+leak-detection = ["backtrace"]
 runtime_shaders = []
 macos-blade = [
     "blade-graphics",

crates/gpui/src/app/entity_map.rs

@@ -19,7 +19,7 @@ use std::{
     thread::panicking,
 };
 
-#[cfg(any(test, feature = "test-support"))]
+#[cfg(any(test, feature = "leak-detection"))]
 use collections::HashMap;
 
 use super::Context;
@@ -62,7 +62,7 @@ pub(crate) struct EntityMap {
 struct EntityRefCounts {
     counts: SlotMap<EntityId, AtomicUsize>,
     dropped_entity_ids: Vec<EntityId>,
-    #[cfg(any(test, feature = "test-support"))]
+    #[cfg(any(test, feature = "leak-detection"))]
     leak_detector: LeakDetector,
 }
 
@@ -74,7 +74,7 @@ impl EntityMap {
             ref_counts: Arc::new(RwLock::new(EntityRefCounts {
                 counts: SlotMap::with_key(),
                 dropped_entity_ids: Vec::new(),
-                #[cfg(any(test, feature = "test-support"))]
+                #[cfg(any(test, feature = "leak-detection"))]
                 leak_detector: LeakDetector {
                     next_handle_id: 0,
                     entity_handles: HashMap::default(),
@@ -221,7 +221,7 @@ pub struct AnyEntity {
     pub(crate) entity_id: EntityId,
     pub(crate) entity_type: TypeId,
     entity_map: Weak<RwLock<EntityRefCounts>>,
-    #[cfg(any(test, feature = "test-support"))]
+    #[cfg(any(test, feature = "leak-detection"))]
     handle_id: HandleId,
 }
 
@@ -231,7 +231,7 @@ impl AnyEntity {
             entity_id: id,
             entity_type,
             entity_map: entity_map.clone(),
-            #[cfg(any(test, feature = "test-support"))]
+            #[cfg(any(test, feature = "leak-detection"))]
             handle_id: entity_map
                 .upgrade()
                 .unwrap()
@@ -290,7 +290,7 @@ impl Clone for AnyEntity {
             entity_id: self.entity_id,
             entity_type: self.entity_type,
             entity_map: self.entity_map.clone(),
-            #[cfg(any(test, feature = "test-support"))]
+            #[cfg(any(test, feature = "leak-detection"))]
             handle_id: self
                 .entity_map
                 .upgrade()
@@ -319,7 +319,7 @@ impl Drop for AnyEntity {
             }
         }
 
-        #[cfg(any(test, feature = "test-support"))]
+        #[cfg(any(test, feature = "leak-detection"))]
         if let Some(entity_map) = self.entity_map.upgrade() {
             entity_map
                 .write()
@@ -535,7 +535,7 @@ impl AnyWeakEntity {
             entity_id: self.entity_id,
             entity_type: self.entity_type,
             entity_map: self.entity_ref_counts.clone(),
-            #[cfg(any(test, feature = "test-support"))]
+            #[cfg(any(test, feature = "leak-detection"))]
             handle_id: self
                 .entity_ref_counts
                 .upgrade()
@@ -547,7 +547,7 @@ impl AnyWeakEntity {
     }
 
     /// Assert that entity referenced by this weak handle has been released.
-    #[cfg(any(test, feature = "test-support"))]
+    #[cfg(any(test, feature = "leak-detection"))]
     pub fn assert_released(&self) {
         self.entity_ref_counts
             .upgrade()
@@ -710,23 +710,23 @@ impl<T> PartialEq<Entity<T>> for WeakEntity<T> {
     }
 }
 
-#[cfg(any(test, feature = "test-support"))]
+#[cfg(any(test, feature = "leak-detection"))]
 static LEAK_BACKTRACE: std::sync::LazyLock<bool> =
     std::sync::LazyLock::new(|| std::env::var("LEAK_BACKTRACE").map_or(false, |b| !b.is_empty()));
 
-#[cfg(any(test, feature = "test-support"))]
+#[cfg(any(test, feature = "leak-detection"))]
 #[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
 pub(crate) struct HandleId {
     id: u64, // id of the handle itself, not the pointed at object
 }
 
-#[cfg(any(test, feature = "test-support"))]
+#[cfg(any(test, feature = "leak-detection"))]
 pub(crate) struct LeakDetector {
     next_handle_id: u64,
     entity_handles: HashMap<EntityId, HashMap<HandleId, Option<backtrace::Backtrace>>>,
 }
 
-#[cfg(any(test, feature = "test-support"))]
+#[cfg(any(test, feature = "leak-detection"))]
 impl LeakDetector {
     #[track_caller]
     pub fn handle_created(&mut self, entity_id: EntityId) -> HandleId {

crates/zed/src/zed/inline_completion_registry.rs

@@ -267,7 +267,7 @@ fn assign_edit_prediction_provider(
                 }
 
                 let zeta = zeta::Zeta::register(
-                    Some(cx.entity()),
+                    editor.workspace().map(|w| w.downgrade()),
                     worktree,
                     client.clone(),
                     user_store,

crates/zeta/src/zeta.rs

@@ -9,7 +9,6 @@ mod rate_completion_modal;
 
 pub(crate) use completion_diff_element::*;
 use db::kvp::KEY_VALUE_STORE;
-use editor::Editor;
 pub use init::*;
 use inline_completion::DataCollectionState;
 pub use license_detection::is_license_eligible_for_data_collection;
@@ -24,7 +23,7 @@ use collections::{HashMap, HashSet, VecDeque};
 use futures::AsyncReadExt;
 use gpui::{
     actions, App, AppContext as _, AsyncApp, Context, Entity, EntityId, Global, SemanticVersion,
-    Subscription, Task,
+    Subscription, Task, WeakEntity,
 };
 use http_client::{HttpClient, Method};
 use input_excerpt::excerpt_for_cursor_position;
@@ -186,7 +185,7 @@ impl std::fmt::Debug for InlineCompletion {
 }
 
 pub struct Zeta {
-    editor: Option<Entity<Editor>>,
+    workspace: Option<WeakEntity<Workspace>>,
     client: Arc<Client>,
     events: VecDeque<Event>,
     registered_buffers: HashMap<gpui::EntityId, RegisteredBuffer>,
@@ -209,14 +208,14 @@ impl Zeta {
     }
 
     pub fn register(
-        editor: Option<Entity<Editor>>,
+        workspace: Option<WeakEntity<Workspace>>,
         worktree: Option<Entity<Worktree>>,
         client: Arc<Client>,
         user_store: Entity<UserStore>,
         cx: &mut App,
     ) -> Entity<Self> {
         let this = Self::global(cx).unwrap_or_else(|| {
-            let entity = cx.new(|cx| Self::new(editor, client, user_store, cx));
+            let entity = cx.new(|cx| Self::new(workspace, client, user_store, cx));
             cx.set_global(ZetaGlobal(entity.clone()));
             entity
         });
@@ -239,7 +238,7 @@ impl Zeta {
     }
 
     fn new(
-        editor: Option<Entity<Editor>>,
+        workspace: Option<WeakEntity<Workspace>>,
         client: Arc<Client>,
         user_store: Entity<UserStore>,
         cx: &mut Context<Self>,
@@ -250,7 +249,7 @@ impl Zeta {
         let data_collection_choice = cx.new(|_| data_collection_choice);
 
         Self {
-            editor,
+            workspace,
             client,
             events: VecDeque::new(),
             shown_completions: VecDeque::new(),
@@ -705,10 +704,7 @@ and then another
         can_collect_data: bool,
         cx: &mut Context<Self>,
     ) -> Task<Result<Option<InlineCompletion>>> {
-        let workspace = self
-            .editor
-            .as_ref()
-            .and_then(|editor| editor.read(cx).workspace());
+        let workspace = self.workspace.as_ref().and_then(|w| w.upgrade());
         self.request_completion_impl(
             workspace,
             project,