5 files changed, 101 insertions, 56 deletions

diff --git a/crates/ra_lsp_server/src/lib.rs b/crates/ra_lsp_server/src/lib.rs
index 2c5d7c72d..69a577b3e 100644
--- a/crates/ra_lsp_server/src/lib.rs
+++ b/crates/ra_lsp_server/src/lib.rs
@@ -8,6 +8,7 @@ mod project_model;
 pub mod req;
 pub mod config;
 mod world;
+mod thread_worker;
 
 pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
 pub use crate::{
diff --git a/crates/ra_lsp_server/src/main.rs b/crates/ra_lsp_server/src/main.rs
index 1debe7660..88504bb89 100644
--- a/crates/ra_lsp_server/src/main.rs
+++ b/crates/ra_lsp_server/src/main.rs
@@ -1,5 +1,5 @@
 use flexi_logger::{Duplicate, Logger};
-use lsp_server::{run_server, stdio_transport, LspServerError};
+use lsp_server::Connection;
 
 use ra_lsp_server::{show_message, Result, ServerConfig};
 use ra_prof;
@@ -29,46 +29,46 @@ fn main() -> Result<()> {
 }
 
 fn main_inner() -> Result<()> {
-    let (sender, receiver, io_threads) = stdio_transport();
     let cwd = std::env::current_dir()?;
-    let caps = serde_json::to_value(ra_lsp_server::server_capabilities()).unwrap();
-    run_server(caps, sender, receiver, |params, s, r| {
-        let params: lsp_types::InitializeParams = serde_json::from_value(params)?;
-        let root = params.root_uri.and_then(|it| it.to_file_path().ok()).unwrap_or(cwd);
+    let (connection, io_threads) = Connection::stdio();
+    let server_capabilities = serde_json::to_value(ra_lsp_server::server_capabilities()).unwrap();
 
-        let workspace_roots = params
-            .workspace_folders
-            .map(|workspaces| {
-                workspaces
-                    .into_iter()
-                    .filter_map(|it| it.uri.to_file_path().ok())
-                    .collect::<Vec<_>>()
-            })
-            .filter(|workspaces| !workspaces.is_empty())
-            .unwrap_or_else(|| vec![root]);
+    let initialize_params = connection.initialize(server_capabilities)?;
+    let initialize_params: lsp_types::InitializeParams = serde_json::from_value(initialize_params)?;
 
-        let server_config: ServerConfig = params
-            .initialization_options
-            .and_then(|v| {
-                serde_json::from_value(v)
-                    .map_err(|e| {
-                        log::error!("failed to deserialize config: {}", e);
-                        show_message(
-                            lsp_types::MessageType::Error,
-                            format!("failed to deserialize config: {}", e),
-                            s,
-                        );
-                    })
-                    .ok()
-            })
-            .unwrap_or_default();
+    let root = initialize_params.root_uri.and_then(|it| it.to_file_path().ok()).unwrap_or(cwd);
+
+    let workspace_roots = initialize_params
+        .workspace_folders
+        .map(|workspaces| {
+            workspaces.into_iter().filter_map(|it| it.uri.to_file_path().ok()).collect::<Vec<_>>()
+        })
+        .filter(|workspaces| !workspaces.is_empty())
+        .unwrap_or_else(|| vec![root]);
+
+    let server_config: ServerConfig = initialize_params
+        .initialization_options
+        .and_then(|v| {
+            serde_json::from_value(v)
+                .map_err(|e| {
+                    log::error!("failed to deserialize config: {}", e);
+                    show_message(
+                        lsp_types::MessageType::Error,
+                        format!("failed to deserialize config: {}", e),
+                        &connection.sender,
+                    );
+                })
+                .ok()
+        })
+        .unwrap_or_default();
+
+    ra_lsp_server::main_loop(
+        workspace_roots,
+        initialize_params.capabilities,
+        server_config,
+        &connection,
+    )?;
 
-        ra_lsp_server::main_loop(workspace_roots, params.capabilities, server_config, r, s)
-    })
-    .map_err(|err| match err {
-        LspServerError::ProtocolError(err) => err.into(),
-        LspServerError::ServerError(err) => err,
-    })?;
     log::info!("shutting down IO...");
     io_threads.join()?;
     log::info!("... IO is down");
diff --git a/crates/ra_lsp_server/src/main_loop.rs b/crates/ra_lsp_server/src/main_loop.rs
index fb357b36b..42ebb5cdf 100644
--- a/crates/ra_lsp_server/src/main_loop.rs
+++ b/crates/ra_lsp_server/src/main_loop.rs
@@ -5,7 +5,7 @@ pub(crate) mod pending_requests;
 use std::{error::Error, fmt, path::PathBuf, sync::Arc, time::Instant};
 
 use crossbeam_channel::{select, unbounded, Receiver, RecvError, Sender};
-use lsp_server::{handle_shutdown, ErrorCode, Message, Notification, Request, RequestId, Response};
+use lsp_server::{Connection, ErrorCode, Message, Notification, Request, RequestId, Response};
 use lsp_types::{ClientCapabilities, NumberOrString};
 use ra_ide_api::{Canceled, FeatureFlags, FileId, LibraryData};
 use ra_prof::profile;
@@ -51,8 +51,7 @@ pub fn main_loop(
     ws_roots: Vec<PathBuf>,
     client_caps: ClientCapabilities,
     config: ServerConfig,
-    msg_receiver: &Receiver<Message>,
-    msg_sender: &Sender<Message>,
+    connection: &Connection,
 ) -> Result<()> {
     log::info!("server_config: {:#?}", config);
     // FIXME: support dynamic workspace loading.
@@ -69,7 +68,7 @@ pub fn main_loop(
                     show_message(
                         req::MessageType::Error,
                         format!("rust-analyzer failed to load workspace: {}", e),
-                        msg_sender,
+                        &connection.sender,
                     );
                 }
             }
@@ -89,7 +88,7 @@ pub fn main_loop(
                 show_message(
                     req::MessageType::Error,
                     format!("unknown feature flag: {:?}", flag),
-                    msg_sender,
+                    &connection.sender,
                 );
             }
         }
@@ -119,8 +118,7 @@ pub fn main_loop(
     log::info!("server initialized, serving requests");
     let main_res = main_loop_inner(
         &pool,
-        msg_sender,
-        msg_receiver,
+        connection,
         task_sender,
         task_receiver.clone(),
         &mut state,
@@ -130,7 +128,7 @@ pub fn main_loop(
     log::info!("waiting for tasks to finish...");
     task_receiver
         .into_iter()
-        .for_each(|task| on_task(task, msg_sender, &mut pending_requests, &mut state));
+        .for_each(|task| on_task(task, &connection.sender, &mut pending_requests, &mut state));
     log::info!("...tasks have finished");
     log::info!("joining threadpool...");
     drop(pool);
@@ -196,8 +194,7 @@ impl fmt::Debug for Event {
 
 fn main_loop_inner(
     pool: &ThreadPool,
-    msg_sender: &Sender<Message>,
-    msg_receiver: &Receiver<Message>,
+    connection: &Connection,
     task_sender: Sender<Task>,
     task_receiver: Receiver<Task>,
     state: &mut WorldState,
@@ -214,7 +211,7 @@ fn main_loop_inner(
     loop {
         log::trace!("selecting");
         let event = select! {
-            recv(msg_receiver) -> msg => match msg {
+            recv(&connection.receiver) -> msg => match msg {
                 Ok(msg) => Event::Msg(msg),
                 Err(RecvError) => Err("client exited without shutdown")?,
             },
@@ -238,7 +235,7 @@ fn main_loop_inner(
         let mut state_changed = false;
         match event {
             Event::Task(task) => {
-                on_task(task, msg_sender, pending_requests, state);
+                on_task(task, &connection.sender, pending_requests, state);
                 state.maybe_collect_garbage();
             }
             Event::Vfs(task) => {
@@ -252,7 +249,7 @@ fn main_loop_inner(
             }
             Event::Msg(msg) => match msg {
                 Message::Request(req) => {
-                    if handle_shutdown(&req, msg_sender) {
+                    if connection.handle_shutdown(&req)? {
                         return Ok(());
                     };
                     on_request(
@@ -260,13 +257,13 @@ fn main_loop_inner(
                         pending_requests,
                         pool,
                         &task_sender,
-                        msg_sender,
+                        &connection.sender,
                         loop_start,
                         req,
                     )?
                 }
                 Message::Notification(not) => {
-                    on_notification(msg_sender, state, pending_requests, &mut subs, not)?;
+                    on_notification(&connection.sender, state, pending_requests, &mut subs, not)?;
                     state_changed = true;
                 }
                 Message::Response(resp) => log::error!("unexpected response: {:?}", resp),
@@ -294,7 +291,7 @@ fn main_loop_inner(
             let n_packages: usize = state.workspaces.iter().map(|it| it.n_packages()).sum();
             if state.feature_flags().get("notifications.workspace-loaded") {
                 let msg = format!("workspace loaded, {} rust packages", n_packages);
-                show_message(req::MessageType::Info, msg, msg_sender);
+                show_message(req::MessageType::Info, msg, &connection.sender);
             }
             // Only send the notification first time
             send_workspace_notification = false;
diff --git a/crates/ra_lsp_server/src/project_model.rs b/crates/ra_lsp_server/src/project_model.rs
index ad59cde64..6234563f2 100644
--- a/crates/ra_lsp_server/src/project_model.rs
+++ b/crates/ra_lsp_server/src/project_model.rs
@@ -1,8 +1,6 @@
 use std::path::PathBuf;
 
-use thread_worker::Worker;
-
-use crate::Result;
+use crate::{thread_worker::Worker, Result};
 
 pub use ra_project_model::{
     CargoWorkspace, Package, ProjectWorkspace, Sysroot, Target, TargetKind,
diff --git a/crates/ra_lsp_server/src/thread_worker.rs b/crates/ra_lsp_server/src/thread_worker.rs
new file mode 100644
index 000000000..68e5c124d
--- /dev/null
+++ b/crates/ra_lsp_server/src/thread_worker.rs
@@ -0,0 +1,49 @@
+//! Small utility to correctly spawn crossbeam-channel based worker threads.
+
+use crossbeam_channel::{bounded, unbounded, Receiver, Sender};
+
+/// A wrapper around event-processing thread with automatic shutdown semantics.
+pub struct Worker<I, O> {
+    // XXX: field order is significant here.
+    //
+    // In Rust, fields are dropped in the declaration order, and we rely on this
+    // here. We must close input first, so that the  `thread` (who holds the
+    // opposite side of the channel) noticed shutdown. Then, we must join the
+    // thread, but we must keep out alive so that the thread does not panic.
+    //
+    // Note that a potential problem here is that we might drop some messages
+    // from receiver on the floor. This is ok for rust-analyzer: we have only a
+    // single client, so, if we are shutting down, nobody is interested in the
+    // unfinished work anyway!
+    sender: Sender<I>,
+    _thread: jod_thread::JoinHandle<()>,
+    receiver: Receiver<O>,
+}
+
+impl<I, O> Worker<I, O> {
+    pub fn spawn<F>(name: &'static str, buf: usize, f: F) -> Worker<I, O>
+    where
+        F: FnOnce(Receiver<I>, Sender<O>) + Send + 'static,
+        I: Send + 'static,
+        O: Send + 'static,
+    {
+        // Set up worker channels in a deadlock-avoiding way. If one sets both input
+        // and output buffers to a fixed size, a worker might get stuck.
+        let (sender, input_receiver) = bounded::<I>(buf);
+        let (output_sender, receiver) = unbounded::<O>();
+        let _thread = jod_thread::Builder::new()
+            .name(name.to_string())
+            .spawn(move || f(input_receiver, output_sender))
+            .expect("failed to spawn a thread");
+        Worker { sender, _thread, receiver }
+    }
+}
+
+impl<I, O> Worker<I, O> {
+    pub fn sender(&self) -> &Sender<I> {
+        &self.sender
+    }
+    pub fn receiver(&self) -> &Receiver<O> {
+        &self.receiver
+    }
+}