diff --git a/crates/squawk_server/src/dispatch.rs b/crates/squawk_server/src/dispatch.rs index 97fc1ef6..d8f873ef 100644 --- a/crates/squawk_server/src/dispatch.rs +++ b/crates/squawk_server/src/dispatch.rs @@ -4,11 +4,16 @@ use std::panic::UnwindSafe; use anyhow::Result; use log::{error, info}; -use lsp_server::{RequestId, Response}; +use lsp_server::{Request, Response}; use lsp_types::{notification::Notification as LspNotification, request::Request as LspRequest}; +use salsa::Cancelled; +use serde::{Serialize, de::DeserializeOwned}; use squawk_thread::ThreadIntent; -use crate::global_state::{GlobalState, Snapshot}; +use crate::{ + global_state::{GlobalState, Snapshot, TaskResult}, + panic::PanicError, +}; pub(crate) struct RequestDispatcher<'a> { req: Option, @@ -23,15 +28,14 @@ impl<'a> RequestDispatcher<'a> { } } - fn parse(&mut self) -> Option<(lsp_server::RequestId, R::Params)> + fn parse(&mut self) -> Option<(Request, R::Params)> where R: LspRequest, { let req = self.req.take_if(|req| req.method.as_str() == R::METHOD)?; let id = req.id.clone(); - - match req.extract(R::METHOD) { - Ok((id, params)) => Some((id, params)), + match from_json(R::METHOD, &req.params) { + Ok(params) => Some((req, params)), Err(err) => { let response = Response::new_err( id, @@ -44,6 +48,26 @@ impl<'a> RequestDispatcher<'a> { } } + pub(crate) fn on_sync( + mut self, + handler: fn(&Snapshot, R::Params) -> Result, + ) -> Self + where + R: LspRequest, + R::Params: Send + 'static + UnwindSafe, + { + let Some((request, params)) = self.parse::() else { + return self; + }; + + let snapshot = self.global_state.snapshot(); + let result = crate::panic::catch_unwind(|| handler(&snapshot, params)); + if let Ok(response) = thread_result_to_response::(request.id.clone(), result) { + self.global_state.respond(response); + } + self + } + pub(crate) fn on_sync_mut( mut self, handler: fn(&mut GlobalState, R::Params) -> Result, @@ -52,23 +76,29 @@ impl<'a> RequestDispatcher<'a> { R: LspRequest, R::Params: Send + 'static, { - if let Some((id, params)) = self.parse::() { - let result = handler(self.global_state, params); - let response = result_to_response::(id, result); + let Some((request, params)) = self.parse::() else { + return self; + }; + + let result = handler(self.global_state, params); + if let Ok(response) = result_to_response::(request.id.clone(), result) { self.global_state.respond(response); } self } - pub(crate) fn on(self, handler: fn(&Snapshot, R::Params) -> Result) -> Self + pub(crate) fn on( + self, + handler: fn(&Snapshot, R::Params) -> Result, + ) -> Self where R: LspRequest, R::Params: Send + 'static + UnwindSafe, { - self.on_with_thread_intent::(ThreadIntent::Worker, handler) + self.on_with_thread_intent::(ThreadIntent::Worker, handler) } - pub(crate) fn on_latency_sensitive( + pub(crate) fn on_latency_sensitive( self, handler: fn(&Snapshot, R::Params) -> Result, ) -> Self @@ -76,10 +106,10 @@ impl<'a> RequestDispatcher<'a> { R: LspRequest, R::Params: Send + 'static + UnwindSafe, { - self.on_with_thread_intent::(ThreadIntent::LatencySensitive, handler) + self.on_with_thread_intent::(ThreadIntent::LatencySensitive, handler) } - fn on_with_thread_intent( + fn on_with_thread_intent( mut self, intent: ThreadIntent, handler: fn(&Snapshot, R::Params) -> Result, @@ -88,15 +118,20 @@ impl<'a> RequestDispatcher<'a> { R: LspRequest, R::Params: Send + 'static + UnwindSafe, { - if let Some((id, params)) = self.parse::() { + if let Some((request, params)) = self.parse::() { let snapshot = self.global_state.snapshot(); self.global_state.task_pool.handle.spawn(intent, move || { - crate::panic::catch_unwind(|| { - let result = handler(&snapshot, params); - result_to_response::(id.clone(), result) - }) - .unwrap_or_else(|error| panic_response(id, &error)) + let result = crate::panic::catch_unwind(|| handler(&snapshot, params)); + match thread_result_to_response::(request.id.clone(), result) { + Ok(response) => TaskResult::Response(response), + Err(_cancelled) if ALLOW_RETRYING => TaskResult::Retry(request), + Err(_cancelled) => TaskResult::Response(Response::new_err( + request.id, + lsp_server::ErrorCode::ContentModified as i32, + "content modified".to_owned(), + )), + } }); } @@ -110,45 +145,75 @@ impl<'a> RequestDispatcher<'a> { } } -fn panic_response(id: RequestId, error: &crate::panic::PanicError) -> Response { - // Check if the request was canceled due to some modifications to the salsa database. - if error.payload.downcast_ref::().is_some() { - // TODO: trigger retries when we have that setup, we'll reenque the task - log::debug!( - "request id={} was cancelled by salsa, sending content modified", - id - ); - Response::new_err( - id, - lsp_server::ErrorCode::ContentModified as i32, - "content modified".to_string(), - ) - } else { - Response::new_err( - id, - lsp_server::ErrorCode::InternalError as i32, - "request handler error".to_string(), - ) +fn thread_result_to_response( + id: lsp_server::RequestId, + result: Result, PanicError>, +) -> Result +where + R: lsp_types::request::Request, + R::Params: DeserializeOwned, + R::Result: Serialize, +{ + match result { + Ok(handler_result) => match handler_result { + Ok(result) => Ok(Response::new_ok(id, result)), + Err(error) => Ok(Response::new_err( + id, + lsp_server::ErrorCode::InternalError as i32, + error.to_string(), + )), + }, + Err(panic) => { + // Check if the request was canceled due to some modifications to the salsa database. + if panic.payload.downcast_ref::().is_some() { + log::debug!( + "request id={} was cancelled by salsa, sending content modified", + id + ); + Err(panic) + } else { + let error = panic.to_string(); + // we don't retry non-salsa cancellation panics + Ok(Response::new_err( + id, + lsp_server::ErrorCode::InternalError as i32, + format!("request handler error: {error}"), + )) + } + } } } -fn result_to_response(id: RequestId, result: Result) -> Response +fn result_to_response( + id: lsp_server::RequestId, + result: anyhow::Result, +) -> std::result::Result where - R: LspRequest, + R: lsp_types::request::Request, { match result { - Ok(result) => Response::new_ok(id, result), - Err(err) => { - error!("Request handler failed: {err}"); - Response::new_err( + Ok(resp) => Ok(lsp_server::Response::new_ok(id, &resp)), + Err(e) => match e.downcast::() { + Ok(cancelled) => Err(cancelled), + Err(e) => Ok(lsp_server::Response::new_err( id, lsp_server::ErrorCode::InternalError as i32, - err.to_string(), - ) - } + e.to_string(), + )), + }, } } +// lsp-server has req.extract(R::METHOD), but it doesn't work for us due to +// ownership so we use this instead. +pub fn from_json( + what: &'static str, + json: &serde_json::Value, +) -> anyhow::Result { + serde_json::from_value(json.clone()) + .map_err(|e| anyhow::format_err!("Failed to deserialize {what}: {e}; {json}")) +} + pub(crate) struct NotificationDispatcher<'a> { notif: Option, state: &'a mut GlobalState, diff --git a/crates/squawk_server/src/global_state.rs b/crates/squawk_server/src/global_state.rs index 3f17fa43..0315ba62 100644 --- a/crates/squawk_server/src/global_state.rs +++ b/crates/squawk_server/src/global_state.rs @@ -2,7 +2,7 @@ use std::{num::NonZeroUsize, sync::Arc, time::Instant}; use crossbeam_channel::{Receiver, Sender, select, unbounded}; use log::info; -use lsp_server::{Message, Response}; +use lsp_server::{Message, Request, Response}; use lsp_types::Url; use lsp_types::notification::Notification as _; use lsp_types::notification::{ @@ -40,7 +40,7 @@ pub(super) struct GlobalState { files: Arc>, req_queue: ReqQueue, sender: Sender, - pub(crate) task_pool: Handle, Receiver>, + pub(crate) task_pool: Handle, Receiver>, shutdown_requested: bool, } @@ -119,6 +119,10 @@ impl GlobalState { } } + pub(crate) fn is_completed(&self, request: &lsp_server::Request) -> bool { + self.req_queue.incoming.is_completed(&request.id) + } + pub(crate) fn request_shutdown(&mut self) { self.shutdown_requested = true; } @@ -134,41 +138,7 @@ impl GlobalState { let loop_start = Instant::now(); match event { Event::Inbox(msg) => match msg { - Message::Request(req) => { - info!("Received request: method={}, id={:?}", req.method, req.id); - - self.register_request(&req, loop_start); - - if self.shutdown_requested { - tracing::warn!( - "Received request `{}` after server shutdown was requested, discarding", - &req.method - ); - - self.respond(Response::new_err( - req.id, - lsp_server::ErrorCode::InvalidRequest as i32, - "Shutdown already requested".to_owned(), - )); - continue; - } - - RequestDispatcher::new(req, self) - .on_sync_mut::(handle_shutdown) - .on::(handle_goto_definition) - .on::(handle_hover) - .on::(handle_code_action) - .on::(handle_selection_range) - .on::(handle_inlay_hints) - .on::(handle_document_symbol) - .on::(handle_folding_range) - .on_latency_sensitive::(handle_completion) - .on::(handle_document_diagnostic) - .on::(handle_syntax_tree) - .on::(handle_tokens) - .on::(handle_references) - .finish(); - } + Message::Request(req) => self.handle_request(req, loop_start), Message::Response(resp) => { info!("Received response: id={:?}", resp.id); } @@ -187,11 +157,19 @@ impl GlobalState { .finish(); } }, - Event::Outbox(response) => { - // Instead of having the tasks send directly via the sender - // channel, we handle them on the main thread so we can check - // for cancellation first. - self.respond(response) + Event::Outbox(result) => { + match result { + TaskResult::Response(resp) => { + // Instead of having the tasks send directly via the sender + // channel, we handle them on the main thread so we can check + // for cancellation first. + self.respond(resp) + } + TaskResult::Retry(req) if !self.is_completed(&req) => { + self.handle_request(req, loop_start) + } + TaskResult::Retry(_) => (), + } } } } @@ -202,13 +180,58 @@ impl GlobalState { fn next_event( &self, inbox: &Receiver, - outbox: &Receiver, + outbox: &Receiver, ) -> Result { select! { recv(inbox) -> msg => msg.map(Event::Inbox), recv(outbox) -> task => task.map(Event::Outbox), } } + + fn handle_request(&mut self, req: Request, loop_start: Instant) { + info!("Received request: method={}, id={:?}", req.method, req.id); + + self.register_request(&req, loop_start); + + if self.shutdown_requested { + tracing::warn!( + "Received request `{}` after server shutdown was requested, discarding", + &req.method + ); + + self.respond(Response::new_err( + req.id, + lsp_server::ErrorCode::InvalidRequest as i32, + "Shutdown already requested".to_owned(), + )); + return; + } + + const RETRY: bool = true; + const NO_RETRY: bool = false; + + RequestDispatcher::new(req, self) + // Request handlers that must run on the main thread because they + // mutate GlobalState: + .on_sync_mut::(handle_shutdown) + // Request handlers which are related to the user typing are run on + // the main thread to reduce latency: + .on_sync::(handle_selection_range) + // latency-sensitive but can't run on the main thread due to + // semantic analysis, so we use a higher priority thread + .on_latency_sensitive::(handle_completion) + .on::(handle_goto_definition) + .on::(handle_hover) + .on::(handle_code_action) + .on::(handle_inlay_hints) + .on::(handle_document_symbol) + .on::(handle_folding_range) + .on::(handle_document_diagnostic) + .on::(handle_syntax_tree) + .on::(handle_tokens) + .on::(handle_references) + .finish(); + } } pub(crate) struct Snapshot { @@ -228,5 +251,10 @@ impl Snapshot { enum Event { Inbox(Message), - Outbox(Response), + Outbox(TaskResult), +} + +pub(crate) enum TaskResult { + Response(Response), + Retry(Request), }