threadpool: add panic_handler (#1052)

Author: ry

tokio-threadpool/src/builder.rs

@@ -6,6 +6,7 @@ use shutdown::ShutdownTrigger;
 use thread_pool::ThreadPool;
 use worker::{self, Worker, WorkerId};
 
+use std::any::Any;
 use std::cmp::max;
 use std::error::Error;
 use std::fmt;
@@ -106,6 +107,7 @@ impl Builder {
                 around_worker: None,
                 after_start: None,
                 before_stop: None,
+                panic_handler: None,
             },
             new_park,
         }
@@ -199,6 +201,34 @@ impl Builder {
         self
     }
 
+    /// Sets a callback to be triggered when a panic during a future bubbles up
+    /// to Tokio. By default Tokio catches these panics, and they will be
+    /// ignored. The parameter passed to this callback is the same error value
+    /// returned from std::panic::catch_unwind(). To abort the process on
+    /// panics, use std::panic::resume_unwind() in this callback as shown
+    /// below.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # extern crate tokio_threadpool;
+    /// # extern crate futures;
+    /// # use tokio_threadpool::Builder;
+    ///
+    /// # pub fn main() {
+    /// let thread_pool = Builder::new()
+    ///     .panic_handler(|err| std::panic::resume_unwind(err))
+    ///     .build();
+    /// # }
+    /// ```
+    pub fn panic_handler<F>(&mut self, f: F) -> &mut Self
+    where
+        F: Fn(Box<Any + Send>) + Send + Sync + 'static,
+    {
+        self.config.panic_handler = Some(Arc::new(f));
+        self
+    }
+
     /// Set name prefix of threads spawned by the scheduler
     ///
     /// Thread name prefix is used for generating thread names. For example, if

tokio-threadpool/src/config.rs

@@ -1,5 +1,6 @@
 use callback::Callback;
 
+use std::any::Any;
 use std::fmt;
 use std::sync::Arc;
 use std::time::Duration;
@@ -14,6 +15,7 @@ pub(crate) struct Config {
     pub around_worker: Option<Callback>,
     pub after_start: Option<Arc<Fn() + Send + Sync>>,
     pub before_stop: Option<Arc<Fn() + Send + Sync>>,
+    pub panic_handler: Option<Arc<Fn(Box<Any + Send>) + Send + Sync>>,
 }
 
 /// Max number of workers that can be part of a pool. This is the most that can

tokio-threadpool/src/task/mod.rs

@@ -165,6 +165,12 @@ impl Task {
                 // Transition to the completed state
                 self.state.store(State::Complete.into(), Release);
 
+                if let Err(panic_err) = res {
+                    if let Some(ref f) = unpark.pool.config.panic_handler {
+                        f(panic_err);
+                    }
+                }
+
                 Run::Complete
             }
             Ok(Ok(Async::NotReady)) => {

tokio-threadpool/tests/threadpool.rs

@@ -399,6 +399,23 @@ fn panic_in_task() {
     pool.shutdown_on_idle().wait().unwrap();
 }
 
+#[test]
+fn count_panics() {
+    let counter = Arc::new(AtomicUsize::new(0));
+    let counter_ = counter.clone();
+    let pool = tokio_threadpool::Builder::new()
+        .panic_handler(move |_err| {
+            // We caught a panic.
+            counter_.fetch_add(1, Relaxed);
+        })
+        .build();
+    // Spawn a future that will panic.
+    pool.spawn(lazy(|| -> Result<(), ()> { panic!() }));
+    pool.shutdown_on_idle().wait().unwrap();
+    let counter = counter.load(Relaxed);
+    assert_eq!(counter, 1);
+}
+
 #[test]
 fn multi_threadpool() {
     use futures::sync::oneshot;