[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,
@@ -8,11 +12,7 @@ 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};
+use std::sync::Arc;
 
 mod api;
 mod connection;
@@ -66,7 +66,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(),
@@ -130,47 +130,12 @@ impl Server {
     }
 
     pub(crate) fn run(mut self) -> crate::Result<()> {
-        type PanicHook = Box<dyn Fn(&PanicHookInfo<'_>) + 'static + Sync + Send>;
-        struct RestorePanicHook {
-            hook: Option<PanicHook>,
-        }
-
-        impl Drop for RestorePanicHook {
-            fn drop(&mut self) {
-                if let Some(hook) = self.hook.take() {
-                    std::panic::set_hook(hook);
-                }
-            }
-        }
-
-        // unregister any previously registered panic hook
-        // The hook will be restored when this function exits.
-        let _ = RestorePanicHook {
-            hook: Some(std::panic::take_hook()),
-        };
-
-        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;
-
-            let backtrace = std::backtrace::Backtrace::force_capture();
-            tracing::error!("{panic_info}\n{backtrace}");
-
-            // we also need to print to stderr directly for when using `$logTrace` because
-            // the message won't be sent to the client.
-            // But don't use `eprintln` because `eprintln` itself may panic if the pipe is broken.
-            let mut stderr = std::io::stderr().lock();
-            writeln!(stderr, "{panic_info}\n{backtrace}").ok();
+        let client = Client::new(
+            self.main_loop_sender.clone(),
+            self.connection.sender.clone(),
+        );
 
-            client
-                .show_message(
-                    "The ty language server exited with a panic. See the logs for more details.",
-                    MessageType::ERROR,
-                )
-                .ok();
-        }));
+        let _panic_hook = ServerPanicHookHandler::new(client);
 
         spawn_main_loop(move || self.main_loop())?.join()
     }
@@ -221,3 +186,63 @@ impl Server {
         }
     }
 }
+
+type PanicHook = Box<dyn Fn(&PanicHookInfo<'_>) + 'static + Sync + Send>;
+
+struct ServerPanicHookHandler {
+    hook: Option<PanicHook>,
+    // Hold on to the strong reference for as long as the panic hook is set.
+    _client: Arc<Client>,
+}
+
+impl ServerPanicHookHandler {
+    fn new(client: Client) -> Self {
+        let hook = std::panic::take_hook();
+        let client = Arc::new(client);
+
+        // Use a weak reference to the client because it must be dropped when exiting or the
+        // io-threads join hangs forever (because client has a reference to the connection sender).
+        let hook_client = Arc::downgrade(&client);
+
+        // When we panic, try to notify the client.
+        std::panic::set_hook(Box::new(move |panic_info| {
+            use std::io::Write;
+
+            let backtrace = std::backtrace::Backtrace::force_capture();
+            tracing::error!("{panic_info}\n{backtrace}");
+
+            // we also need to print to stderr directly for when using `$logTrace` because
+            // the message won't be sent to the client.
+            // But don't use `eprintln` because `eprintln` itself may panic if the pipe is broken.
+            let mut stderr = std::io::stderr().lock();
+            writeln!(stderr, "{panic_info}\n{backtrace}").ok();
+
+            if let Some(client) = hook_client.upgrade() {
+                client
+                    .show_message(
+                        "The ty language server exited with a panic. See the logs for more details.",
+                        MessageType::ERROR,
+                    )
+                    .ok();
+            }
+        }));
+
+        Self {
+            hook: Some(hook),
+            _client: client,
+        }
+    }
+}
+
+impl Drop for ServerPanicHookHandler {
+    fn drop(&mut self) {
+        if std::thread::panicking() {
+            // Calling `std::panic::set_hook` while panicking results in a panic.
+            return;
+        }
+
+        if let Some(hook) = self.hook.take() {
+            std::panic::set_hook(hook);
+        }
+    }
+}

crates/ty_server/src/server/api.rs

@@ -49,6 +49,7 @@ pub(super) fn request(req: server::Request) -> Task {
         >(
             req, BackgroundSchedule::LatencySensitive
         ),
+        lsp_types::request::Shutdown::METHOD => sync_request_task::<requests::ShutdownHandler>(req),
 
         method => {
             tracing::warn!("Received request {method} which does not have a handler");
@@ -62,7 +63,7 @@ pub(super) fn request(req: server::Request) -> Task {
     .unwrap_or_else(|err| {
         tracing::error!("Encountered error when routing request with ID {id}: {err}");
 
-        Task::local(move |_session, client| {
+        Task::sync(move |_session, client| {
             client.show_error_message(
                 "ty failed to handle a request from the editor. Check the logs for more details.",
             );
@@ -82,25 +83,25 @@ pub(super) fn request(req: server::Request) -> Task {
 pub(super) fn notification(notif: server::Notification) -> Task {
     match notif.method.as_str() {
         notifications::DidCloseTextDocumentHandler::METHOD => {
-            local_notification_task::<notifications::DidCloseTextDocumentHandler>(notif)
+            sync_notification_task::<notifications::DidCloseTextDocumentHandler>(notif)
         }
         notifications::DidOpenTextDocumentHandler::METHOD => {
-            local_notification_task::<notifications::DidOpenTextDocumentHandler>(notif)
+            sync_notification_task::<notifications::DidOpenTextDocumentHandler>(notif)
         }
         notifications::DidChangeTextDocumentHandler::METHOD => {
-            local_notification_task::<notifications::DidChangeTextDocumentHandler>(notif)
+            sync_notification_task::<notifications::DidChangeTextDocumentHandler>(notif)
         }
         notifications::DidOpenNotebookHandler::METHOD => {
-            local_notification_task::<notifications::DidOpenNotebookHandler>(notif)
+            sync_notification_task::<notifications::DidOpenNotebookHandler>(notif)
         }
         notifications::DidCloseNotebookHandler::METHOD => {
-            local_notification_task::<notifications::DidCloseNotebookHandler>(notif)
+            sync_notification_task::<notifications::DidCloseNotebookHandler>(notif)
         }
         notifications::DidChangeWatchedFiles::METHOD => {
-            local_notification_task::<notifications::DidChangeWatchedFiles>(notif)
+            sync_notification_task::<notifications::DidChangeWatchedFiles>(notif)
         }
         lsp_types::notification::Cancel::METHOD => {
-            local_notification_task::<notifications::CancelNotificationHandler>(notif)
+            sync_notification_task::<notifications::CancelNotificationHandler>(notif)
         }
         lsp_types::notification::SetTrace::METHOD => {
             tracing::trace!("Ignoring `setTrace` notification");
@@ -114,20 +115,20 @@ pub(super) fn notification(notif: server::Notification) -> Task {
     }
         .unwrap_or_else(|err| {
             tracing::error!("Encountered error when routing notification: {err}");
-            Task::local(|_session, client| {
+            Task::sync(|_session, client| {
                 client.show_error_message(
                     "ty failed to handle a notification from the editor. Check the logs for more details."
                 );
             })
         })
 }
 
-fn _local_request_task<R: traits::SyncRequestHandler>(req: server::Request) -> Result<Task>
+fn sync_request_task<R: traits::SyncRequestHandler>(req: server::Request) -> Result<Task>
 where
     <<R as RequestHandler>::RequestType as Request>::Params: UnwindSafe,
 {
     let (id, params) = cast_request::<R>(req)?;
-    Ok(Task::local(move |session, client: &Client| {
+    Ok(Task::sync(move |session, client: &Client| {
         let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
         let result = R::run(session, client, params);
         respond::<R>(&id, result, client);
@@ -245,11 +246,11 @@ where
     }
 }
 
-fn local_notification_task<N: traits::SyncNotificationHandler>(
+fn sync_notification_task<N: traits::SyncNotificationHandler>(
     notif: server::Notification,
 ) -> Result<Task> {
     let (id, params) = cast_notification::<N>(notif)?;
-    Ok(Task::local(move |session, client| {
+    Ok(Task::sync(move |session, client| {
         let _span = tracing::debug_span!("notification", method = N::METHOD).entered();
         if let Err(err) = N::run(session, client, params) {
             tracing::error!("An error occurred while running {id}: {err}");

crates/ty_server/src/server/api/requests.rs

@@ -3,9 +3,11 @@ mod diagnostic;
 mod goto_type_definition;
 mod hover;
 mod inlay_hints;
+mod shutdown;
 
 pub(super) use completion::CompletionRequestHandler;
 pub(super) use diagnostic::DocumentDiagnosticRequestHandler;
 pub(super) use goto_type_definition::GotoTypeDefinitionRequestHandler;
 pub(super) use hover::HoverRequestHandler;
 pub(super) use inlay_hints::InlayHintRequestHandler;
+pub(super) use shutdown::ShutdownHandler;

crates/ty_server/src/server/api/requests/shutdown.rs

@@ -0,0 +1,17 @@
+use crate::Session;
+use crate::server::api::traits::{RequestHandler, SyncRequestHandler};
+use crate::session::client::Client;
+
+pub(crate) struct ShutdownHandler;
+
+impl RequestHandler for ShutdownHandler {
+    type RequestType = lsp_types::request::Shutdown;
+}
+
+impl SyncRequestHandler for ShutdownHandler {
+    fn run(session: &mut Session, _client: &Client, _params: ()) -> crate::server::Result<()> {
+        tracing::debug!("Received shutdown request, waiting for shutdown notification");
+        session.set_shutdown_requested(true);
+        Ok(())
+    }
+}

crates/ty_server/src/server/api/traits.rs

@@ -17,7 +17,6 @@ pub(super) trait RequestHandler {
 /// This will block the main message receiver loop, meaning that no
 /// incoming requests or notifications will be handled while `run` is
 /// executing. Try to avoid doing any I/O or long-running computations.
-#[expect(dead_code)]
 pub(super) trait SyncRequestHandler: RequestHandler {
     fn run(
         session: &mut Session,