diff --git a/crates/ty_server/src/lib.rs b/crates/ty_server/src/lib.rs index 0cf68dfdf988cd..5589c60464c380 100644 --- a/crates/ty_server/src/lib.rs +++ b/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 } diff --git a/crates/ty_server/src/server.rs b/crates/ty_server/src/server.rs index 3ceb5402021eaf..84093d6ffa65b0 100644 --- a/crates/ty_server/src/server.rs +++ b/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) + 'static + Sync + Send>; - struct RestorePanicHook { - hook: Option, - } - - 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) + 'static + Sync + Send>; + +struct ServerPanicHookHandler { + hook: Option, + // Hold on to the strong reference for as long as the panic hook is set. + _client: Arc, +} + +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); + } + } +} diff --git a/crates/ty_server/src/server/api.rs b/crates/ty_server/src/server/api.rs index d6a934c9489fbf..db0c1bb1a729fd 100644 --- a/crates/ty_server/src/server/api.rs +++ b/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::(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::(notif) + sync_notification_task::(notif) } notifications::DidOpenTextDocumentHandler::METHOD => { - local_notification_task::(notif) + sync_notification_task::(notif) } notifications::DidChangeTextDocumentHandler::METHOD => { - local_notification_task::(notif) + sync_notification_task::(notif) } notifications::DidOpenNotebookHandler::METHOD => { - local_notification_task::(notif) + sync_notification_task::(notif) } notifications::DidCloseNotebookHandler::METHOD => { - local_notification_task::(notif) + sync_notification_task::(notif) } notifications::DidChangeWatchedFiles::METHOD => { - local_notification_task::(notif) + sync_notification_task::(notif) } lsp_types::notification::Cancel::METHOD => { - local_notification_task::(notif) + sync_notification_task::(notif) } lsp_types::notification::SetTrace::METHOD => { tracing::trace!("Ignoring `setTrace` notification"); @@ -114,7 +115,7 @@ 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." ); @@ -122,12 +123,12 @@ pub(super) fn notification(notif: server::Notification) -> Task { }) } -fn _local_request_task(req: server::Request) -> Result +fn sync_request_task(req: server::Request) -> Result where <::RequestType as Request>::Params: UnwindSafe, { let (id, params) = cast_request::(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::(&id, result, client); @@ -245,11 +246,11 @@ where } } -fn local_notification_task( +fn sync_notification_task( notif: server::Notification, ) -> Result { let (id, params) = cast_notification::(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}"); diff --git a/crates/ty_server/src/server/api/requests.rs b/crates/ty_server/src/server/api/requests.rs index bfdec09623f7f8..12e7243893f777 100644 --- a/crates/ty_server/src/server/api/requests.rs +++ b/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; diff --git a/crates/ty_server/src/server/api/requests/shutdown.rs b/crates/ty_server/src/server/api/requests/shutdown.rs new file mode 100644 index 00000000000000..73b956168c7ecd --- /dev/null +++ b/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(()) + } +} diff --git a/crates/ty_server/src/server/api/traits.rs b/crates/ty_server/src/server/api/traits.rs index 7132251bf07af0..460da654dbffd8 100644 --- a/crates/ty_server/src/server/api/traits.rs +++ b/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, diff --git a/crates/ty_server/src/server/connection.rs b/crates/ty_server/src/server/connection.rs index ded1fcee7b257b..f347cf8eb47a61 100644 --- a/crates/ty_server/src/server/connection.rs +++ b/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; -type ConnectionReceiver = crossbeam::channel::Receiver; /// 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 { + ) -> crate::Result { 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 { - &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 { - 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) } } diff --git a/crates/ty_server/src/server/main_loop.rs b/crates/ty_server/src/server/main_loop.rs index 9f0aefaf5b41d0..456e9d88b21354 100644 --- a/crates/ty_server/src/server/main_loop.rs +++ b/crates/ty_server/src/server/main_loop.rs @@ -2,8 +2,10 @@ use crate::Session; use crate::server::schedule::Scheduler; use crate::server::{Server, api}; use crate::session::client::Client; +use anyhow::anyhow; use crossbeam::select; use lsp_server::Message; +use lsp_types::notification::Notification; use lsp_types::{DidChangeWatchedFilesRegistrationOptions, FileSystemWatcher}; pub(crate) type MainLoopSender = crossbeam::channel::Sender; @@ -13,7 +15,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,9 +27,11 @@ 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 @@ -35,9 +39,37 @@ impl Server { .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 { + if !self.session.is_shutdown_requested() { + return Err(anyhow!( + "Received exit notification before a shutdown request" + )); + } + + tracing::debug!("Received exit notification, exiting"); + return Ok(()); + } + + api::notification(notification) + } // Handle the response from the client to a server request Message::Response(response) => { @@ -59,8 +91,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 +105,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={}", @@ -112,12 +142,13 @@ impl Server { /// /// Returns `Ok(None)` if the client connection is closed. fn next_event(&self) -> Result, 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()), - ); - - next.map(Some) + select!( + recv(self.connection.receiver) -> msg => { + // Ignore disconnect errors, they're handled by the main loop (it will exit). + Ok(msg.ok().map(Event::Message)) + }, + recv(self.main_loop_receiver) -> event => event.map(Some), + ) } fn initialize(&mut self, client: &Client) { diff --git a/crates/ty_server/src/server/schedule/task.rs b/crates/ty_server/src/server/schedule/task.rs index e781ffcec700a0..024a76fe13121b 100644 --- a/crates/ty_server/src/server/schedule/task.rs +++ b/crates/ty_server/src/server/schedule/task.rs @@ -68,7 +68,7 @@ impl Task { }) } /// Creates a new local task. - pub(crate) fn local(func: F) -> Self + pub(crate) fn sync(func: F) -> Self where F: FnOnce(&mut Session, &Client) + 'static, { @@ -82,7 +82,7 @@ impl Task { where R: Serialize + Send + 'static, { - Self::local(move |_, client| { + Self::sync(move |_, client| { if let Err(err) = client.respond(&id, result) { tracing::error!("Unable to send immediate response: {err}"); } @@ -91,6 +91,6 @@ impl Task { /// Creates a local task that does nothing. pub(crate) fn nothing() -> Self { - Self::local(move |_, _| {}) + Self::sync(move |_, _| {}) } } diff --git a/crates/ty_server/src/session.rs b/crates/ty_server/src/session.rs index f995b2f19f69f0..cced715ad56f1e 100644 --- a/crates/ty_server/src/session.rs +++ b/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