From 58eefcd33102879011aa0228934858a441167804 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Mon, 5 Jul 2021 17:37:19 -0700
Subject: [PATCH] Add ModelHandle::spawn_weak, fix worktree leak from spawn
 call

---
 gpui/src/app.rs     | 10 ++++++++++
 zed/src/worktree.rs | 48 +++++++++++++++++++--------------------------
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/gpui/src/app.rs b/gpui/src/app.rs
index 0470881dd52a3bd6c71685a5288065dcbe88622a..f499207a09a19337ad19856114f2e1ff6b8f4fe1 100644
--- a/gpui/src/app.rs
+++ b/gpui/src/app.rs
@@ -1761,6 +1761,16 @@ impl<'a, T: Entity> ModelContext<'a, T> {
         let handle = self.handle();
         self.app.spawn(|cx| f(handle, cx))
     }
+
+    pub fn spawn_weak<F, Fut, S>(&self, f: F) -> Task<S>
+    where
+        F: FnOnce(WeakModelHandle<T>, AsyncAppContext) -> Fut,
+        Fut: 'static + Future<Output = S>,
+        S: 'static,
+    {
+        let handle = self.handle().downgrade();
+        self.app.spawn(|cx| f(handle, cx))
+    }
 }
 
 impl<M> AsRef<AppContext> for ModelContext<'_, M> {
diff --git a/zed/src/worktree.rs b/zed/src/worktree.rs
index 99a66efd393bd8deaa43ec4e02c8ca4e7d2e04c2..2d436ab669fedc2a1afda5569d3adb0929c86b2f 100644
--- a/zed/src/worktree.rs
+++ b/zed/src/worktree.rs
@@ -196,14 +196,18 @@ impl Worktree {
                     })
                     .detach();
 
-                cx.spawn(|this, mut cx| async move {
+                cx.spawn_weak(|this, mut cx| async move {
                     while let Some(snapshot) = snapshot_rx.recv().await {
-                        this.update(&mut cx, |this, cx| {
-                            let this = this.as_remote_mut().unwrap();
-                            this.snapshot = snapshot;
-                            cx.notify();
-                            this.update_open_buffers(cx);
-                        })
+                        if let Some(this) = cx.read(|cx| this.upgrade(cx)) {
+                            this.update(&mut cx, |this, cx| {
+                                let this = this.as_remote_mut().unwrap();
+                                this.snapshot = snapshot;
+                                cx.notify();
+                                this.update_open_buffers(cx);
+                            });
+                        } else {
+                            break;
+                        }
                     }
                 })
                 .detach();
@@ -460,28 +464,16 @@ impl LocalWorktree {
             scanner.run(event_stream)
         });
 
-        cx.spawn(|this, mut cx| {
-            let this = this.downgrade();
-            async move {
-                while let Ok(scan_state) = scan_state_rx.recv().await {
-                    let alive = cx.update(|cx| {
-                        if let Some(handle) = this.upgrade(&cx) {
-                            handle.update(cx, |this, cx| {
-                                if let Worktree::Local(worktree) = this {
-                                    worktree.observe_scan_state(scan_state, cx)
-                                } else {
-                                    unreachable!()
-                                }
-                            });
-                            true
-                        } else {
-                            false
-                        }
+        cx.spawn_weak(|this, mut cx| async move {
+            while let Ok(scan_state) = scan_state_rx.recv().await {
+                if let Some(handle) = cx.read(|cx| this.upgrade(&cx)) {
+                    handle.update(&mut cx, |this, cx| {
+                        this.as_local_mut()
+                            .unwrap()
+                            .observe_scan_state(scan_state, cx)
                     });
-
-                    if !alive {
-                        break;
-                    }
+                } else {
+                    break;
                 }
             }
         })