diff --git a/tokio-threadpool/src/builder.rs b/tokio-threadpool/src/builder.rs index 1cb2ec520..a51e8e520 100644 --- a/tokio-threadpool/src/builder.rs +++ b/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(&mut self, f: F) -> &mut Self + where + F: Fn(Box) + 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 diff --git a/tokio-threadpool/src/config.rs b/tokio-threadpool/src/config.rs index e51efc2ba..5cbd0156e 100644 --- a/tokio-threadpool/src/config.rs +++ b/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, pub after_start: Option>, pub before_stop: Option>, + pub panic_handler: Option) + Send + Sync>>, } /// Max number of workers that can be part of a pool. This is the most that can diff --git a/tokio-threadpool/src/task/mod.rs b/tokio-threadpool/src/task/mod.rs index 2dd00c591..d9bdb07a4 100644 --- a/tokio-threadpool/src/task/mod.rs +++ b/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)) => { diff --git a/tokio-threadpool/tests/threadpool.rs b/tokio-threadpool/tests/threadpool.rs index dd482625d..507fd03ea 100644 --- a/tokio-threadpool/tests/threadpool.rs +++ b/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;