Merge pull request #702 from alexheretic/work-pool-dos

Address work-pool DoS with long running tasks

Author: nrc

Cargo.lock

@@ -32,6 +32,7 @@ serde_json = "1.0"
 serde_derive = "1.0"
 url = "1.1.0"
 rayon = "0.9"
+num_cpus = "1"
 
 [dev-dependencies]
 json = "0.11"

Cargo.toml

@@ -53,6 +53,7 @@ macro_rules! parse_file_path {
     }
 }
 
+pub mod work_pool;
 pub mod post_build;
 pub mod requests;
 pub mod notifications;

src/actions/mod.rs

@@ -22,8 +22,9 @@ use rustfmt::{FileName, format_input, Input as FmtInput};
 use rustfmt::file_lines::{FileLines, Range as RustfmtRange};
 use serde_json;
 use span;
-use rayon;
 
+use actions::work_pool;
+use actions::work_pool::WorkDescription;
 use lsp_data;
 use lsp_data::*;
 use server;
@@ -50,7 +51,6 @@ pub use lsp_data::FindImpls;
 
 use std::collections::HashMap;
 use std::path::Path;
-use std::sync::mpsc;
 
 /// Represent the result of a deglob action for a single wildcard import.
 ///
@@ -222,14 +222,14 @@ impl RequestAction for Definition {
         // If configured start racer concurrently and fallback to racer result
         let racer_receiver = {
             if ctx.config.lock().unwrap().goto_def_racer_fallback {
-                Some(receive_from_thread(move || {
+                Some(work_pool::receive_from_thread(move || {
                     let cache = racer::FileCache::new(vfs);
                     let session = racer::Session::new(&cache);
                     let location = pos_to_racer_location(params.position);
 
                     racer::find_definition(file_path, location, &session)
                         .and_then(location_from_racer_match)
-                }))
+                }, WorkDescription("textDocument/definition-racer")))
             } else {
                 None
             }
@@ -850,29 +850,3 @@ fn location_from_racer_match(a_match: racer::Match) -> Option<Location> {
         ls_util::rls_location_to_location(&loc)
     })
 }
-
-lazy_static! {
-    /// Thread pool for request execution allowing concurrent request processing.
-    static ref WORK_POOL: rayon::ThreadPool = rayon::ThreadPool::new(
-        rayon::Configuration::default()
-            .thread_name(|num| format!("request-worker-{}", num))
-            // panic details will be on stderr, otherwise ignore the work panic as it
-            // will already cause a mpsc disconnect-error & there isn't anything else to log
-            .panic_handler(|_| {})
-    ).unwrap();
-}
-
-/// Runs work in a new thread on the `WORK_POOL` returning a result `Receiver`
-/// Panicking work will receive `Err(RecvError)` / `Err(RecvTimeoutError::Disconnected)`
-pub fn receive_from_thread<T, F>(work_fn: F) -> mpsc::Receiver<T>
-where
-    T: Send + 'static,
-    F: FnOnce() -> T + Send + 'static,
-{
-    let (sender, receiver) = mpsc::channel();
-    WORK_POOL.spawn(move || {
-        // an error here simply means the work took too long and the receiver has been dropped
-        let _ = sender.send(work_fn());
-    });
-    receiver
-}

src/actions/requests.rs

@@ -0,0 +1,102 @@
+use rayon;
+use server::DEFAULT_REQUEST_TIMEOUT;
+use std::{fmt, panic};
+use std::time::{Duration, Instant};
+use std::sync::{mpsc, Mutex};
+
+/// Description of work on the request work pool. Equality implies two pieces of work are the same
+/// kind of thing. The `str` should be human readable for logging, ie the language server protocol
+/// request message name or similar.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct WorkDescription(pub &'static str);
+
+impl fmt::Display for WorkDescription {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+lazy_static! {
+    /// Maximum total concurrent working tasks
+    static ref NUM_THREADS: usize = ::num_cpus::get();
+
+    /// Duration of work after which we should warn something is taking a long time
+    static ref WARN_TASK_DURATION: Duration = *DEFAULT_REQUEST_TIMEOUT * 5;
+
+    /// Current work descriptions active on the work pool
+    static ref WORK: Mutex<Vec<WorkDescription>> = Mutex::new(vec![]);
+
+    /// Thread pool for request execution allowing concurrent request processing.
+    static ref WORK_POOL: rayon::ThreadPool = rayon::ThreadPool::new(
+        rayon::Configuration::default()
+            .thread_name(|num| format!("request-worker-{}", num))
+            .num_threads(*NUM_THREADS)
+    ).unwrap();
+}
+
+/// Maximum concurrent working tasks of the same type
+/// Note: `2` allows a single task to run immediately after a similar task has timed out.
+/// Once multiple tasks have timed out but remain running we start refusing to start new ones.
+const MAX_SIMILAR_CONCURRENT_WORK: usize = 2;
+
+/// Runs work in a new thread on the `WORK_POOL` returning a result `Receiver`
+///
+/// Panicking work will receive `Err(RecvError)` / `Err(RecvTimeoutError::Disconnected)`
+///
+/// If too many tasks are already running the work will not be done and the receiver will
+/// immediately return `Err(RecvTimeoutError::Disconnected)`
+pub fn receive_from_thread<T, F>(work_fn: F, description: WorkDescription) -> mpsc::Receiver<T>
+where
+    T: Send + 'static,
+    F: FnOnce() -> T + Send + panic::UnwindSafe + 'static,
+{
+    let (sender, receiver) = mpsc::channel();
+
+    {
+        let mut work = WORK.lock().unwrap();
+        if work.len() >= *NUM_THREADS {
+            // there are already N ongoing tasks, that may or may not have timed out
+            // don't add yet more to the queue fail fast to allow the work pool to recover
+            warn!(
+                "Could not start `{}` as at work capacity, {:?} in progress",
+                description, *work,
+            );
+            return receiver;
+        }
+        if work.iter().filter(|desc| *desc == &description).count() >= MAX_SIMILAR_CONCURRENT_WORK {
+            // this type of work is already filling around half the work pool, so there's
+            // good reason to believe it may fill the entire pool => fail fast to allow
+            // other task-types to run
+            info!(
+                "Could not start `{}` as same work-type is filling half capacity, {:?} in progress",
+                description, *work,
+            );
+            return receiver;
+        }
+        work.push(description);
+    }
+
+    WORK_POOL.spawn(move || {
+        let start = Instant::now();
+
+        // panic details will be on stderr, otherwise ignore the work panic as it
+        // will already cause a mpsc disconnect-error & there isn't anything else to log
+        if let Ok(work_result) = panic::catch_unwind(work_fn) {
+            // an error here simply means the work took too long and the receiver has been dropped
+            let _ = sender.send(work_result);
+        }
+
+        let mut work = WORK.lock().unwrap();
+        if let Some(index) = work.iter().position(|desc| desc == &description) {
+            work.swap_remove(index);
+        }
+
+        let elapsed = start.elapsed();
+        if elapsed >= *WARN_TASK_DURATION {
+            let secs =
+                elapsed.as_secs() as f64 + f64::from(elapsed.subsec_nanos()) / 1_000_000_000_f64;
+            warn!("`{}` took {:.1}s", description, secs);
+        }
+    });
+    receiver
+}

src/actions/work_pool.rs

@@ -30,6 +30,7 @@ extern crate languageserver_types as ls_types;
 extern crate lazy_static;
 #[macro_use]
 extern crate log;
+extern crate num_cpus;
 extern crate racer;
 extern crate rayon;
 extern crate rls_analysis as analysis;

src/main.rs

@@ -9,6 +9,8 @@
 // except according to those terms.
 
 use super::requests::*;
+use actions::work_pool;
+use actions::work_pool::WorkDescription;
 use jsonrpc_core as jsonrpc;
 use server;
 use server::{Request, Response};
@@ -20,7 +22,7 @@ use std::thread;
 use std::time::Duration;
 
 lazy_static! {
-    static ref TIMEOUT: Duration = Duration::from_millis(::COMPILER_TIMEOUT);
+    pub static ref DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_millis(::COMPILER_TIMEOUT);
 }
 
 /// Macro enum `DispatchRequest` packing in various similar `Request` types
@@ -48,7 +50,7 @@ macro_rules! define_dispatch_request_enum {
                         let Request { id, params, received, .. } = req;
                         let timeout = $request_type::timeout();
 
-                        let receiver = receive_from_thread(move || {
+                        let receiver = work_pool::receive_from_thread(move || {
                             // checking timeout here can prevent starting expensive work that has
                             // already timed out due to previous long running requests
                             // Note: done here on the threadpool as pool scheduling may incur
@@ -59,7 +61,7 @@ macro_rules! define_dispatch_request_enum {
                             else {
                                 $request_type::handle(ctx, params)
                             }
-                        });
+                        }, WorkDescription($request_type::METHOD));
 
                         match receiver.recv_timeout(timeout)
                             .unwrap_or_else(|_| $request_type::fallback_response()) {
@@ -160,7 +162,7 @@ pub trait RequestAction: LSPRequest {
 
     /// Max duration this request should finish within, also see `fallback_response()`
     fn timeout() -> Duration {
-        *TIMEOUT
+        *DEFAULT_REQUEST_TIMEOUT
     }
 
     /// Returns a response used in timeout scenarios

src/server/dispatch.rs

@@ -28,7 +28,7 @@ use config::Config;
 pub use server::io::{MessageReader, Output};
 use server::io::{StdioMsgReader, StdioOutput};
 use server::dispatch::Dispatcher;
-pub use server::dispatch::{RequestAction, ResponseError};
+pub use server::dispatch::{RequestAction, ResponseError, DEFAULT_REQUEST_TIMEOUT};
 
 pub use ls_types::request::Shutdown as ShutdownRequest;
 pub use ls_types::request::Initialize as InitializeRequest;