Stop extensions' servers and message loops before removing their files (#34208)

Fixes an issue that caused Windows to fail when removing extension's
directories, as Zed had never stop any related processes.

Now:

* Zed shuts down and waits until the end when the language servers are
shut down

* Adds `impl Drop for WasmExtension` where does
`self.tx.close_channel();` to stop a receiver loop that holds the "lock"
on the extension's work dir.
The extension was dropped, but the channel was not closed for some
reason.

* Does more unregistration to ensure `Arc<WasmExtension>` with the `tx`
does not leak further

* Tidies up the related errors which had never reported a problematic
path before

Release Notes:

- N/A

---------

Co-authored-by: Smit Barmase <heysmitbarmase@gmail.com>
Co-authored-by: Smit <smit@zed.dev>

crates/extension_host/src/extension_host.rs

@@ -20,6 +20,7 @@ use extension::{
     ExtensionSnippetProxy, ExtensionThemeProxy,
 };
 use fs::{Fs, RemoveOptions};
+use futures::future::join_all;
 use futures::{
     AsyncReadExt as _, Future, FutureExt as _, StreamExt as _,
     channel::{
@@ -860,8 +861,8 @@ impl ExtensionStore {
             btree_map::Entry::Vacant(e) => e.insert(ExtensionOperation::Remove),
         };
 
-        cx.spawn(async move |this, cx| {
-            let _finish = cx.on_drop(&this, {
+        cx.spawn(async move |extension_store, cx| {
+            let _finish = cx.on_drop(&extension_store, {
                 let extension_id = extension_id.clone();
                 move |this, cx| {
                     this.outstanding_operations.remove(extension_id.as_ref());
@@ -876,22 +877,39 @@ impl ExtensionStore {
                     ignore_if_not_exists: true,
                 },
             )
-            .await?;
+            .await
+            .with_context(|| format!("Removing extension dir {extension_dir:?}"))?;
 
-            // todo(windows)
-            // Stop the server here.
-            this.update(cx, |this, cx| this.reload(None, cx))?.await;
+            extension_store
+                .update(cx, |extension_store, cx| extension_store.reload(None, cx))?
+                .await;
 
-            fs.remove_dir(
-                &work_dir,
-                RemoveOptions {
-                    recursive: true,
-                    ignore_if_not_exists: true,
-                },
-            )
-            .await?;
+            // There's a race between wasm extension fully stopping and the directory removal.
+            // On Windows, it's impossible to remove a directory that has a process running in it.
+            for i in 0..3 {
+                cx.background_executor()
+                    .timer(Duration::from_millis(i * 100))
+                    .await;
+                let removal_result = fs
+                    .remove_dir(
+                        &work_dir,
+                        RemoveOptions {
+                            recursive: true,
+                            ignore_if_not_exists: true,
+                        },
+                    )
+                    .await;
+                match removal_result {
+                    Ok(()) => break,
+                    Err(e) => {
+                        if i == 2 {
+                            log::error!("Failed to remove extension work dir {work_dir:?} : {e}");
+                        }
+                    }
+                }
+            }
 
-            this.update(cx, |_, cx| {
+            extension_store.update(cx, |_, cx| {
                 cx.emit(Event::ExtensionUninstalled(extension_id.clone()));
                 if let Some(events) = ExtensionEvents::try_global(cx) {
                     if let Some(manifest) = extension_manifest {
@@ -1143,27 +1161,38 @@ impl ExtensionStore {
             })
             .collect::<Vec<_>>();
         let mut grammars_to_remove = Vec::new();
+        let mut server_removal_tasks = Vec::with_capacity(extensions_to_unload.len());
         for extension_id in &extensions_to_unload {
             let Some(extension) = old_index.extensions.get(extension_id) else {
                 continue;
             };
             grammars_to_remove.extend(extension.manifest.grammars.keys().cloned());
-            for (language_server_name, config) in extension.manifest.language_servers.iter() {
+            for (language_server_name, config) in &extension.manifest.language_servers {
                 for language in config.languages() {
-                    self.proxy
-                        .remove_language_server(&language, language_server_name);
+                    server_removal_tasks.push(self.proxy.remove_language_server(
+                        &language,
+                        language_server_name,
+                        cx,
+                    ));
                 }
             }
 
-            for (server_id, _) in extension.manifest.context_servers.iter() {
+            for (server_id, _) in &extension.manifest.context_servers {
                 self.proxy.unregister_context_server(server_id.clone(), cx);
             }
-            for (adapter, _) in extension.manifest.debug_adapters.iter() {
+            for (adapter, _) in &extension.manifest.debug_adapters {
                 self.proxy.unregister_debug_adapter(adapter.clone());
             }
-            for (locator, _) in extension.manifest.debug_locators.iter() {
+            for (locator, _) in &extension.manifest.debug_locators {
                 self.proxy.unregister_debug_locator(locator.clone());
             }
+            for (command_name, _) in &extension.manifest.slash_commands {
+                self.proxy.unregister_slash_command(command_name.clone());
+            }
+            for (provider_id, _) in &extension.manifest.indexed_docs_providers {
+                self.proxy
+                    .unregister_indexed_docs_provider(provider_id.clone());
+            }
         }
 
         self.wasm_extensions
@@ -1268,14 +1297,15 @@ impl ExtensionStore {
             cx.background_spawn({
                 let fs = fs.clone();
                 async move {
-                    for theme_path in themes_to_add.into_iter() {
+                    let _ = join_all(server_removal_tasks).await;
+                    for theme_path in themes_to_add {
                         proxy
                             .load_user_theme(theme_path, fs.clone())
                             .await
                             .log_err();
                     }
 
-                    for (icon_theme_path, icons_root_path) in icon_themes_to_add.into_iter() {
+                    for (icon_theme_path, icons_root_path) in icon_themes_to_add {
                         proxy
                             .load_icon_theme(icon_theme_path, icons_root_path, fs.clone())
                             .await

crates/extension_host/src/extension_store_test.rs

@@ -11,6 +11,7 @@ use futures::{AsyncReadExt, StreamExt, io::BufReader};
 use gpui::{AppContext as _, SemanticVersion, TestAppContext};
 use http_client::{FakeHttpClient, Response};
 use language::{BinaryStatus, LanguageMatcher, LanguageRegistry};
+use language_extension::LspAccess;
 use lsp::LanguageServerName;
 use node_runtime::NodeRuntime;
 use parking_lot::Mutex;
@@ -271,7 +272,7 @@ async fn test_extension_store(cx: &mut TestAppContext) {
     let theme_registry = Arc::new(ThemeRegistry::new(Box::new(())));
     theme_extension::init(proxy.clone(), theme_registry.clone(), cx.executor());
     let language_registry = Arc::new(LanguageRegistry::test(cx.executor()));
-    language_extension::init(proxy.clone(), language_registry.clone());
+    language_extension::init(LspAccess::Noop, proxy.clone(), language_registry.clone());
     let node_runtime = NodeRuntime::unavailable();
 
     let store = cx.new(|cx| {
@@ -554,7 +555,11 @@ async fn test_extension_store_with_test_extension(cx: &mut TestAppContext) {
     let theme_registry = Arc::new(ThemeRegistry::new(Box::new(())));
     theme_extension::init(proxy.clone(), theme_registry.clone(), cx.executor());
     let language_registry = project.read_with(cx, |project, _cx| project.languages().clone());
-    language_extension::init(proxy.clone(), language_registry.clone());
+    language_extension::init(
+        LspAccess::ViaLspStore(project.update(cx, |project, _| project.lsp_store())),
+        proxy.clone(),
+        language_registry.clone(),
+    );
     let node_runtime = NodeRuntime::unavailable();
 
     let mut status_updates = language_registry.language_server_binary_statuses();
@@ -815,7 +820,6 @@ async fn test_extension_store_with_test_extension(cx: &mut TestAppContext) {
     extension_store
         .update(cx, |store, cx| store.reload(Some("gleam".into()), cx))
         .await;
-
     cx.executor().run_until_parked();
     project.update(cx, |project, cx| {
         project.restart_language_servers_for_buffers(vec![buffer.clone()], HashSet::default(), cx)

crates/extension_host/src/headless_host.rs

@@ -11,6 +11,7 @@ use extension::{
     ExtensionLanguageServerProxy, ExtensionManifest,
 };
 use fs::{Fs, RemoveOptions, RenameOptions};
+use futures::future::join_all;
 use gpui::{App, AppContext as _, AsyncApp, Context, Entity, Task, WeakEntity};
 use http_client::HttpClient;
 use language::{LanguageConfig, LanguageName, LanguageQueries, LoadedLanguage};
@@ -230,18 +231,27 @@ impl HeadlessExtensionStore {
             .unwrap_or_default();
         self.proxy.remove_languages(&languages_to_remove, &[]);
 
-        for (language_server_name, language) in self
+        let servers_to_remove = self
             .loaded_language_servers
             .remove(extension_id)
-            .unwrap_or_default()
-        {
-            self.proxy
-                .remove_language_server(&language, &language_server_name);
-        }
-
+            .unwrap_or_default();
+        let proxy = self.proxy.clone();
         let path = self.extension_dir.join(&extension_id.to_string());
         let fs = self.fs.clone();
-        cx.spawn(async move |_, _| {
+        cx.spawn(async move |_, cx| {
+            let mut removal_tasks = Vec::with_capacity(servers_to_remove.len());
+            cx.update(|cx| {
+                for (language_server_name, language) in servers_to_remove {
+                    removal_tasks.push(proxy.remove_language_server(
+                        &language,
+                        &language_server_name,
+                        cx,
+                    ));
+                }
+            })
+            .ok();
+            let _ = join_all(removal_tasks).await;
+
             fs.remove_dir(
                 &path,
                 RemoveOptions {
@@ -250,6 +260,7 @@ impl HeadlessExtensionStore {
                 },
             )
             .await
+            .with_context(|| format!("Removing directory {path:?}"))
         })
     }
 

crates/extension_host/src/wasm_host.rs

@@ -54,7 +54,7 @@ pub struct WasmHost {
     main_thread_message_tx: mpsc::UnboundedSender<MainThreadCall>,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct WasmExtension {
     tx: UnboundedSender<ExtensionCall>,
     pub manifest: Arc<ExtensionManifest>,
@@ -63,6 +63,12 @@ pub struct WasmExtension {
     pub zed_api_version: SemanticVersion,
 }
 
+impl Drop for WasmExtension {
+    fn drop(&mut self) {
+        self.tx.close_channel();
+    }
+}
+
 #[async_trait]
 impl extension::Extension for WasmExtension {
     fn manifest(&self) -> Arc<ExtensionManifest> {
@@ -742,7 +748,6 @@ impl WasmExtension {
     {
         let (return_tx, return_rx) = oneshot::channel();
         self.tx
-            .clone()
             .unbounded_send(Box::new(move |extension, store| {
                 async {
                     let result = f(extension, store).await;