fix: allow for recursive block-on calls

Author: dignifiedquire

Cargo.toml

@@ -98,4 +98,7 @@ required-features = ["unstable"]
 
 [[example]]
 name = "surf-web"
-required-features = ["surf"]
+required-features = ["surf"]
+
+[patch.crates-io]
+smol = { git = "https://github.com/dignifiedquire/smol-1", branch = "feat/recursive-block-on" }

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);
+}