diff options
Diffstat (limited to 'tokio/src/runtime/thread_pool')
-rw-r--r-- | tokio/src/runtime/thread_pool/worker.rs | 27 |
1 files changed, 22 insertions, 5 deletions
diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs index 400e2a93..2213ec6c 100644 --- a/tokio/src/runtime/thread_pool/worker.rs +++ b/tokio/src/runtime/thread_pool/worker.rs @@ -177,21 +177,33 @@ cfg_blocking! { F: FnOnce() -> R, { // Try to steal the worker core back - struct Reset; + struct Reset(bool); impl Drop for Reset { fn drop(&mut self) { CURRENT.with(|maybe_cx| { + if !self.0 { + // We were not the ones to give away the core, + // so we do not get to restore it either. + // This is necessary so that with a nested + // block_in_place, the inner block_in_place + // does not restore the core. + return; + } + if let Some(cx) = maybe_cx { let core = cx.worker.core.take(); - *cx.core.borrow_mut() = core; + let mut cx_core = cx.core.borrow_mut(); + assert!(cx_core.is_none()); + *cx_core = core; } }); } } + let mut had_core = false; CURRENT.with(|maybe_cx| { - let cx = maybe_cx.expect("can call blocking only when running in a spawned task"); + let cx = maybe_cx.expect("can call blocking only when running in a spawned task on the multi-threaded runtime"); // Get the worker core. If none is set, then blocking is fine! let core = match cx.core.borrow_mut().take() { @@ -212,6 +224,7 @@ cfg_blocking! { // // First, move the core back into the worker's shared core slot. cx.worker.core.set(core); + had_core = true; // Next, clone the worker handle and send it to a new thread for // processing. @@ -222,9 +235,13 @@ cfg_blocking! { runtime::spawn_blocking(move || run(worker)); }); - let _reset = Reset; + let _reset = Reset(had_core); - f() + if had_core { + crate::runtime::enter::exit(f) + } else { + f() + } } } |