diff --git a/Cargo.lock b/Cargo.lock index a89e743a..8692bc82 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2496,6 +2496,7 @@ name = "squawk-server" version = "2.44.0" dependencies = [ "anyhow", + "crossbeam-channel", "etcetera", "insta", "line-index", @@ -2512,6 +2513,7 @@ dependencies = [ "squawk-lexer", "squawk-linter", "squawk-syntax", + "squawk-thread", ] [[package]] diff --git a/crates/squawk_ide/src/db.rs b/crates/squawk_ide/src/db.rs index d3af2b01..675b9a10 100644 --- a/crates/squawk_ide/src/db.rs +++ b/crates/squawk_ide/src/db.rs @@ -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, } diff --git a/crates/squawk_server/Cargo.toml b/crates/squawk_server/Cargo.toml index f3ded112..93050679 100644 --- a/crates/squawk_server/Cargo.toml +++ b/crates/squawk_server/Cargo.toml @@ -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 diff --git a/crates/squawk_server/src/dispatch.rs b/crates/squawk_server/src/dispatch.rs index 44547602..d294ecc3 100644 --- a/crates/squawk_server/src/dispatch.rs +++ b/crates/squawk_server/src/dispatch.rs @@ -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, - 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( + self, + handler: fn(&dyn System, R::Params) -> Result, + ) -> Result + where + R: LspRequest, + R::Params: Send + 'static, + { + self.on_with_thread_intent::(ThreadIntent::Worker, handler) + } + + pub(crate) fn on_latency_sensitive( + self, + handler: fn(&dyn System, R::Params) -> Result, + ) -> Result + where + R: LspRequest, + R::Params: Send + 'static, + { + self.on_with_thread_intent::(ThreadIntent::LatencySensitive, handler) + } + + fn on_with_thread_intent( mut self, + intent: ThreadIntent, handler: fn(&dyn System, R::Params) -> Result, ) -> Result where R: LspRequest, + R::Params: Send + 'static, { if let Some((id, params)) = self.parse::() { - 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, - 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( mut self, - handler: fn(&Connection, N::Params, &mut dyn System) -> Result<()>, + handler: fn(&Connection, N::Params, &mut dyn MutableSystem) -> Result<()>, ) -> Result where N: LspNotification, diff --git a/crates/squawk_server/src/handlers/notifications.rs b/crates/squawk_server/src/handlers/notifications.rs index 515268a7..af836408 100644 --- a/crates/squawk_server/src/handlers/notifications.rs +++ b/crates/squawk_server/src/handlers/notifications.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; diff --git a/crates/squawk_server/src/server.rs b/crates/squawk_server/src/server.rs index d69a86d1..d8b5bf97 100644 --- a/crates/squawk_server/src/server.rs +++ b/crates/squawk_server/src/server.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,7 +105,7 @@ fn main_loop(connection: Connection, params: serde_json::Value) -> Result<()> { return Ok(()); } - RequestDispatcher::new(&connection, req, &system) + RequestDispatcher::new(req, &mut system) .on::(handle_goto_definition)? .on::(handle_hover)? .on::(handle_code_action)? @@ -110,7 +113,7 @@ fn main_loop(connection: Connection, params: serde_json::Value) -> Result<()> { .on::(handle_inlay_hints)? .on::(handle_document_symbol)? .on::(handle_folding_range)? - .on::(handle_completion)? + .on_latency_sensitive::(handle_completion)? .on::(handle_document_diagnostic)? .on::(handle_syntax_tree)? .on::(handle_tokens)? diff --git a/crates/squawk_server/src/system.rs b/crates/squawk_server/src/system.rs index 25e55657..b240d8d9 100644 --- a/crates/squawk_server/src/system.rs +++ b/crates/squawk_server/src/system.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; +} + +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>, +} + +impl System for Snapshot { + fn db(&self) -> &Database { + &self.db + } + + fn file(&self, uri: &Url) -> Option { + self.files.get(uri).copied() + } +} + pub(super) struct GlobalState { - pub db: Database, - files: FxHashMap, + db: Database, + files: Arc>, + pub(crate) sender: Sender, + pub(crate) task_pool: TaskPool, } impl GlobalState { - pub(super) fn new() -> Self { + pub(super) fn new(sender: Sender, 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 { 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); } } diff --git a/crates/squawk_thread/src/pool.rs b/crates/squawk_thread/src/pool.rs index a2b08334..9a06ce10 100644 --- a/crates/squawk_thread/src/pool.rs +++ b/crates/squawk_thread/src/pool.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({ diff --git a/crates/squawk_thread/src/taskpool.rs b/crates/squawk_thread/src/taskpool.rs index ff603a31..20464399 100644 --- a/crates/squawk_thread/src/taskpool.rs +++ b/crates/squawk_thread/src/taskpool.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 TaskPool { pub fn spawn(&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 TaskPool { pub fn spawn_with_sender(&mut self, intent: ThreadIntent, task: F) where - F: FnOnce(Sender) + Send + UnwindSafe + 'static, + F: FnOnce(Sender) + Send + 'static, T: Send + 'static, { self.pool.spawn(intent, {