From be7b24fcf74c2c2f219976c8d0e06b103cb3964b Mon Sep 17 00:00:00 2001
From: Thorsten Ball <mrnugget@gmail.com>
Date: Tue, 15 Oct 2024 11:26:23 +0200
Subject: [PATCH] ssh remote: Shutdown SSH & server process correctly on app
 quit (#19210)

Release Notes:

- N/A

Co-authored-by: Bennet <bennet@zed.dev>
---
 crates/project/src/project.rs    | 42 +++++++++++++++------
 crates/remote/src/ssh_session.rs | 65 +++++++++++++++++---------------
 2 files changed, 65 insertions(+), 42 deletions(-)

diff --git a/crates/project/src/project.rs b/crates/project/src/project.rs
index 68f07f1ca9191..2692c711bf531 100644
--- a/crates/project/src/project.rs
+++ b/crates/project/src/project.rs
@@ -746,17 +746,6 @@ impl Project {
             });
             cx.subscribe(&lsp_store, Self::on_lsp_store_event).detach();
 
-            cx.on_release(|this, cx| {
-                if let Some(ssh_client) = this.ssh_client.as_ref() {
-                    ssh_client
-                        .read(cx)
-                        .to_proto_client()
-                        .send(proto::ShutdownRemoteServer {})
-                        .log_err();
-                }
-            })
-            .detach();
-
             cx.subscribe(&ssh, Self::on_ssh_event).detach();
             cx.observe(&ssh, |_, _, cx| cx.notify()).detach();
 
@@ -769,7 +758,22 @@ impl Project {
                 join_project_response_message_id: 0,
                 client_state: ProjectClientState::Local,
                 client_subscriptions: Vec::new(),
-                _subscriptions: vec![cx.on_release(Self::release)],
+                _subscriptions: vec![
+                    cx.on_release(Self::release),
+                    cx.on_app_quit(|this, cx| {
+                        let shutdown = this.ssh_client.take().and_then(|client| {
+                            client
+                                .read(cx)
+                                .shutdown_processes(Some(proto::ShutdownRemoteServer {}))
+                        });
+
+                        cx.background_executor().spawn(async move {
+                            if let Some(shutdown) = shutdown {
+                                shutdown.await;
+                            }
+                        })
+                    }),
+                ],
                 active_entry: None,
                 snippets,
                 languages,
@@ -1094,6 +1098,20 @@ impl Project {
     }
 
     fn release(&mut self, cx: &mut AppContext) {
+        if let Some(client) = self.ssh_client.take() {
+            let shutdown = client
+                .read(cx)
+                .shutdown_processes(Some(proto::ShutdownRemoteServer {}));
+
+            cx.background_executor()
+                .spawn(async move {
+                    if let Some(shutdown) = shutdown {
+                        shutdown.await;
+                    }
+                })
+                .detach()
+        }
+
         match &self.client_state {
             ProjectClientState::Local => {}
             ProjectClientState::Shared { .. } => {
diff --git a/crates/remote/src/ssh_session.rs b/crates/remote/src/ssh_session.rs
index 7a2073bb55ef9..ccb66a0c626e0 100644
--- a/crates/remote/src/ssh_session.rs
+++ b/crates/remote/src/ssh_session.rs
@@ -424,12 +424,6 @@ pub struct SshRemoteClient {
     state: Arc<Mutex<Option<State>>>,
 }
 
-impl Drop for SshRemoteClient {
-    fn drop(&mut self) {
-        self.shutdown_processes();
-    }
-}
-
 #[derive(Debug)]
 pub enum SshRemoteEvent {
     Disconnected,
@@ -449,19 +443,11 @@ impl SshRemoteClient {
             let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
 
             let client = cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx))?;
-            let this = cx.new_model(|cx| {
-                cx.on_app_quit(|this: &mut Self, _| {
-                    this.shutdown_processes();
-                    futures::future::ready(())
-                })
-                .detach();
-
-                Self {
-                    client: client.clone(),
-                    unique_identifier: unique_identifier.clone(),
-                    connection_options: connection_options.clone(),
-                    state: Arc::new(Mutex::new(Some(State::Connecting))),
-                }
+            let this = cx.new_model(|_| Self {
+                client: client.clone(),
+                unique_identifier: unique_identifier.clone(),
+                connection_options: connection_options.clone(),
+                state: Arc::new(Mutex::new(Some(State::Connecting))),
             })?;
 
             let (proxy, proxy_incoming_tx, proxy_outgoing_rx) =
@@ -506,25 +492,44 @@ impl SshRemoteClient {
         })
     }
 
-    fn shutdown_processes(&self) {
-        let Some(state) = self.state.lock().take() else {
-            return;
-        };
+    pub fn shutdown_processes<T: RequestMessage>(
+        &self,
+        shutdown_request: Option<T>,
+    ) -> Option<impl Future<Output = ()>> {
+        let state = self.state.lock().take()?;
         log::info!("shutting down ssh processes");
 
         let State::Connected {
             multiplex_task,
             heartbeat_task,
-            ..
+            ssh_connection,
+            delegate,
+            forwarder,
         } = state
         else {
-            return;
+            return None;
         };
-        // Drop `multiplex_task` because it owns our ssh_proxy_process, which is a
-        // child of master_process.
-        drop(multiplex_task);
-        // Now drop the rest of state, which kills master process.
-        drop(heartbeat_task);
+
+        let client = self.client.clone();
+
+        Some(async move {
+            if let Some(shutdown_request) = shutdown_request {
+                client.send(shutdown_request).log_err();
+                // We wait 50ms instead of waiting for a response, because
+                // waiting for a response would require us to wait on the main thread
+                // which we want to avoid in an `on_app_quit` callback.
+                Timer::after(Duration::from_millis(50)).await;
+            }
+
+            // Drop `multiplex_task` because it owns our ssh_proxy_process, which is a
+            // child of master_process.
+            drop(multiplex_task);
+            // Now drop the rest of state, which kills master process.
+            drop(heartbeat_task);
+            drop(ssh_connection);
+            drop(delegate);
+            drop(forwarder);
+        })
     }
 
     fn reconnect(&mut self, cx: &mut ModelContext<Self>) -> Result<()> {