diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs index 0c9be1639..0b65606c2 100644 --- a/src/runtime/mod.rs +++ b/src/runtime/mod.rs @@ -209,11 +209,13 @@ struct Inner { pub fn run(future: F) where F: Future + Send + 'static, { - let mut runtime = Runtime::new().unwrap(); + // Check enter before creating a new Runtime... + let mut entered = enter().expect("nested tokio::run"); + let mut runtime = Runtime::new().expect("failed to start new Runtime"); runtime.spawn(future); - enter().expect("nested tokio::run") + entered .block_on(runtime.shutdown_on_idle()) - .unwrap(); + .expect("shutdown cannot error") } impl Runtime { @@ -362,9 +364,10 @@ impl Runtime { R: Send + 'static, E: Send + 'static, { + let mut entered = enter().expect("nested block_on"); let (tx, rx) = futures::sync::oneshot::channel(); self.spawn(future.then(move |r| tx.send(r).map_err(|_| unreachable!()))); - rx.wait().unwrap() + entered.block_on(rx).unwrap() } /// Run a future to completion on the Tokio runtime, then wait for all @@ -387,9 +390,16 @@ impl Runtime { R: Send + 'static, E: Send + 'static, { - let res = self.block_on(future); - self.shutdown_on_idle().wait().unwrap(); - res + let mut entered = enter().expect("nested block_on_all"); + let (tx, rx) = futures::sync::oneshot::channel(); + self.spawn(future.then(move |r| tx.send(r).map_err(|_| unreachable!()))); + let block = rx + .map_err(|_| unreachable!()) + .and_then(move |r| { + self.shutdown_on_idle() + .map(move |()| r) + }); + entered.block_on(block).unwrap() } /// Signals the runtime to shutdown once it becomes idle. diff --git a/tests/runtime.rs b/tests/runtime.rs index 66d10b951..a528ef1c5 100644 --- a/tests/runtime.rs +++ b/tests/runtime.rs @@ -391,14 +391,78 @@ mod from_block_on_all { } } -#[test] -fn run_in_run() { +mod nested_enter { + use super::*; + use tokio::runtime::current_thread; use std::panic; - tokio::run(lazy(|| { - panic::catch_unwind(|| { - tokio::run(lazy(|| { Ok::<(), ()>(()) })) - }).unwrap_err(); - Ok::<(), ()>(()) - })); + fn test(first: F1, nested: F2) + where + F1: Fn(Box + Send>) + Send + 'static, + F2: Fn(Box + Send>) + panic::UnwindSafe + Send + 'static, + { + let panicked = Arc::new(Mutex::new(false)); + let panicked2 = panicked.clone(); + + // Since this is testing panics in other threads, printing about panics + // is noisy and can give the impression that the test is ignoring panics. + // + // It *is* ignoring them, but on purpose. + let prev_hook = panic::take_hook(); + panic::set_hook(Box::new(|info| { + let s = info.to_string(); + if s.starts_with("panicked at 'nested ") + || s.starts_with("panicked at 'Multiple executors at once") + { + // expected, noop + } else { + println!("{}", s); + } + })); + + first(Box::new(lazy(move || { + panic::catch_unwind(move || { + nested(Box::new(lazy(|| { Ok::<(), ()>(()) }))) + }).expect_err("nested should panic"); + *panicked2.lock().unwrap() = true; + Ok::<(), ()>(()) + }))); + + panic::set_hook(prev_hook); + + assert!(*panicked.lock().unwrap(), "nested call should have panicked"); + } + + fn threadpool_new() -> Runtime { + Runtime::new().expect("rt new") + } + + #[test] + fn run_in_run() { + test(tokio::run, tokio::run); + } + + #[test] + fn threadpool_block_on_in_run() { + test(tokio::run, |fut| { + let mut rt = threadpool_new(); + rt.block_on(fut).unwrap(); + }); + } + + #[test] + fn threadpool_block_on_all_in_run() { + test(tokio::run, |fut| { + let rt = threadpool_new(); + rt.block_on_all(fut).unwrap(); + }); + } + + #[test] + fn current_thread_block_on_all_in_run() { + test(tokio::run, |fut| { + current_thread::block_on_all(fut).unwrap(); + }); + } } +