join router threads on exit

4 files changed, 53 insertions, 43 deletions

diff --git a/zellij-client/src/lib.rs b/zellij-client/src/lib.rs
index 7042bed72..45b60d44a 100644
--- a/zellij-client/src/lib.rs
+++ b/zellij-client/src/lib.rs
@@ -241,12 +241,10 @@ pub fn start_client(
         err_ctx.add_call(ContextType::Client((&client_instruction).into()));
         match client_instruction {
             ClientInstruction::Exit(reason) => {
-                match reason {
-                    ExitReason::Error(_) => handle_error(format!("{}", reason)),
-                    ExitReason::ForceDetached => {
-                        os_input.send_to_server(ClientToServerMsg::ClientDetached);
-                    }
-                    _ => {}
+                os_input.send_to_server(ClientToServerMsg::ClientExited);
+
+                if let ExitReason::Error(_) = reason {
+                    handle_error(format!("{}", reason));
                 }
                 exit_msg = format!("{}", reason);
                 break;
diff --git a/zellij-server/src/lib.rs b/zellij-server/src/lib.rs
index e598d2384..7a5aed1ec 100644
--- a/zellij-server/src/lib.rs
+++ b/zellij-server/src/lib.rs
@@ -11,7 +11,7 @@ mod wasm_vm;
 
 use zellij_utils::zellij_tile;
 
-use std::sync::{Arc, RwLock};
+use std::sync::{Arc, Mutex, RwLock};
 use std::thread;
 use std::{path::PathBuf, sync::mpsc};
 use wasmer::Store;
@@ -132,18 +132,22 @@ pub fn start_server(os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
         })
     });
 
-    #[cfg(any(feature = "test", test))]
-    thread::Builder::new()
-        .name("server_router".to_string())
-        .spawn({
-            let session_data = session_data.clone();
-            let os_input = os_input.clone();
-            let to_server = to_server.clone();
-            let session_state = session_state.clone();
+    let thread_handles = Arc::new(Mutex::new(Vec::new()));
 
-            move || route_thread_main(session_data, session_state, os_input, to_server)
-        })
-        .unwrap();
+    #[cfg(any(feature = "test", test))]
+    thread_handles.lock().unwrap().push(
+        thread::Builder::new()
+            .name("server_router".to_string())
+            .spawn({
+                let session_data = session_data.clone();
+                let os_input = os_input.clone();
+                let to_server = to_server.clone();
+                let session_state = session_state.clone();
+
+                move || route_thread_main(session_data, session_state, os_input, to_server)
+            })
+            .unwrap(),
+    );
     #[cfg(not(any(feature = "test", test)))]
     let _ = thread::Builder::new()
         .name("server_listener".to_string())
@@ -157,6 +161,7 @@ pub fn start_server(os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
             let session_state = session_state.clone();
             let to_server = to_server.clone();
             let socket_path = socket_path.clone();
+            let thread_handles = thread_handles.clone();
             move || {
                 drop(std::fs::remove_file(&socket_path));
                 let listener = LocalSocketListener::bind(&*socket_path).unwrap();
@@ -169,23 +174,25 @@ pub fn start_server(os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                             let session_data = session_data.clone();
                             let session_state = session_state.clone();
                             let to_server = to_server.clone();
-                            thread::Builder::new()
-                                .name("server_router".to_string())
-                                .spawn({
-                                    let session_data = session_data.clone();
-                                    let os_input = os_input.clone();
-                                    let to_server = to_server.clone();
-
-                                    move || {
-                                        route_thread_main(
-                                            session_data,
-                                            session_state,
-                                            os_input,
-                                            to_server,
-                                        )
-                                    }
-                                })
-                                .unwrap();
+                            thread_handles.lock().unwrap().push(
+                                thread::Builder::new()
+                                    .name("server_router".to_string())
+                                    .spawn({
+                                        let session_data = session_data.clone();
+                                        let os_input = os_input.clone();
+                                        let to_server = to_server.clone();
+
+                                        move || {
+                                            route_thread_main(
+                                                session_data,
+                                                session_state,
+                                                os_input,
+                                                to_server,
+                                            )
+                                        }
+                                    })
+                                    .unwrap(),
+                            );
                         }
                         Err(err) => {
                             panic!("err {:?}", err);
@@ -274,6 +281,11 @@ pub fn start_server(os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
             }
         }
     }
+    thread_handles
+        .lock()
+        .unwrap()
+        .drain(..)
+        .for_each(|h| drop(h.join()));
     #[cfg(not(any(feature = "test", test)))]
     drop(std::fs::remove_file(&socket_path));
 }
diff --git a/zellij-server/src/route.rs b/zellij-server/src/route.rs
index c8f78caae..8dcd98ffa 100644
--- a/zellij-server/src/route.rs
+++ b/zellij-server/src/route.rs
@@ -233,20 +233,20 @@ pub(crate) fn route_thread_main(
                     os_input.send_to_temp_client(ServerToClientMsg::Exit(ExitReason::Error(
                         "Cannot add new client".into(),
                     )));
-                    break;
+                } else {
+                    os_input.add_client_sender();
+                    to_server.send(instruction.into()).unwrap();
                 }
-                os_input.add_client_sender();
-                to_server.send(instruction.into()).unwrap();
             }
             ClientToServerMsg::AttachClient(_, force) => {
                 if *session_state.read().unwrap() == SessionState::Attached && !force {
                     os_input.send_to_temp_client(ServerToClientMsg::Exit(ExitReason::CannotAttach));
-                    break;
+                } else {
+                    os_input.add_client_sender();
+                    to_server.send(instruction.into()).unwrap();
                 }
-                os_input.add_client_sender();
-                to_server.send(instruction.into()).unwrap();
             }
-            ClientToServerMsg::ClientDetached => break,
+            ClientToServerMsg::ClientExited => break,
         }
     }
 }
diff --git a/zellij-utils/src/ipc.rs b/zellij-utils/src/ipc.rs
index 41feef1c0..a160b782f 100644
--- a/zellij-utils/src/ipc.rs
+++ b/zellij-utils/src/ipc.rs
@@ -59,7 +59,7 @@ pub enum ClientToServerMsg {
     NewClient(ClientAttributes, Box<CliArgs>, Box<Options>),
     AttachClient(ClientAttributes, bool),
     Action(Action),
-    ClientDetached,
+    ClientExited,
 }
 
 // Types of messages sent from the server to the client