server: multi threading via task pool & crossbeam (#1026)

We're only running completions in latency sensitive threads, but we'll always want to run semantic tokens in there once we build that.

Author: sbdchd

Cargo.lock

@@ -31,7 +31,7 @@ pub fn bind(db: &dyn Db, file: File) -> Binder {
 }
 
 #[salsa::db]
-#[derive(Default)]
+#[derive(Clone, Default)]
 pub struct Database {
     storage: Storage<Self>,
 }

crates/squawk_ide/src/db.rs

@@ -31,6 +31,8 @@ line-index.workspace = true
 insta.workspace = true
 etcetera.workspace = true
 rustc-hash.workspace = true
+squawk-thread.workspace = true
+crossbeam-channel.workspace = true
 
 [lints]
 workspace = true

crates/squawk_server/Cargo.toml

@@ -4,23 +4,18 @@ use anyhow::Result;
 use log::{error, info};
 use lsp_server::{Connection, Message, Response};
 use lsp_types::{notification::Notification as LspNotification, request::Request as LspRequest};
+use squawk_thread::ThreadIntent;
 
-use crate::system::System;
+use crate::system::{GlobalState, MutableSystem, System};
 
 pub(crate) struct RequestDispatcher<'a> {
-    connection: &'a Connection,
     req: Option<lsp_server::Request>,
-    system: &'a dyn System,
+    system: &'a mut GlobalState,
 }
 
 impl<'a> RequestDispatcher<'a> {
-    pub(crate) fn new(
-        connection: &'a Connection,
-        req: lsp_server::Request,
-        system: &'a dyn System,
-    ) -> Self {
+    pub(crate) fn new(req: lsp_server::Request, system: &'a mut GlobalState) -> Self {
         Self {
-            connection,
             req: Some(req),
             system,
         }
@@ -41,7 +36,7 @@ impl<'a> RequestDispatcher<'a> {
                     lsp_server::ErrorCode::ParseError as i32,
                     err.to_string(),
                 );
-                if let Err(err) = self.connection.sender.send(Message::Response(response)) {
+                if let Err(err) = self.system.sender.send(Message::Response(response)) {
                     error!("Failed to send parse error response: {err}");
                 }
                 None
@@ -50,25 +45,54 @@ impl<'a> RequestDispatcher<'a> {
     }
 
     pub(crate) fn on<R>(
+        self,
+        handler: fn(&dyn System, R::Params) -> Result<R::Result>,
+    ) -> Result<Self>
+    where
+        R: LspRequest,
+        R::Params: Send + 'static,
+    {
+        self.on_with_thread_intent::<R>(ThreadIntent::Worker, handler)
+    }
+
+    pub(crate) fn on_latency_sensitive<R>(
+        self,
+        handler: fn(&dyn System, R::Params) -> Result<R::Result>,
+    ) -> Result<Self>
+    where
+        R: LspRequest,
+        R::Params: Send + 'static,
+    {
+        self.on_with_thread_intent::<R>(ThreadIntent::LatencySensitive, handler)
+    }
+
+    fn on_with_thread_intent<R>(
         mut self,
+        intent: ThreadIntent,
         handler: fn(&dyn System, R::Params) -> Result<R::Result>,
     ) -> Result<Self>
     where
         R: LspRequest,
+        R::Params: Send + 'static,
     {
         if let Some((id, params)) = self.parse::<R>() {
-            let resp = match handler(self.system, params) {
-                Ok(result) => Response::new_ok(id, result),
-                Err(err) => {
-                    error!("Request handler failed: {err}");
-                    Response::new_err(
-                        id,
-                        lsp_server::ErrorCode::InternalError as i32,
-                        err.to_string(),
-                    )
-                }
-            };
-            self.connection.sender.send(Message::Response(resp))?;
+            let snapshot = self.system.snapshot();
+
+            self.system.task_pool.spawn(intent, move || {
+                let resp = match handler(&snapshot, params) {
+                    Ok(result) => Response::new_ok(id, result),
+                    Err(err) => {
+                        error!("Request handler failed: {err}");
+                        Response::new_err(
+                            id,
+                            lsp_server::ErrorCode::InternalError as i32,
+                            err.to_string(),
+                        )
+                    }
+                };
+
+                Message::Response(resp)
+            });
         }
 
         Ok(self)
@@ -84,14 +108,14 @@ impl<'a> RequestDispatcher<'a> {
 pub(crate) struct NotificationDispatcher<'a> {
     connection: &'a Connection,
     notif: Option<lsp_server::Notification>,
-    system: &'a mut dyn System,
+    system: &'a mut dyn MutableSystem,
 }
 
 impl<'a> NotificationDispatcher<'a> {
     pub(crate) fn new(
         connection: &'a Connection,
         notif: lsp_server::Notification,
-        system: &'a mut dyn System,
+        system: &'a mut dyn MutableSystem,
     ) -> Self {
         Self {
             connection,
@@ -119,7 +143,7 @@ impl<'a> NotificationDispatcher<'a> {
 
     pub(crate) fn on<N>(
         mut self,
-        handler: fn(&Connection, N::Params, &mut dyn System) -> Result<()>,
+        handler: fn(&Connection, N::Params, &mut dyn MutableSystem) -> Result<()>,
     ) -> Result<Self>
     where
         N: LspNotification,

crates/squawk_server/src/dispatch.rs

@@ -7,12 +7,12 @@ use lsp_types::{
 };
 
 use crate::lsp_utils;
-use crate::system::System;
+use crate::system::MutableSystem;
 
 pub(crate) fn handle_did_open(
     _connection: &Connection,
     params: DidOpenTextDocumentParams,
-    system: &mut dyn System,
+    system: &mut dyn MutableSystem,
 ) -> Result<()> {
     let uri = params.text_document.uri;
     let content = params.text_document.text;
@@ -25,7 +25,7 @@ pub(crate) fn handle_did_open(
 pub(crate) fn handle_did_change(
     _connection: &Connection,
     params: DidChangeTextDocumentParams,
-    system: &mut dyn System,
+    system: &mut dyn MutableSystem,
 ) -> Result<()> {
     let uri = params.text_document.uri;
 
@@ -43,7 +43,7 @@ pub(crate) fn handle_did_change(
 pub(crate) fn handle_did_close(
     connection: &Connection,
     params: DidCloseTextDocumentParams,
-    system: &mut dyn System,
+    system: &mut dyn MutableSystem,
 ) -> Result<()> {
     let uri = params.text_document.uri;
 

crates/squawk_server/src/handlers/notifications.rs

@@ -1,3 +1,5 @@
+use std::num::NonZeroUsize;
+
 use anyhow::Result;
 use log::info;
 use lsp_server::{Connection, Message};
@@ -90,7 +92,8 @@ fn main_loop(connection: Connection, params: serde_json::Value) -> Result<()> {
     let client_name = init_params.client_info.map(|x| x.name);
     info!("Client name: {client_name:?}");
 
-    let mut system = GlobalState::new();
+    let threads = std::thread::available_parallelism().unwrap_or(NonZeroUsize::MIN);
+    let mut system = GlobalState::new(connection.sender.clone(), threads);
 
     for msg in &connection.receiver {
         match msg {
@@ -102,15 +105,15 @@ fn main_loop(connection: Connection, params: serde_json::Value) -> Result<()> {
                     return Ok(());
                 }
 
-                RequestDispatcher::new(&connection, req, &system)
+                RequestDispatcher::new(req, &mut system)
                     .on::<GotoDefinition>(handle_goto_definition)?
                     .on::<HoverRequest>(handle_hover)?
                     .on::<CodeActionRequest>(handle_code_action)?
                     .on::<SelectionRangeRequest>(handle_selection_range)?
                     .on::<InlayHintRequest>(handle_inlay_hints)?
                     .on::<DocumentSymbolRequest>(handle_document_symbol)?
                     .on::<FoldingRangeRequest>(handle_folding_range)?
-                    .on::<Completion>(handle_completion)?
+                    .on_latency_sensitive::<Completion>(handle_completion)?
                     .on::<DocumentDiagnosticRequest>(handle_document_diagnostic)?
                     .on::<SyntaxTreeRequest>(handle_syntax_tree)?
                     .on::<TokensRequest>(handle_tokens)?

crates/squawk_server/src/server.rs

@@ -1,48 +1,84 @@
+use std::{num::NonZeroUsize, sync::Arc};
+
+use crossbeam_channel::Sender;
+use lsp_server::Message;
 use lsp_types::Url;
 use rustc_hash::FxHashMap;
 use salsa::Setter;
 use squawk_ide::db::{Database, File};
+use squawk_thread::TaskPool;
 
 pub(crate) trait System {
     fn db(&self) -> &Database;
     fn file(&self, uri: &Url) -> Option<File>;
+}
+
+pub(crate) trait MutableSystem: System {
     fn set(&mut self, uri: Url, content: String);
     fn remove(&mut self, uri: &Url);
 }
 
+pub(crate) struct Snapshot {
+    db: Database,
+    files: Arc<FxHashMap<Url, File>>,
+}
+
+impl System for Snapshot {
+    fn db(&self) -> &Database {
+        &self.db
+    }
+
+    fn file(&self, uri: &Url) -> Option<File> {
+        self.files.get(uri).copied()
+    }
+}
+
 pub(super) struct GlobalState {
-    pub db: Database,
-    files: FxHashMap<Url, File>,
+    db: Database,
+    files: Arc<FxHashMap<Url, File>>,
+    pub(crate) sender: Sender<Message>,
+    pub(crate) task_pool: TaskPool<Message>,
 }
 
 impl GlobalState {
-    pub(super) fn new() -> Self {
+    pub(super) fn new(sender: Sender<Message>, threads: NonZeroUsize) -> Self {
         Self {
             db: Database::default(),
-            files: FxHashMap::default(),
+            files: Arc::new(FxHashMap::default()),
+            task_pool: TaskPool::new_with_threads(sender.clone(), threads),
+            sender,
+        }
+    }
+
+    pub(crate) fn snapshot(&self) -> Snapshot {
+        Snapshot {
+            db: self.db.clone(),
+            files: self.files.clone(),
         }
     }
 }
 
 impl System for GlobalState {
     fn db(&self) -> &Database {
-        return &self.db;
+        &self.db
     }
 
     fn file(&self, uri: &Url) -> Option<File> {
         self.files.get(uri).copied()
     }
+}
 
+impl MutableSystem for GlobalState {
     fn set(&mut self, uri: Url, content: String) {
         if let Some(file) = self.files.get(&uri).copied() {
             file.set_content(&mut self.db).to(content.into());
         } else {
             let file = File::new(&self.db, content.into());
-            self.files.insert(uri, file);
+            Arc::make_mut(&mut self.files).insert(uri, file);
         }
     }
 
     fn remove(&mut self, uri: &Url) {
-        self.files.remove(uri);
+        Arc::make_mut(&mut self.files).remove(uri);
     }
 }

crates/squawk_server/src/system.rs

@@ -54,7 +54,7 @@ impl Pool {
 
         let mut handles = Vec::with_capacity(threads.into());
         for idx in 0..threads.into() {
-            let handle = Builder::new(INITIAL_INTENT, format!("squawk:worker:{idx}",))
+            let handle = Builder::new(INITIAL_INTENT, format!("Worker:{idx}",))
                 .stack_size(STACK_SIZE)
                 .allow_leak(true)
                 .spawn({

crates/squawk_thread/src/pool.rs

@@ -2,7 +2,7 @@
 // https://github.com/rust-lang/rust-analyzer/blob/2efc80078029894eec0699f62ec8d5c1a56af763/crates/rust-analyzer/src/task_pool.rs#L6C19-L6C19
 //! A thin wrapper around [`crate::Pool`] which threads a sender through spawned jobs.
 
-use std::{num::NonZeroUsize, panic::UnwindSafe};
+use std::num::NonZeroUsize;
 
 use crossbeam_channel::Sender;
 
@@ -23,7 +23,7 @@ impl<T> TaskPool<T> {
 
     pub fn spawn<F>(&mut self, intent: ThreadIntent, task: F)
     where
-        F: FnOnce() -> T + Send + UnwindSafe + 'static,
+        F: FnOnce() -> T + Send + 'static,
         T: Send + 'static,
     {
         self.pool.spawn(intent, {
@@ -34,7 +34,7 @@ impl<T> TaskPool<T> {
 
     pub fn spawn_with_sender<F>(&mut self, intent: ThreadIntent, task: F)
     where
-        F: FnOnce(Sender<T>) + Send + UnwindSafe + 'static,
+        F: FnOnce(Sender<T>) + Send + 'static,
         T: Send + 'static,
     {
         self.pool.spawn(intent, {