mirror of https://github.com/stjepang/smol
Merge pull request #147 from dignifiedquire/feat/recursive-block-on
feat: allow recursive block_on calls
This commit is contained in:
commit
ee25d12cac
|
@ -57,7 +57,7 @@ pub fn block_on<T>(future: impl Future<Output = T>) -> T {
|
||||||
|
|
||||||
CACHE.with(|cache| {
|
CACHE.with(|cache| {
|
||||||
// Panic if `block_on()` is called recursively.
|
// Panic if `block_on()` is called recursively.
|
||||||
let (parker, waker) = &mut *cache.try_borrow_mut().expect("recursive `block_on()`");
|
let (parker, waker) = &*cache.borrow();
|
||||||
|
|
||||||
// If enabled, set up tokio before execution begins.
|
// If enabled, set up tokio before execution begins.
|
||||||
context::enter(|| {
|
context::enter(|| {
|
||||||
|
|
|
@ -10,10 +10,27 @@ pub(crate) fn enter<T>(f: impl FnOnce() -> T) -> T {
|
||||||
#[cfg(feature = "tokio02")]
|
#[cfg(feature = "tokio02")]
|
||||||
{
|
{
|
||||||
use once_cell::sync::Lazy;
|
use once_cell::sync::Lazy;
|
||||||
|
use std::cell::Cell;
|
||||||
use tokio::runtime::Runtime;
|
use tokio::runtime::Runtime;
|
||||||
|
|
||||||
|
thread_local! {
|
||||||
|
/// The level of nested `enter` calls we are in, to ensure that the outer most always has a
|
||||||
|
/// runtime spawned.
|
||||||
|
static NESTING: Cell<usize> = Cell::new(0);
|
||||||
|
}
|
||||||
|
|
||||||
static RT: Lazy<Runtime> = Lazy::new(|| Runtime::new().expect("cannot initialize tokio"));
|
static RT: Lazy<Runtime> = Lazy::new(|| Runtime::new().expect("cannot initialize tokio"));
|
||||||
|
|
||||||
RT.enter(f)
|
NESTING.with(|nesting| {
|
||||||
|
let res = if nesting.get() == 0 {
|
||||||
|
nesting.replace(1);
|
||||||
|
RT.enter(f)
|
||||||
|
} else {
|
||||||
|
nesting.replace(nesting.get() + 1);
|
||||||
|
f()
|
||||||
|
};
|
||||||
|
nesting.replace(nesting.get() - 1);
|
||||||
|
res
|
||||||
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
@ -0,0 +1,37 @@
|
||||||
|
use futures_util::future;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn smoke() {
|
||||||
|
std::thread::spawn(|| {
|
||||||
|
smol::run(future::pending::<()>());
|
||||||
|
});
|
||||||
|
let res = smol::block_on(async { 1 + 2 });
|
||||||
|
assert_eq!(res, 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[should_panic = "boom"]
|
||||||
|
fn panic() {
|
||||||
|
std::thread::spawn(|| {
|
||||||
|
smol::run(future::pending::<()>());
|
||||||
|
});
|
||||||
|
smol::block_on(async {
|
||||||
|
// This panic should get propagated into the parent thread.
|
||||||
|
panic!("boom");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn nested_block_on() {
|
||||||
|
std::thread::spawn(|| {
|
||||||
|
smol::run(future::pending::<()>());
|
||||||
|
});
|
||||||
|
|
||||||
|
let x = smol::block_on(async {
|
||||||
|
let a = smol::block_on(async { smol::block_on(async { future::ready(3).await }) });
|
||||||
|
let b = smol::block_on(async { smol::block_on(async { future::ready(2).await }) });
|
||||||
|
a + b
|
||||||
|
});
|
||||||
|
|
||||||
|
assert_eq!(x, 3 + 2);
|
||||||
|
}
|
Loading…
Reference in New Issue