fix(server): prioritize shutdown requests (#2624)

refs #2612

5 files changed, 247 insertions, 37 deletions

diff --git a/src/api/mod.rs b/src/api/mod.rs
index a44899da..b4a5d262 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -7,7 +7,8 @@ mod subscriptions;
 mod wait;
 
 pub use event_hub::EventHub;
-pub use server::{start_server, start_server_with_capabilities, ServerHandle};
+pub(crate) use server::start_server_with_stop_control;
+pub use server::{start_server_with_capabilities, ServerHandle};
 pub use status::{read_runtime_status_at, RuntimeStatus};
 
 use std::path::PathBuf;
diff --git a/src/api/server.rs b/src/api/server.rs
index e8b73167..3af8051f 100644
--- a/src/api/server.rs
+++ b/src/api/server.rs
@@ -56,18 +56,12 @@ impl ServerHandle {
     }
 }
 
-pub fn start_server(
+pub(crate) fn start_server_with_stop_control(
     api_tx: ApiRequestSender,
     event_hub: EventHub,
+    server_stop: Arc<AtomicBool>,
 ) -> std::io::Result<ServerHandle> {
-    start_server_with_capabilities(
-        api_tx,
-        event_hub,
-        Some(ServerCapabilities {
-            live_handoff: crate::platform::capabilities().live_handoff,
-            detached_server_daemon: crate::platform::current_process_is_detached_server_daemon(),
-        }),
-    )
+    start_server_inner(api_tx, event_hub, default_capabilities(), Some(server_stop))
 }
 
 pub fn start_server_with_capabilities(
@@ -75,6 +69,22 @@ pub fn start_server_with_capabilities(
     event_hub: EventHub,
     capabilities: Option<ServerCapabilities>,
 ) -> std::io::Result<ServerHandle> {
+    start_server_inner(api_tx, event_hub, capabilities, None)
+}
+
+fn default_capabilities() -> Option<ServerCapabilities> {
+    Some(ServerCapabilities {
+        live_handoff: crate::platform::capabilities().live_handoff,
+        detached_server_daemon: crate::platform::current_process_is_detached_server_daemon(),
+    })
+}
+
+fn start_server_inner(
+    api_tx: ApiRequestSender,
+    event_hub: EventHub,
+    capabilities: Option<ServerCapabilities>,
+    server_stop: Option<Arc<AtomicBool>>,
+) -> std::io::Result<ServerHandle> {
     let path = socket_path();
     prepare_socket_path(&path)?;
 
@@ -92,14 +102,16 @@ pub fn start_server_with_capabilities(
                     let api_tx = api_tx.clone();
                     let event_hub = event_hub.clone();
                     let capabilities = capabilities.clone();
+                    let server_stop = server_stop.clone();
                     let connection_running = Arc::clone(&listener_running);
                     std::thread::spawn(move || {
-                        if let Err(err) = handle_connection(
+                        if let Err(err) = handle_connection_with_stop(
                             stream,
                             &api_tx,
                             &event_hub,
                             &connection_running,
                             capabilities,
+                            server_stop.as_ref(),
                         ) {
                             warn!(err = %err, "api connection failed");
                         }
@@ -135,12 +147,24 @@ fn restrict_socket_permissions(path: &Path) -> std::io::Result<()> {
     crate::ipc::restrict_socket_permissions(path, SOCKET_PERMISSION_MODE)
 }
 
+#[cfg(test)]
 fn handle_connection(
+    stream: LocalStream,
+    api_tx: &ApiRequestSender,
+    event_hub: &EventHub,
+    running: &Arc<AtomicBool>,
+    capabilities: Option<ServerCapabilities>,
+) -> std::io::Result<()> {
+    handle_connection_with_stop(stream, api_tx, event_hub, running, capabilities, None)
+}
+
+fn handle_connection_with_stop(
     mut stream: LocalStream,
     api_tx: &ApiRequestSender,
     event_hub: &EventHub,
     running: &Arc<AtomicBool>,
     capabilities: Option<ServerCapabilities>,
+    server_stop: Option<&Arc<AtomicBool>>,
 ) -> std::io::Result<()> {
     if let Err(err) = stream.set_send_timeout(Some(STREAM_WRITE_TIMEOUT)) {
         debug!(err = %err, "api connection write timeout unavailable");
@@ -263,6 +287,7 @@ fn handle_connection(
                 },
                 api_tx,
                 capabilities,
+                server_stop,
                 Some(response_write_rx),
             );
             let result = write_text_line_allow_disconnect(&mut stream, &response);
@@ -316,10 +341,11 @@ fn handle_request(
     request: Request,
     api_tx: &ApiRequestSender,
     capabilities: Option<ServerCapabilities>,
+    server_stop: Option<&Arc<AtomicBool>>,
     response_write_complete: Option<std::sync::mpsc::Receiver<()>>,
 ) -> String {
-    match request.method {
-        Method::Ping(_) => serde_json::to_string(&SuccessResponse {
+    if matches!(&request.method, Method::Ping(_)) {
+        return serde_json::to_string(&SuccessResponse {
             id: request.id,
             result: ResponseResult::Pong {
                 version: crate::build_info::version(),
@@ -330,9 +356,27 @@ fn handle_request(
         .unwrap_or_else(|_| {
             r#"{"id":"","error":{"code":"internal_error","message":"failed to encode response"}}"#
                 .to_string()
-        }),
-        _ => dispatch_to_app(request, api_tx, None, response_write_complete, None),
+        });
     }
+
+    if matches!(&request.method, Method::ServerStop(_)) {
+        if let Some(server_stop) = server_stop {
+            server_stop.store(true, Ordering::Release);
+            return serde_json::to_string(&SuccessResponse {
+                id: request.id,
+                result: ResponseResult::Ok {},
+            })
+            .unwrap_or_else(|_| "{}".to_string());
+        }
+    } else if server_stop.is_some_and(|stop| stop.load(Ordering::Acquire)) {
+        return error_response_json(
+            request.id,
+            "server_unavailable",
+            "server is shutting down".into(),
+        );
+    }
+
+    dispatch_to_app(request, api_tx, None, response_write_complete, None)
 }
 
 fn api_method_name(method: &Method) -> &'static str {
@@ -1039,6 +1083,7 @@ mod tests {
                 detached_server_daemon: true,
             }),
             None,
+            None,
         );
 
         let parsed: SuccessResponse = serde_json::from_str(&response).unwrap();
@@ -1047,6 +1092,41 @@ mod tests {
     }
 
     #[test]
+    fn server_stop_control_bypasses_app_channel() {
+        let (tx, mut rx) = mpsc::unbounded_channel();
+        let stop = Arc::new(AtomicBool::new(false));
+        let response = handle_request(
+            Request {
+                id: "priority_stop".into(),
+                method: Method::ServerStop(crate::api::schema::EmptyParams::default()),
+            },
+            &tx,
+            None,
+            Some(&stop),
+            None,
+        );
+
+        let response: serde_json::Value = serde_json::from_str(&response).unwrap();
+        assert_eq!(response["id"], "priority_stop");
+        assert_eq!(response["result"]["type"], "ok");
+        assert!(stop.load(Ordering::Acquire));
+
+        let rejected = handle_request(
+            Request {
+                id: "after_stop".into(),
+                method: Method::WorkspaceList(crate::api::schema::EmptyParams::default()),
+            },
+            &tx,
+            None,
+            Some(&stop),
+            None,
+        );
+        let rejected: serde_json::Value = serde_json::from_str(&rejected).unwrap();
+        assert_eq!(rejected["error"]["code"], "server_unavailable");
+        assert!(rx.try_recv().is_err());
+    }
+
+    #[test]
     fn request_dispatches_to_app_channel() {
         let (tx, mut rx) = mpsc::unbounded_channel();
         let request = Request {
@@ -1056,7 +1136,7 @@ mod tests {
 
         let request_for_thread = request.clone();
         let thread =
-            std::thread::spawn(move || handle_request(request_for_thread, &tx, None, None));
+            std::thread::spawn(move || handle_request(request_for_thread, &tx, None, None, None));
 
         let msg = rx.blocking_recv().unwrap();
         assert_eq!(msg.request.id, "req_2");
diff --git a/src/server/client_accept.rs b/src/server/client_accept.rs
index ccfe7b1a..e2a8a42d 100644
--- a/src/server/client_accept.rs
+++ b/src/server/client_accept.rs
@@ -1,5 +1,5 @@
 use std::io;
-use std::sync::{atomic::AtomicBool, Arc};
+use std::sync::{atomic::AtomicBool, atomic::Ordering, Arc};
 
 use interprocess::local_socket::traits::{Listener as _, Stream as _};
 use tokio::sync::mpsc;
@@ -16,6 +16,9 @@ pub(crate) fn accept_pending_client_connections(
     server_event_tx: &mpsc::Sender<ServerEvent>,
 ) -> io::Result<()> {
     loop {
+        if should_quit.load(Ordering::Acquire) {
+            break;
+        }
         match listener.accept() {
             Ok(stream) => {
                 let client_id = *next_client_id;
diff --git a/src/server/client_transport.rs b/src/server/client_transport.rs
index b5023b97..f1237c89 100644
--- a/src/server/client_transport.rs
+++ b/src/server/client_transport.rs
@@ -522,6 +522,10 @@ pub(crate) fn handle_client_handshake(
     server_event_tx: &mpsc::Sender<ServerEvent>,
     should_quit: &Arc<AtomicBool>,
 ) -> io::Result<()> {
+    if should_quit.load(Ordering::Acquire) {
+        return Ok(());
+    }
+
     // Reset to blocking mode — the accept loop sets nonblocking but
     // the handshake thread needs blocking I/O for read_message/write_message.
     stream.set_nonblocking(false)?;
@@ -624,6 +628,10 @@ pub(crate) fn handle_client_handshake(
         }
     };
 
+    if should_quit.load(Ordering::Acquire) {
+        return Ok(());
+    }
+
     // Send Welcome.
     let welcome = ServerMessage::Welcome {
         version: PROTOCOL_VERSION,
@@ -653,8 +661,13 @@ pub(crate) fn handle_client_handshake(
         client_writer_loop(write_stream, client_id, writer_queue, writer_event_tx);
     });
 
+    if should_quit.load(Ordering::Acquire) {
+        send_shutdown_to_unregistered_client(&writer);
+        return Ok(());
+    }
+
     // Notify the main loop about the new client.
-    let _ = server_event_tx.blocking_send(ServerEvent::ClientConnected {
+    let connected = ServerEvent::ClientConnected {
         client_id,
         cols: client_cols,
         rows: client_rows,
@@ -665,12 +678,31 @@ pub(crate) fn handle_client_handshake(
         direct_attach_requested,
         direct_graphics,
         writer,
-    });
+    };
+    if let Err(err) = server_event_tx.blocking_send(connected) {
+        if let ServerEvent::ClientConnected { writer, .. } = err.0 {
+            send_shutdown_to_unregistered_client(&writer);
+        }
+    }
 
     // Enter read loop — read client messages and forward to main loop.
     client_read_loop(stream, client_id, server_event_tx, should_quit)
 }
 
+fn send_shutdown_to_unregistered_client(writer: &ClientWriter) {
+    let mut framed = Vec::new();
+    if protocol::write_message(
+        &mut framed,
+        &ServerMessage::ServerShutdown {
+            reason: Some("server is shutting down".to_owned()),
+        },
+    )
+    .is_ok()
+    {
+        let _ = writer.control.send(framed);
+    }
+}
+
 /// The client writer loop — prioritizes control messages over render frames.
 fn client_writer_loop(
     mut stream: LocalStream,
diff --git a/src/server/headless.rs b/src/server/headless.rs
index 6461b13e..0632861c 100644
--- a/src/server/headless.rs
+++ b/src/server/headless.rs
@@ -466,6 +466,7 @@ impl HeadlessServer {
         config_diagnostics: &[String],
         api_tx: Option<api::ApiRequestSender>,
         api_server: Option<api::ServerHandle>,
+        should_quit: Arc<AtomicBool>,
     ) -> io::Result<Self> {
         let client_path = client_socket_path();
         prepare_socket_path(&client_path)?;
@@ -479,8 +480,6 @@ impl HeadlessServer {
         #[cfg(unix)]
         listener.set_nonblocking(ListenerNonblockingMode::Accept)?;
 
-        let should_quit = Arc::new(AtomicBool::new(false));
-
         // Channel for server events from client threads.
         let (server_event_tx, server_event_rx) = mpsc::channel(64);
         #[cfg(windows)]
@@ -555,12 +554,15 @@ impl HeadlessServer {
 
             // If shutdown has been initiated, complete it and exit.
             if self.shutting_down {
-                self.complete_shutdown()?;
+                self.complete_shutdown().await?;
                 break;
             }
 
             // Check if we should start shutting down.
             if self.app.state.should_quit || self.should_quit.load(Ordering::Acquire) {
+                self.drain_internal_events_with_forwarding_up_to(
+                    crate::app::APP_EVENT_CHANNEL_CAPACITY,
+                );
                 self.initiate_shutdown();
                 continue;
             }
@@ -585,6 +587,9 @@ impl HeadlessServer {
                 needs_graphics_render = false;
                 crate::render_prof::event("full_render_cause.internal_events");
             }
+            if self.should_quit.load(Ordering::Acquire) {
+                continue;
+            }
             if self.app.expire_due_metadata(Instant::now()) {
                 needs_render = true;
                 needs_full_render = true;
@@ -612,6 +617,9 @@ impl HeadlessServer {
                 needs_full_render = true;
                 crate::render_prof::event("full_render_cause.api_requests");
             }
+            if self.should_quit.load(Ordering::Acquire) {
+                continue;
+            }
 
             self.app.sync_focus_events();
             self.app.sync_session_save_schedule();
@@ -640,6 +648,9 @@ impl HeadlessServer {
                 needs_full_render = true;
                 crate::render_prof::event("full_render_cause.server_events");
             }
+            if self.should_quit.load(Ordering::Acquire) {
+                continue;
+            }
 
             // 6. Handle scheduled tasks.
             let now = Instant::now();
@@ -786,6 +797,25 @@ impl HeadlessServer {
                 }
             };
 
+            if self.should_quit.load(Ordering::Acquire) {
+                match event {
+                    LoopEvent::Internal(ev) => {
+                        self.handle_internal_event_with_forwarding(ev);
+                    }
+                    LoopEvent::ServerEvent(ServerEvent::ClientConnected { writer, .. }) => {
+                        if let Ok(message) =
+                            Self::frame_server_message(&ServerMessage::ServerShutdown {
+                                reason: Some("server is shutting down".to_owned()),
+                            })
+                        {
+                            let _ = writer.control.send(message);
+                        }
+                    }
+                    _ => {}
+                }
+                continue;
+            }
+
             match event {
                 LoopEvent::Timer => {}
                 LoopEvent::Internal(ev) => {
@@ -1394,7 +1424,11 @@ impl HeadlessServer {
             .api_tx
             .clone()
             .ok_or_else(|| io::Error::other("cannot restore api socket without api sender"))?;
-        let api_server = api::start_server(api_tx, self.app.event_hub.clone())?;
+        let api_server = api::start_server_with_stop_control(
+            api_tx,
+            self.app.event_hub.clone(),
+            self.should_quit.clone(),
+        )?;
 
         let client_path = client_socket_path();
         prepare_socket_path(&client_path)?;
@@ -1657,7 +1691,10 @@ impl HeadlessServer {
     /// Uses the original full-render semantics when pane graphics are dormant.
     fn drain_server_events(&mut self) -> bool {
         let mut changed = false;
-        while let Ok(ev) = self.server_event_rx.try_recv() {
+        while !self.should_quit.load(Ordering::Acquire) {
+            let Ok(ev) = self.server_event_rx.try_recv() else {
+                break;
+            };
             changed |= self.handle_server_event(ev);
         }
         changed
@@ -1666,12 +1703,28 @@ impl HeadlessServer {
     /// Returns the strongest render impact from the drained event batch.
     fn drain_server_events_with_render_impact(&mut self) -> RenderImpact {
         let mut impact = RenderImpact::None;
-        while let Ok(ev) = self.server_event_rx.try_recv() {
+        while !self.should_quit.load(Ordering::Acquire) {
+            let Ok(ev) = self.server_event_rx.try_recv() else {
+                break;
+            };
             impact.merge(self.handle_server_event_with_render_impact(ev));
         }
         impact
     }
 
+    async fn reject_late_client_connections(&mut self) {
+        self.server_event_rx.close();
+        while let Some(event) = self.server_event_rx.recv().await {
+            if let ServerEvent::ClientConnected { writer, .. } = event {
+                if let Ok(message) = Self::frame_server_message(&ServerMessage::ServerShutdown {
+                    reason: Some("server is shutting down".to_owned()),
+                }) {
+                    let _ = writer.control.send(message);
+                }
+            }
+        }
+    }
+
     fn terminal_id_by_string(&self, terminal_id: &str) -> Option<crate::terminal::TerminalId> {
         self.app
             .state
@@ -2431,7 +2484,7 @@ impl HeadlessServer {
             let (had_event, batch_changed) =
                 self.drain_internal_events_with_forwarding_up_to(crate::app::APP_EVENT_DRAIN_LIMIT);
             changed |= batch_changed;
-            if !had_event {
+            if !had_event || self.should_quit.load(Ordering::Acquire) {
                 break;
             }
         }
@@ -3335,15 +3388,30 @@ impl HeadlessServer {
     /// During shutdown, remaining requests get a `server_unavailable` error.
     fn drain_api_requests_with_shutdown_check(&mut self) -> bool {
         let mut changed = false;
-        while let Ok(msg) = self.app.api_rx.try_recv() {
+        while !self.should_quit.load(Ordering::Acquire) {
+            let Ok(msg) = self.app.api_rx.try_recv() else {
+                break;
+            };
             changed |= self.handle_api_request_with_shutdown_check(msg);
         }
         changed
     }
 
+    fn reject_queued_api_requests_for_shutdown(&mut self) {
+        for _ in 0..self.app.api_rx.len() {
+            let Ok(msg) = self.app.api_rx.try_recv() else {
+                break;
+            };
+            self.handle_api_request_with_shutdown_check(msg);
+        }
+    }
+
     fn drain_api_requests_with_render_impact(&mut self) -> RenderImpact {
         let mut impact = RenderImpact::None;
-        while let Ok(msg) = self.app.api_rx.try_recv() {
+        while !self.should_quit.load(Ordering::Acquire) {
+            let Ok(msg) = self.app.api_rx.try_recv() else {
+                break;
+            };
             impact.merge(self.handle_api_request_with_render_impact(msg));
         }
         impact
@@ -4594,8 +4662,9 @@ impl HeadlessServer {
 
     /// Completes the shutdown sequence: send ServerShutdown to clients,
     /// close client connections, remove socket files, and clean up.
-    fn complete_shutdown(&mut self) -> io::Result<()> {
+    async fn complete_shutdown(&mut self) -> io::Result<()> {
         info!("completing server shutdown");
+        self.reject_late_client_connections().await;
 
         // Send ServerShutdown to all remaining clients.
         if !self.clients.is_empty() {
@@ -4609,8 +4678,8 @@ impl HeadlessServer {
             std::thread::sleep(Duration::from_millis(50));
         }
 
-        // Drain remaining API requests with server_unavailable.
-        self.drain_api_requests_with_shutdown_check();
+        // Reject only the requests already queued when shutdown reached cleanup.
+        self.reject_queued_api_requests_for_shutdown();
 
         // Close all client connections.
         let staged_files = self
@@ -4810,9 +4879,14 @@ pub fn run_server() -> io::Result<()> {
     let loaded_config = config::Config::load();
     let (api_tx, api_rx) = tokio::sync::mpsc::unbounded_channel();
     let event_hub = api::EventHub::default();
+    let should_quit = Arc::new(AtomicBool::new(false));
 
     // Start the JSON API socket server.
-    let _api_server = match api::start_server(api_tx.clone(), event_hub.clone()) {
+    let _api_server = match api::start_server_with_stop_control(
+        api_tx.clone(),
+        event_hub.clone(),
+        should_quit.clone(),
+    ) {
         Ok(server) => server,
         Err(err) if err.kind() == io::ErrorKind::AddrInUse => {
             eprintln!("error: herdr server is already running");
@@ -4855,6 +4929,7 @@ pub fn run_server() -> io::Result<()> {
             &loaded_config.diagnostics,
             Some(api_tx.clone()),
             Some(_api_server),
+            should_quit,
         ) {
             Ok(server) => server,
             Err(err) if err.kind() == io::ErrorKind::AddrInUse => {
@@ -4919,6 +4994,7 @@ fn run_handoff_import_server(socket_path: &Path, token: &str) -> io::Result<()>
 
     let (api_tx, api_rx) = tokio::sync::mpsc::unbounded_channel();
     let event_hub = api::EventHub::default();
+    let should_quit = Arc::new(AtomicBool::new(false));
 
     let mut imports = HashMap::new();
     for (pane, fd) in received.manifest.panes.into_iter().zip(received.fds) {
@@ -4957,12 +5033,17 @@ fn run_handoff_import_server(socket_path: &Path, token: &str) -> io::Result<()>
         }
         wait_for_old_public_sockets_to_close(Duration::from_secs(5))?;
 
-        let api_server = api::start_server(api_tx.clone(), event_hub.clone())?;
+        let api_server = api::start_server_with_stop_control(
+            api_tx.clone(),
+            event_hub.clone(),
+            should_quit.clone(),
+        )?;
         let mut server = HeadlessServer::new(
             app,
             &loaded_config.diagnostics,
             Some(api_tx.clone()),
             Some(api_server),
+            should_quit,
         )?;
         crate::server::handoff::report_ready(&mut received.stream)?;
         crate::server::handoff::wait_committed(&mut received.stream)?;
@@ -5108,7 +5189,6 @@ mod tests {
             .set_nonblocking(ListenerNonblockingMode::Accept)
             .expect("set listener nonblocking");
         let (server_event_tx, server_event_rx) = mpsc::channel(64);
-        #[cfg(windows)]
         let should_quit = Arc::new(AtomicBool::new(false));
         #[cfg(windows)]
         spawn_windows_client_accept_thread(listener, should_quit.clone(), server_event_tx.clone());
@@ -5139,9 +5219,6 @@ mod tests {
             handoff_in_progress: false,
             #[cfg(unix)]
             pending_handoff_repaint_nudge: false,
-            #[cfg(unix)]
-            should_quit: Arc::new(AtomicBool::new(false)),
-            #[cfg(windows)]
             should_quit,
             server_event_rx,
             server_event_tx,
@@ -5276,6 +5353,23 @@ mod tests {
     }
 
     #[test]
+    fn server_stop_interrupts_server_event_backlog() {
+        let mut server = test_headless_server();
+        for client_id in 1..=64 {
+            server
+                .server_event_tx
+                .try_send(ServerEvent::ClientDisconnected { client_id })
+                .unwrap();
+        }
+
+        server.should_quit.store(true, Ordering::Release);
+
+        assert!(!server.drain_server_events());
+        assert!(server.server_event_rx.try_recv().is_ok());
+        shutdown_test_runtimes(&mut server);
+    }
+
+    #[test]
     fn headless_api_request_drains_all_pending_internal_events_before_reading_state() {
         let mut server = test_headless_server();
         for i in 0..=crate::app::APP_EVENT_DRAIN_LIMIT {