[ty] Fix server hang after shutdown request

crates/ty_server/src/lib.rs

@@ -37,10 +37,18 @@ pub fn run_server() -> anyhow::Result<()> {
 
     let io_result = io_threads.join();
 
-    match (server_result, io_result) {
+    let result = match (server_result, io_result) {
         (Ok(()), Ok(())) => Ok(()),
         (Err(server), Err(io)) => Err(server).context(format!("IO thread error: {io}")),
         (Err(server), _) => Err(server),
         (_, Err(io)) => Err(io).context("IO thread error"),
+    };
+
+    if let Err(err) = result.as_ref() {
+        tracing::warn!("Server shut down with an error: {err}");
+    } else {
+        tracing::info!("Server shut down");
     }
+
+    result
 }

crates/ty_server/src/server.rs

@@ -1,5 +1,9 @@
 //! Scheduling, I/O, and API endpoints.
 
+use self::schedule::spawn_main_loop;
+use crate::PositionEncoding;
+use crate::session::{AllSettings, ClientSettings, Experimental, Session};
+use lsp_server::Connection;
 use lsp_types::{
     ClientCapabilities, DiagnosticOptions, DiagnosticServerCapabilities, HoverProviderCapability,
     InlayHintOptions, InlayHintServerCapabilities, MessageType, ServerCapabilities,
@@ -9,11 +13,6 @@ use lsp_types::{
 use std::num::NonZeroUsize;
 use std::panic::PanicHookInfo;
 
-use self::connection::Connection;
-use self::schedule::spawn_main_loop;
-use crate::PositionEncoding;
-use crate::session::{AllSettings, ClientSettings, Experimental, Session};
-
 mod api;
 mod connection;
 mod main_loop;
@@ -66,7 +65,7 @@ impl Server {
         // The number 32 was chosen arbitrarily. The main goal was to have enough capacity to queue
         // some responses before blocking.
         let (main_loop_sender, main_loop_receiver) = crossbeam::channel::bounded(32);
-        let client = Client::new(main_loop_sender.clone(), connection.sender());
+        let client = Client::new(main_loop_sender.clone(), connection.sender.clone());
 
         crate::logging::init_logging(
             global_settings.tracing.log_level.unwrap_or_default(),
@@ -131,26 +130,37 @@ impl Server {
 
     pub(crate) fn run(mut self) -> crate::Result<()> {
         type PanicHook = Box<dyn Fn(&PanicHookInfo<'_>) + 'static + Sync + Send>;
-        struct RestorePanicHook {
+        struct RestorePanicHook<'a> {
             hook: Option<PanicHook>,
+            client: &'a std::sync::Mutex<Option<Client>>,
         }
 
-        impl Drop for RestorePanicHook {
+        impl Drop for RestorePanicHook<'_> {
             fn drop(&mut self) {
+                if let Ok(mut client) = self.client.lock() {
+                    // Take the client and drop it
+                    let _ = client.take();
+                }
                 if let Some(hook) = self.hook.take() {
                     std::panic::set_hook(hook);
                 }
             }
         }
 
+        // It's important that this client gets dropped when exiting the main loop or
+        // joining the io_thread will wait forever because we hold on to a connection sender.
+        let client = std::sync::Mutex::new(Some(Client::new(
+            self.main_loop_sender.clone(),
+            self.connection.sender.clone(),
+        )));
+
         // unregister any previously registered panic hook
         // The hook will be restored when this function exits.
         let _ = RestorePanicHook {
             hook: Some(std::panic::take_hook()),
+            client: &client,
         };
 
-        let client = Client::new(self.main_loop_sender.clone(), self.connection.sender());
-
         // When we panic, try to notify the client.
         std::panic::set_hook(Box::new(move |panic_info| {
             use std::io::Write;
@@ -164,12 +174,13 @@ impl Server {
             let mut stderr = std::io::stderr().lock();
             writeln!(stderr, "{panic_info}\n{backtrace}").ok();
 
-            client
-                .show_message(
+            if let Ok(Some(client)) = client.lock().as_deref() {
+                client.show_message(
                     "The ty language server exited with a panic. See the logs for more details.",
                     MessageType::ERROR,
                 )
                 .ok();
+            }
         }));
 
         spawn_main_loop(move || self.main_loop())?.join()

crates/ty_server/src/server/api.rs

@@ -49,6 +49,15 @@ pub(super) fn request(req: server::Request) -> Task {
         >(
             req, BackgroundSchedule::LatencySensitive
         ),
+        lsp_types::request::Shutdown::METHOD => {
+            tracing::debug!("Received shutdown request, waiting for shutdown notification.");
+            Ok(Task::local(move |session, client| {
+                session.set_shutdown_requested(true);
+                if let Err(error) = client.respond(&req.id, Ok(())) {
+                    tracing::debug!("Failed to send shutdown response: {error}");
+                }
+            }))
+        }
 
         method => {
             tracing::warn!("Received request {method} which does not have a handler");

crates/ty_server/src/server/connection.rs

@@ -1,20 +1,12 @@
 use lsp_server as lsp;
-use lsp_types::{notification::Notification, request::Request};
 
 pub(crate) type ConnectionSender = crossbeam::channel::Sender<lsp::Message>;
-type ConnectionReceiver = crossbeam::channel::Receiver<lsp::Message>;
 
 /// A builder for `Connection` that handles LSP initialization.
 pub(crate) struct ConnectionInitializer {
     connection: lsp::Connection,
 }
 
-/// Handles inbound and outbound messages with the client.
-pub(crate) struct Connection {
-    sender: ConnectionSender,
-    receiver: ConnectionReceiver,
-}
-
 impl ConnectionInitializer {
     /// Create a new LSP server connection over stdin/stdout.
     pub(crate) fn stdio() -> (Self, lsp::IoThreads) {
@@ -40,7 +32,7 @@ impl ConnectionInitializer {
         server_capabilities: &lsp_types::ServerCapabilities,
         name: &str,
         version: &str,
-    ) -> crate::Result<Connection> {
+    ) -> crate::Result<lsp_server::Connection> {
         self.connection.initialize_finish(
             id,
             serde_json::json!({
@@ -51,76 +43,7 @@ impl ConnectionInitializer {
                 }
             }),
         )?;
-        let Self {
-            connection: lsp::Connection { sender, receiver },
-        } = self;
-        Ok(Connection { sender, receiver })
-    }
-}
-
-impl Connection {
-    /// Make a new `ClientSender` for sending messages to the client.
-    pub(super) fn sender(&self) -> ConnectionSender {
-        self.sender.clone()
-    }
 
-    pub(super) fn send(&self, msg: lsp::Message) -> crate::Result<()> {
-        self.sender.send(msg)?;
-        Ok(())
-    }
-
-    /// An iterator over incoming messages from the client.
-    pub(super) fn incoming(&self) -> &crossbeam::channel::Receiver<lsp::Message> {
-        &self.receiver
-    }
-
-    /// Check and respond to any incoming shutdown requests; returns`true` if the server should be shutdown.
-    pub(super) fn handle_shutdown(&self, message: &lsp::Message) -> crate::Result<bool> {
-        match message {
-            lsp::Message::Request(lsp::Request { id, method, .. })
-                if method == lsp_types::request::Shutdown::METHOD =>
-            {
-                self.sender
-                    .send(lsp::Response::new_ok(id.clone(), ()).into())?;
-                tracing::info!("Shutdown request received. Waiting for an exit notification...");
-
-                loop {
-                    match &self
-                        .receiver
-                        .recv_timeout(std::time::Duration::from_secs(30))?
-                    {
-                        lsp::Message::Notification(lsp::Notification { method, .. })
-                            if method == lsp_types::notification::Exit::METHOD =>
-                        {
-                            tracing::info!("Exit notification received. Server shutting down...");
-                            return Ok(true);
-                        }
-                        lsp::Message::Request(lsp::Request { id, method, .. }) => {
-                            tracing::warn!(
-                                "Server received unexpected request {method} ({id}) while waiting for exit notification",
-                            );
-                            self.sender.send(lsp::Message::Response(lsp::Response::new_err(
-                                id.clone(),
-                                lsp::ErrorCode::InvalidRequest as i32,
-                                "Server received unexpected request while waiting for exit notification".to_string(),
-                            )))?;
-                        }
-                        message => {
-                            tracing::warn!(
-                                "Server received unexpected message while waiting for exit notification: {message:?}"
-                            );
-                        }
-                    }
-                }
-            }
-            lsp::Message::Notification(lsp::Notification { method, .. })
-                if method == lsp_types::notification::Exit::METHOD =>
-            {
-                anyhow::bail!(
-                    "Server received an exit notification before a shutdown request was sent. Exiting..."
-                );
-            }
-            _ => Ok(false),
-        }
+        Ok(self.connection)
     }
 }

crates/ty_server/src/server/main_loop.rs

@@ -4,6 +4,7 @@ use crate::server::{Server, api};
 use crate::session::client::Client;
 use crossbeam::select;
 use lsp_server::Message;
+use lsp_types::notification::Notification;
 use lsp_types::{DidChangeWatchedFilesRegistrationOptions, FileSystemWatcher};
 
 pub(crate) type MainLoopSender = crossbeam::channel::Sender<Event>;
@@ -13,7 +14,7 @@ impl Server {
     pub(super) fn main_loop(&mut self) -> crate::Result<()> {
         self.initialize(&Client::new(
             self.main_loop_sender.clone(),
-            self.connection.sender(),
+            self.connection.sender.clone(),
         ));
 
         let mut scheduler = Scheduler::new(self.worker_threads);
@@ -25,19 +26,48 @@ impl Server {
 
             match next_event {
                 Event::Message(msg) => {
-                    if self.connection.handle_shutdown(&msg)? {
-                        break;
-                    }
+                    let client = Client::new(
+                        self.main_loop_sender.clone(),
+                        self.connection.sender.clone(),
+                    );
+
                     let task = match msg {
                         Message::Request(req) => {
                             self.session
                                 .request_queue_mut()
                                 .incoming_mut()
                                 .register(req.id.clone(), req.method.clone());
 
+                            if self.session.is_shutdown_requested() {
+                                tracing::warn!(
+                                    "Received request after server shutdown was requested, discarding"
+                                );
+                                client.respond_err(
+                                    req.id,
+                                    lsp_server::ResponseError {
+                                        code: lsp_server::ErrorCode::InvalidRequest as i32,
+                                        message: "Shutdown already requested.".to_owned(),
+                                        data: None,
+                                    },
+                                )?;
+                                continue;
+                            }
+
                             api::request(req)
                         }
-                        Message::Notification(notification) => api::notification(notification),
+                        Message::Notification(notification) => {
+                            if notification.method == lsp_types::notification::Exit::METHOD {
+                                tracing::debug!("Received exit notification, exiting");
+                                if !self.session.is_shutdown_requested() {
+                                    tracing::warn!(
+                                        "Received exit notification before shutdown request"
+                                    );
+                                }
+                                return Ok(());
+                            }
+
+                            api::notification(notification)
+                        }
 
                         // Handle the response from the client to a server request
                         Message::Response(response) => {
@@ -59,8 +89,6 @@ impl Server {
                         }
                     };
 
-                    let client =
-                        Client::new(self.main_loop_sender.clone(), self.connection.sender());
                     scheduler.dispatch(task, &mut self.session, client);
                 }
                 Event::Action(action) => match action {
@@ -75,7 +103,7 @@ impl Server {
                             let duration = start_time.elapsed();
                             tracing::trace!(name: "message response", method, %response.id, duration = format_args!("{:0.2?}", duration));
 
-                            self.connection.send(Message::Response(response))?;
+                            self.connection.sender.send(Message::Response(response))?;
                         } else {
                             tracing::trace!(
                                 "Ignoring response for canceled request id={}",
@@ -113,8 +141,8 @@ impl Server {
     /// Returns `Ok(None)` if the client connection is closed.
     fn next_event(&self) -> Result<Option<Event>, crossbeam::channel::RecvError> {
         let next = select!(
-            recv(self.connection.incoming()) -> msg => msg.map(Event::Message),
-            recv(self.main_loop_receiver) -> event => return Ok(event.ok()),
+            recv(self.connection.receiver) -> msg => return Ok(msg.ok().map(Event::Message)),
+            recv(self.main_loop_receiver) -> event => event,
         );
 
         next.map(Some)

crates/ty_server/src/session.rs

@@ -50,6 +50,9 @@ pub struct Session {
 
     /// Tracks the pending requests between client and server.
     request_queue: RequestQueue,
+
+    /// Has the client requested the server to shutdown.
+    shutdown_requested: bool,
 }
 
 impl Session {
@@ -86,6 +89,7 @@ impl Session {
                 client_capabilities,
             )),
             request_queue: RequestQueue::new(),
+            shutdown_requested: false,
         })
     }
 
@@ -97,6 +101,14 @@ impl Session {
         &mut self.request_queue
     }
 
+    pub(crate) fn is_shutdown_requested(&self) -> bool {
+        self.shutdown_requested
+    }
+
+    pub(crate) fn set_shutdown_requested(&mut self, requested: bool) {
+        self.shutdown_requested = requested;
+    }
+
     // TODO(dhruvmanila): Ideally, we should have a single method for `workspace_db_for_path_mut`
     // and `default_workspace_db_mut` but the borrow checker doesn't allow that.
     // https://github.com/astral-sh/ruff/pull/13041#discussion_r1726725437