server: add retry support (#1032)

Author: sbdchd

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<lsp_server::Request>,
@@ -23,15 +28,14 @@ impl<'a> RequestDispatcher<'a> {
         }
     }
 
-    fn parse<R>(&mut self) -> Option<(lsp_server::RequestId, R::Params)>
+    fn parse<R>(&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<R>(
+        mut self,
+        handler: fn(&Snapshot, R::Params) -> Result<R::Result>,
+    ) -> Self
+    where
+        R: LspRequest,
+        R::Params: Send + 'static + UnwindSafe,
+    {
+        let Some((request, params)) = self.parse::<R>() 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::<R>(request.id.clone(), result) {
+            self.global_state.respond(response);
+        }
+        self
+    }
+
     pub(crate) fn on_sync_mut<R>(
         mut self,
         handler: fn(&mut GlobalState, R::Params) -> Result<R::Result>,
@@ -52,34 +76,40 @@ impl<'a> RequestDispatcher<'a> {
         R: LspRequest,
         R::Params: Send + 'static,
     {
-        if let Some((id, params)) = self.parse::<R>() {
-            let result = handler(self.global_state, params);
-            let response = result_to_response::<R>(id, result);
+        let Some((request, params)) = self.parse::<R>() else {
+            return self;
+        };
+
+        let result = handler(self.global_state, params);
+        if let Ok(response) = result_to_response::<R>(request.id.clone(), result) {
             self.global_state.respond(response);
         }
         self
     }
 
-    pub(crate) fn on<R>(self, handler: fn(&Snapshot, R::Params) -> Result<R::Result>) -> Self
+    pub(crate) fn on<const ALLOW_RETRYING: bool, R>(
+        self,
+        handler: fn(&Snapshot, R::Params) -> Result<R::Result>,
+    ) -> Self
     where
         R: LspRequest,
         R::Params: Send + 'static + UnwindSafe,
     {
-        self.on_with_thread_intent::<R>(ThreadIntent::Worker, handler)
+        self.on_with_thread_intent::<ALLOW_RETRYING, R>(ThreadIntent::Worker, handler)
     }
 
-    pub(crate) fn on_latency_sensitive<R>(
+    pub(crate) fn on_latency_sensitive<const ALLOW_RETRYING: bool, R>(
         self,
         handler: fn(&Snapshot, R::Params) -> Result<R::Result>,
     ) -> Self
     where
         R: LspRequest,
         R::Params: Send + 'static + UnwindSafe,
     {
-        self.on_with_thread_intent::<R>(ThreadIntent::LatencySensitive, handler)
+        self.on_with_thread_intent::<ALLOW_RETRYING, R>(ThreadIntent::LatencySensitive, handler)
     }
 
-    fn on_with_thread_intent<R>(
+    fn on_with_thread_intent<const ALLOW_RETRYING: bool, R>(
         mut self,
         intent: ThreadIntent,
         handler: fn(&Snapshot, R::Params) -> Result<R::Result>,
@@ -88,15 +118,20 @@ impl<'a> RequestDispatcher<'a> {
         R: LspRequest,
         R::Params: Send + 'static + UnwindSafe,
     {
-        if let Some((id, params)) = self.parse::<R>() {
+        if let Some((request, params)) = self.parse::<R>() {
             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::<R>(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::<R>(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::<salsa::Cancelled>().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<R>(
+    id: lsp_server::RequestId,
+    result: Result<anyhow::Result<R::Result>, PanicError>,
+) -> Result<lsp_server::Response, PanicError>
+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::<salsa::Cancelled>().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<R>(id: RequestId, result: Result<R::Result>) -> Response
+fn result_to_response<R>(
+    id: lsp_server::RequestId,
+    result: anyhow::Result<R::Result>,
+) -> std::result::Result<lsp_server::Response, Cancelled>
 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::<Cancelled>() {
+            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<T: DeserializeOwned>(
+    what: &'static str,
+    json: &serde_json::Value,
+) -> anyhow::Result<T> {
+    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<lsp_server::Notification>,
     state: &'a mut GlobalState,