fix: allow for recursive block-on calls

Cargo.toml

@@ -75,7 +75,7 @@ futures-timer = { version = "3.0.2", optional = true }
 surf = { version = "1.0.3", optional = true }
 
 [target.'cfg(not(target_os = "unknown"))'.dependencies]
-smol = { version = "0.1.10", optional = true }
+smol = { version = "0.1.11", optional = true }
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 futures-timer = { version = "3.0.2", optional = true, features = ["wasm-bindgen"] }
@@ -103,3 +103,4 @@ required-features = ["unstable"]
 [[example]]
 name = "surf-web"
 required-features = ["surf"]
+

src/task/builder.rs

@@ -1,3 +1,4 @@
+use std::cell::Cell;
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -8,6 +9,11 @@ use pin_project_lite::pin_project;
 use crate::io;
 use crate::task::{JoinHandle, Task, TaskLocalsWrapper};
 
+#[cfg(not(target_os = "unknown"))]
+thread_local! {
+    static IS_CURRENT_BLOCKING: Cell<usize> = Cell::new(0);
+}
+
 /// Task builder that configures the settings of a new task.
 #[derive(Debug, Default)]
 pub struct Builder {
@@ -151,7 +157,36 @@ impl Builder {
         });
 
         // Run the future as a task.
-        unsafe { TaskLocalsWrapper::set_current(&wrapped.tag, || smol::run(wrapped)) }
+        IS_CURRENT_BLOCKING.with(|is_current_blocking| {
+            let count = is_current_blocking.get();
+            let res = if count == 0 {
+                // increase the count
+                is_current_blocking.replace(1);
+
+                // The first call should use run.
+                unsafe {
+                    TaskLocalsWrapper::set_current(&wrapped.tag, || {
+                        let res = smol::run(wrapped);
+                        is_current_blocking.replace(is_current_blocking.get() - 1);
+                        res
+                    })
+                }
+            } else {
+                // increase the count
+                is_current_blocking.replace(is_current_blocking.get() + 1);
+
+                // Subsequent calls should be using `block_on`.
+                unsafe {
+                    TaskLocalsWrapper::set_current(&wrapped.tag, || {
+                        let res = smol::block_on(wrapped);
+                        is_current_blocking.replace(is_current_blocking.get() - 1);
+                        res
+                    })
+                }
+            };
+
+            res
+        })
     }
 }
 

tests/block_on.rs

@@ -16,3 +16,62 @@ fn panic() {
         panic!("boom");
     });
 }
+
+#[cfg(feature = "unstable")]
+#[test]
+fn nested_block_on_local() {
+    let x = task::block_on(async {
+        let a =
+            task::block_on(async { task::block_on(async { async_std::future::ready(3).await }) });
+        let b = task::spawn_local(async {
+            task::block_on(async { async_std::future::ready(2).await })
+        })
+        .await;
+        let c =
+            task::block_on(async { task::block_on(async { async_std::future::ready(1).await }) });
+        a + b + c
+    });
+
+    assert_eq!(x, 3 + 2 + 1);
+
+    let y = task::block_on(async {
+        let a =
+            task::block_on(async { task::block_on(async { async_std::future::ready(3).await }) });
+        let b = task::spawn_local(async {
+            task::block_on(async { async_std::future::ready(2).await })
+        })
+        .await;
+        let c =
+            task::block_on(async { task::block_on(async { async_std::future::ready(1).await }) });
+        a + b + c
+    });
+
+    assert_eq!(y, 3 + 2 + 1);
+}
+
+#[test]
+fn nested_block_on() {
+    let x = task::block_on(async {
+        let a =
+            task::block_on(async { task::block_on(async { async_std::future::ready(3).await }) });
+        let b =
+            task::block_on(async { task::block_on(async { async_std::future::ready(2).await }) });
+        let c =
+            task::block_on(async { task::block_on(async { async_std::future::ready(1).await }) });
+        a + b + c
+    });
+
+    assert_eq!(x, 3 + 2 + 1);
+
+    let y = task::block_on(async {
+        let a =
+            task::block_on(async { task::block_on(async { async_std::future::ready(3).await }) });
+        let b =
+            task::block_on(async { task::block_on(async { async_std::future::ready(2).await }) });
+        let c =
+            task::block_on(async { task::block_on(async { async_std::future::ready(1).await }) });
+        a + b + c
+    });
+
+    assert_eq!(y, 3 + 2 + 1);
+}