repl: Shutdown all kernels on app quit (#48760)

Kyle Kelley created 2 months ago

Closes #17985
Closes #21911

Force the shutdown of the kernel by ensuring the kernel sessions are
dropped on app quit.

Release Notes:

- Fixed shutdown of kernels on app exit

Change summary

crates/repl/src/kernels/mod.rs            |  1 
crates/repl/src/kernels/native_kernel.rs  | 10 ++++++--
crates/repl/src/kernels/remote_kernels.rs |  4 +++
crates/repl/src/repl_store.rs             | 29 +++++++++++++++++++++---
4 files changed, 37 insertions(+), 7 deletions(-)

Detailed changes

crates/repl/src/kernels/mod.rs

@@ -170,6 +170,7 @@ pub trait RunningKernel: Send + Debug {
     fn kernel_info(&self) -> Option<&KernelInfoReply>;
     fn set_kernel_info(&mut self, info: KernelInfoReply);
     fn force_shutdown(&mut self, window: &mut Window, cx: &mut App) -> Task<anyhow::Result<()>>;
+    fn kill(&mut self);
 }
 
 #[derive(Debug, Clone)]

crates/repl/src/kernels/native_kernel.rs

@@ -370,17 +370,21 @@ impl RunningKernel for NativeRunningKernel {
     }
 
     fn force_shutdown(&mut self, _window: &mut Window, _cx: &mut App) -> Task<anyhow::Result<()>> {
+        self.kill();
+        Task::ready(Ok(()))
+    }
+
+    fn kill(&mut self) {
         self._process_status_task.take();
         self.request_tx.close_channel();
-        Task::ready(self.process.kill().context("killing the kernel process"))
+        self.process.kill().ok();
     }
 }
 
 impl Drop for NativeRunningKernel {
     fn drop(&mut self) {
         std::fs::remove_file(&self.connection_path).ok();
-        self.request_tx.close_channel();
-        self.process.kill().ok();
+        self.kill();
     }
 }
 

crates/repl/src/kernels/remote_kernels.rs

@@ -289,4 +289,8 @@ impl RunningKernel for RemoteRunningKernel {
             Ok(())
         })
     }
+
+    fn kill(&mut self) {
+        self.request_tx.close_channel();
+    }
 }

crates/repl/src/repl_store.rs

@@ -1,3 +1,4 @@
+use std::future::Future;
 use std::sync::Arc;
 
 use anyhow::{Context as _, Result};
@@ -10,7 +11,7 @@ use project::{Fs, Project, WorktreeId};
 use settings::{Settings, SettingsStore};
 
 use crate::kernels::{
-    list_remote_kernelspecs, local_kernel_specifications, python_env_kernel_specifications,
+    Kernel, list_remote_kernelspecs, local_kernel_specifications, python_env_kernel_specifications,
 };
 use crate::{JupyterSettings, KernelSpecification, Session};
 
@@ -47,9 +48,12 @@ impl ReplStore {
     }
 
     pub fn new(fs: Arc<dyn Fs>, cx: &mut Context<Self>) -> Self {
-        let subscriptions = vec![cx.observe_global::<SettingsStore>(move |this, cx| {
-            this.set_enabled(JupyterSettings::enabled(cx), cx);
-        })];
+        let subscriptions = vec![
+            cx.observe_global::<SettingsStore>(move |this, cx| {
+                this.set_enabled(JupyterSettings::enabled(cx), cx);
+            }),
+            cx.on_app_quit(Self::shutdown_all_sessions),
+        ];
 
         let this = Self {
             fs,
@@ -281,6 +285,23 @@ impl ReplStore {
         self.sessions.remove(&entity_id);
     }
 
+    fn shutdown_all_sessions(
+        &mut self,
+        cx: &mut Context<Self>,
+    ) -> impl Future<Output = ()> + use<> {
+        for session in self.sessions.values() {
+            session.update(cx, |session, _cx| {
+                if let Kernel::RunningKernel(mut kernel) =
+                    std::mem::replace(&mut session.kernel, Kernel::Shutdown)
+                {
+                    kernel.kill();
+                }
+            });
+        }
+        self.sessions.clear();
+        futures::future::ready(())
+    }
+
     #[cfg(test)]
     pub fn set_kernel_specs_for_testing(
         &mut self,