From e04b1b968d8a030eded6d4e081c802decbe5712a Mon Sep 17 00:00:00 2001
From: oxalica <oxalicc@pm.me>
Date: Mon, 12 Sep 2022 10:20:27 +0800
Subject: [PATCH] Wrap panicking in handlers into RPC failure

---
 crates/nil/src/state.rs | 52 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 48 insertions(+), 4 deletions(-)

diff --git a/crates/nil/src/state.rs b/crates/nil/src/state.rs
index 00dc529..d9c94e1 100644
--- a/crates/nil/src/state.rs
+++ b/crates/nil/src/state.rs
@@ -5,10 +5,12 @@ use lsp_server::{ErrorCode, Message, Notification, Request, RequestId, Response}
 use lsp_types::notification::Notification as _;
 use lsp_types::{notification as notif, request as req, PublishDiagnosticsParams, Url};
 use serde::Serialize;
+use std::cell::Cell;
 use std::collections::HashSet;
-use std::fs;
+use std::panic::UnwindSafe;
 use std::path::PathBuf;
-use std::sync::{Arc, RwLock};
+use std::sync::{Arc, Once, RwLock};
+use std::{fs, panic};
 
 const MAX_DIAGNOSTICS_CNT: usize = 128;
 
@@ -202,11 +204,17 @@ impl<'s> RequestDispatcher<'s> {
         self
     }
 
-    fn on<R: req::Request>(mut self, f: fn(StateSnapshot, R::Params) -> Result<R::Result>) -> Self {
+    fn on<R: req::Request>(mut self, f: fn(StateSnapshot, R::Params) -> Result<R::Result>) -> Self
+    where
+        R::Params: UnwindSafe,
+    {
         if matches!(&self.1, Some(notif) if notif.method == R::METHOD) {
             let req = self.1.take().unwrap();
             let ret = match serde_json::from_value::<R::Params>(req.params) {
-                Ok(params) => result_to_response(req.id, f(self.0.snapshot(), params)),
+                Ok(params) => {
+                    let snap = self.0.snapshot();
+                    result_to_response(req.id, with_catch_unwind(R::METHOD, || f(snap, params)))
+                }
                 Err(err) => Ok(Response::new_err(
                     req.id,
                     ErrorCode::InvalidParams as i32,
@@ -250,6 +258,42 @@ impl<'s> NotificationDispatcher<'s> {
     }
 }
 
+fn with_catch_unwind<T>(ctx: &str, f: impl FnOnce() -> Result<T> + UnwindSafe) -> Result<T> {
+    static INSTALL_PANIC_HOOK: Once = Once::new();
+    thread_local! {
+        static PANIC_LOCATION: Cell<String> = Cell::new(String::new());
+    }
+
+    INSTALL_PANIC_HOOK.call_once(|| {
+        let old_hook = panic::take_hook();
+        panic::set_hook(Box::new(move |info| {
+            if let Some(loc) = info.location() {
+                PANIC_LOCATION.with(|inner| {
+                    inner.set(loc.to_string());
+                });
+            }
+            old_hook(info);
+        }))
+    });
+
+    match panic::catch_unwind(f) {
+        Ok(ret) => ret,
+        Err(payload) => {
+            let reason = payload
+                .downcast_ref::<String>()
+                .map(|s| &**s)
+                .or_else(|| payload.downcast_ref::<&str>().map(|s| &**s))
+                .unwrap_or("unknown");
+            let mut loc = PANIC_LOCATION.with(|inner| inner.take());
+            if loc.is_empty() {
+                loc = "unknown".into();
+            }
+            let msg = format!("Request handler of {} panicked at {}: {}", ctx, loc, reason);
+            Err(msg.into())
+        }
+    }
+}
+
 fn result_to_response(id: RequestId, ret: Result<impl Serialize>) -> Result<Response, Cancelled> {
     match ret {
         Ok(ret) => Ok(Response::new_ok(id, ret)),