support async return

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

Author: waynexia

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 = []
 

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;

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() {