repl: Avoid duplicate kernelspec fetches (#51024)

MostlyK created 3 weeks ago

Switch python_env_kernel_specifications to a buffered stream (using
buffer_unordered(4) and filter_map) to run discovery concurrently with a
bounded parallelism and skip None results.

Closes #50988

Before you mark this PR as ready for review, make sure that you have:
- [x] Added a solid test coverage and/or screenshots from doing manual
testing
- [x] Done a self-review taking into account security and performance
aspects
- [x] Aligned any UI changes with the [UI
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)

Release Notes:

- N/A

Change summary

crates/repl/src/kernels/mod.rs | 10 +++++-----
crates/repl/src/repl_store.rs  | 20 ++++++++++++++++----
2 files changed, 21 insertions(+), 9 deletions(-)

Detailed changes

crates/repl/src/kernels/mod.rs

@@ -503,11 +503,11 @@ pub fn python_env_kernel_specifications(
             });
 
         #[allow(unused_mut)]
-        let mut kernel_specs: Vec<KernelSpecification> = futures::future::join_all(kernelspecs)
-            .await
-            .into_iter()
-            .flatten()
-            .collect();
+        let mut kernel_specs: Vec<KernelSpecification> = futures::stream::iter(kernelspecs)
+            .buffer_unordered(4)
+            .filter_map(|x| async move { x })
+            .collect::<Vec<_>>()
+            .await;
 
         #[cfg(target_os = "windows")]
         if kernel_specs.is_empty() && !is_remote {

crates/repl/src/repl_store.rs

@@ -32,6 +32,7 @@ pub struct ReplStore {
     kernel_specifications_for_worktree: HashMap<WorktreeId, Vec<KernelSpecification>>,
     active_python_toolchain_for_worktree: HashMap<WorktreeId, SharedString>,
     remote_worktrees: HashSet<WorktreeId>,
+    fetching_python_kernelspecs: HashSet<WorktreeId>,
     _subscriptions: Vec<Subscription>,
 }
 
@@ -66,6 +67,7 @@ impl ReplStore {
             selected_kernel_for_worktree: HashMap::default(),
             active_python_toolchain_for_worktree: HashMap::default(),
             remote_worktrees: HashSet::default(),
+            fetching_python_kernelspecs: HashSet::default(),
         };
         this.on_enabled_changed(cx);
         this
@@ -140,6 +142,10 @@ impl ReplStore {
         project: &Entity<Project>,
         cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
+        if !self.fetching_python_kernelspecs.insert(worktree_id) {
+            return Task::ready(Ok(()));
+        }
+
         let is_remote = project.read(cx).is_remote();
         // WSL does require access to global kernel specs, so we only exclude remote worktrees that aren't WSL.
         // TODO: a better way to handle WSL vs SSH/remote projects,
@@ -149,7 +155,7 @@ impl ReplStore {
             .map_or(false, |opts| {
                 matches!(opts, RemoteConnectionOptions::Wsl(_))
             });
-        let kernel_specifications = python_env_kernel_specifications(project, worktree_id, cx);
+        let kernel_specifications_task = python_env_kernel_specifications(project, worktree_id, cx);
         let active_toolchain = project.read(cx).active_toolchain(
             ProjectPath {
                 worktree_id,
@@ -160,9 +166,15 @@ impl ReplStore {
         );
 
         cx.spawn(async move |this, cx| {
-            let kernel_specifications = kernel_specifications
-                .await
-                .context("getting python kernelspecs")?;
+            let kernel_specifications_res = kernel_specifications_task.await;
+
+            this.update(cx, |this, _cx| {
+                this.fetching_python_kernelspecs.remove(&worktree_id);
+            })
+            .ok();
+
+            let kernel_specifications =
+                kernel_specifications_res.context("getting python kernelspecs")?;
 
             let active_toolchain_path = active_toolchain.await.map(|toolchain| toolchain.path);