Unregister worktree when its last handle to it gets released

crates/project/src/project.rs

@@ -160,16 +160,13 @@ impl Project {
 
                                 if let Some(project_id) = remote_id {
                                     let mut registrations = Vec::new();
-                                    this.read_with(&cx, |this, cx| {
+                                    this.update(&mut cx, |this, cx| {
                                         for worktree in &this.worktrees {
-                                            let worktree_id = worktree.id() as u64;
-                                            let worktree = worktree.read(cx).as_local().unwrap();
-                                            registrations.push(rpc.request(
-                                                proto::RegisterWorktree {
-                                                    project_id,
-                                                    worktree_id,
-                                                    root_name: worktree.root_name().to_string(),
-                                                    authorized_logins: worktree.authorized_logins(),
+                                            registrations.push(worktree.update(
+                                                cx,
+                                                |worktree, cx| {
+                                                    let worktree = worktree.as_local_mut().unwrap();
+                                                    worktree.register(project_id, cx)
                                                 },
                                             ));
                                         }
@@ -407,7 +404,7 @@ impl Project {
                 for worktree in &this.worktrees {
                     worktree.update(cx, |worktree, cx| {
                         let worktree = worktree.as_local_mut().unwrap();
-                        tasks.push(worktree.share(project_id, cx));
+                        tasks.push(worktree.share(cx));
                     });
                 }
             });
@@ -897,21 +894,15 @@ impl Project {
             });
 
             if let Some(project_id) = remote_project_id {
-                let worktree_id = worktree.id() as u64;
-                let register_message = worktree.update(&mut cx, |worktree, _| {
-                    let worktree = worktree.as_local_mut().unwrap();
-                    proto::RegisterWorktree {
-                        project_id,
-                        worktree_id,
-                        root_name: worktree.root_name().to_string(),
-                        authorized_logins: worktree.authorized_logins(),
-                    }
-                });
-                client.request(register_message).await?;
+                worktree
+                    .update(&mut cx, |worktree, cx| {
+                        worktree.as_local_mut().unwrap().register(project_id, cx)
+                    })
+                    .await?;
                 if is_shared {
                     worktree
                         .update(&mut cx, |worktree, cx| {
-                            worktree.as_local_mut().unwrap().share(project_id, cx)
+                            worktree.as_local_mut().unwrap().share(cx)
                         })
                         .await?;
                 }
@@ -921,6 +912,12 @@ impl Project {
         })
     }
 
+    pub fn remove_worktree(&mut self, id: WorktreeId, cx: &mut ModelContext<Self>) {
+        self.worktrees
+            .retain(|worktree| worktree.read(cx).id() != id);
+        cx.notify();
+    }
+
     fn add_worktree(&mut self, worktree: ModelHandle<Worktree>, cx: &mut ModelContext<Self>) {
         cx.observe(&worktree, |_, _, cx| cx.notify()).detach();
         self.worktrees.push(worktree);
@@ -1104,9 +1101,7 @@ impl Project {
         cx: &mut ModelContext<Self>,
     ) -> Result<()> {
         let worktree_id = WorktreeId::from_proto(envelope.payload.worktree_id);
-        self.worktrees
-            .retain(|worktree| worktree.read(cx).as_remote().unwrap().id() != worktree_id);
-        cx.notify();
+        self.remove_worktree(worktree_id, cx);
         Ok(())
     }
 

crates/project/src/worktree.rs

@@ -66,6 +66,24 @@ pub enum Worktree {
 
 impl Entity for Worktree {
     type Event = ();
+
+    fn release(&mut self, cx: &mut MutableAppContext) {
+        if let Some(worktree) = self.as_local_mut() {
+            if let Registration::Done { project_id } = worktree.registration {
+                let client = worktree.client.clone();
+                let unregister_message = proto::UnregisterWorktree {
+                    project_id,
+                    worktree_id: worktree.id().to_proto(),
+                };
+                cx.foreground()
+                    .spawn(async move {
+                        client.send(unregister_message).await?;
+                        Ok::<_, anyhow::Error>(())
+                    })
+                    .detach_and_log_err(cx);
+            }
+        }
+    }
 }
 
 impl Worktree {
@@ -747,6 +765,7 @@ pub struct LocalWorktree {
     last_scan_state_rx: watch::Receiver<ScanState>,
     _background_scanner_task: Option<Task<()>>,
     poll_task: Option<Task<()>>,
+    registration: Registration,
     share: Option<ShareState>,
     loading_buffers: LoadingBuffers,
     open_buffers: HashMap<usize, WeakModelHandle<Buffer>>,
@@ -759,6 +778,13 @@ pub struct LocalWorktree {
     fs: Arc<dyn Fs>,
 }
 
+#[derive(Debug, Eq, PartialEq)]
+enum Registration {
+    None,
+    Pending,
+    Done { project_id: u64 },
+}
+
 struct ShareState {
     project_id: u64,
     snapshots_tx: Sender<Snapshot>,
@@ -851,6 +877,7 @@ impl LocalWorktree {
                 background_snapshot: Arc::new(Mutex::new(snapshot)),
                 last_scan_state_rx,
                 _background_scanner_task: None,
+                registration: Registration::None,
                 share: None,
                 poll_task: None,
                 loading_buffers: Default::default(),
@@ -1316,11 +1343,48 @@ impl LocalWorktree {
         })
     }
 
-    pub fn share(
+    pub fn register(
         &mut self,
         project_id: u64,
         cx: &mut ModelContext<Worktree>,
     ) -> Task<anyhow::Result<()>> {
+        if self.registration != Registration::None {
+            return Task::ready(Ok(()));
+        }
+
+        self.registration = Registration::Pending;
+        let client = self.client.clone();
+        let register_message = proto::RegisterWorktree {
+            project_id,
+            worktree_id: self.id().to_proto(),
+            root_name: self.root_name().to_string(),
+            authorized_logins: self.authorized_logins(),
+        };
+        cx.spawn(|this, mut cx| async move {
+            let response = client.request(register_message).await;
+            this.update(&mut cx, |this, _| {
+                let worktree = this.as_local_mut().unwrap();
+                match response {
+                    Ok(_) => {
+                        worktree.registration = Registration::Done { project_id };
+                        Ok(())
+                    }
+                    Err(error) => {
+                        worktree.registration = Registration::None;
+                        Err(error)
+                    }
+                }
+            })
+        })
+    }
+
+    pub fn share(&mut self, cx: &mut ModelContext<Worktree>) -> Task<anyhow::Result<()>> {
+        let project_id = if let Registration::Done { project_id } = self.registration {
+            project_id
+        } else {
+            return Task::ready(Err(anyhow!("cannot share worktree before registering it")));
+        };
+
         if self.share.is_some() {
             return Task::ready(Ok(()));
         }