1 files changed, 18 insertions, 18 deletions

diff --git a/cli/tokio_util.rs b/cli/tokio_util.rs
index 5b975ff40..e1f8587c3 100644
--- a/cli/tokio_util.rs
+++ b/cli/tokio_util.rs
@@ -8,25 +8,28 @@ use std::mem;
 use std::net::SocketAddr;
 use tokio;
 use tokio::net::TcpStream;
+use tokio::runtime;
+
+pub fn create_threadpool_runtime() -> tokio::runtime::Runtime {
+  // This code can be simplified once the following PR is landed and
+  // released: https://github.com/tokio-rs/tokio/pull/1055
+  use tokio_threadpool::Builder as ThreadPoolBuilder;
+  let mut threadpool_builder = ThreadPoolBuilder::new();
+  threadpool_builder.panic_handler(|err| std::panic::resume_unwind(err));
+  #[allow(deprecated)]
+  runtime::Builder::new()
+    .threadpool_builder(threadpool_builder)
+    .build()
+    .unwrap()
+}
 
 pub fn run<F>(future: F)
 where
   F: Future<Item = (), Error = ()> + Send + 'static,
 {
-  abort_on_panic();
   // tokio::runtime::current_thread::run(future)
-  tokio::run(future)
-}
-
-// Tokio swallows panics. In order to actually crash when we panic, we
-// have to set this custom hook.
-// https://github.com/tokio-rs/tokio/issues/495
-// https://github.com/tokio-rs/tokio/issues/209
-pub fn abort_on_panic() {
-  std::panic::set_hook(Box::new(|panic_info| {
-    eprintln!("{}", panic_info.to_string());
-    std::process::abort();
-  }));
+  let rt = create_threadpool_runtime();
+  rt.block_on_all(future).unwrap();
 }
 
 pub fn block_on<F, R, E>(future: F) -> Result<R, E>
@@ -48,13 +51,10 @@ pub fn init<F>(f: F)
 where
   F: FnOnce(),
 {
-  let rt = tokio::runtime::Runtime::new().unwrap();
+  let rt = create_threadpool_runtime();
   let mut executor = rt.executor();
   let mut enter = tokio_executor::enter().expect("Multiple executors at once");
-  tokio_executor::with_default(&mut executor, &mut enter, move |_enter| {
-    abort_on_panic();
-    f()
-  });
+  tokio_executor::with_default(&mut executor, &mut enter, move |_enter| f());
 }
 
 #[derive(Debug)]