From 609d94ced3f5f986625c986681d362a6d1c8da1b Mon Sep 17 00:00:00 2001
From: Ruihang Xia <waynestxia@gmail.com>
Date: Tue, 28 Sep 2021 23:43:29 +0800
Subject: [PATCH] support async return

---
 Cargo.toml     |  6 ++++++
 src/lib.rs     |  5 +++++
 tests/tests.rs | 18 ++++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/Cargo.toml b/Cargo.toml
index fa097a6..b6c15f8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,6 +18,12 @@ log = { version = "0.4", features = ["std"] }
 lazy_static = "1.2"
 rand = "0.8"
 
+[dev-dependencies]
+tokio = { version = "1.12.0", default-features = false, features = [
+    "rt",
+    "macros",
+] }
+
 [features]
 failpoints = []
 
diff --git a/src/lib.rs b/src/lib.rs
index 7f2bac1..069dfbf 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -797,6 +797,11 @@ macro_rules! fail_point {
             panic!("Return is not supported for the fail point \"{}\"", $name);
         });
     }};
+    ($name:expr, | $arg:ident  $(: $t:ty )? | async $block:tt ) => {{
+        if let Some(res) = $crate::eval($name, {| $arg $(: $t)?| async $block}) {
+            return res.await;
+        }
+    }};
     ($name:expr, $e:expr) => {{
         if let Some(res) = $crate::eval($name, $e) {
             return res;
diff --git a/tests/tests.rs b/tests/tests.rs
index fecb53c..f14d2ba 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -36,6 +36,24 @@ fn test_return() {
     assert_eq!(f(), 2);
 }
 
+#[tokio::test]
+#[cfg_attr(not(feature = "failpoints"), ignore)]
+async fn test_async_return() {
+    async fn async_fn() -> usize {
+        fail_point!("async_return", |s: Option<String>| async {
+            (async {}).await;
+            s.map_or(2, |s| s.parse().unwrap())
+        });
+        0
+    }
+
+    fail::cfg("async_return", "return(1000)").unwrap();
+    assert_eq!(async_fn().await, 1000);
+
+    fail::cfg("async_return", "return").unwrap();
+    assert_eq!(async_fn().await, 2);
+}
+
 #[test]
 #[cfg_attr(not(feature = "failpoints"), ignore)]
 fn test_sleep() {