From 7e7f42a581e2bfc05fb9b2a989cccc64fdf133f3 Mon Sep 17 00:00:00 2001
From: Richard Pringle <richard.pringle@avalabs.org>
Date: Thu, 28 Mar 2024 17:32:39 -0400
Subject: [PATCH] Handle return type properly

---
 tokio-macros/src/entry.rs | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs
index ed782ad38f6..544130473af 100644
--- a/tokio-macros/src/entry.rs
+++ b/tokio-macros/src/entry.rs
@@ -380,6 +380,23 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
         }
     };
 
+    let output_type = match &input.sig.output {
+        // For functions with no return value syn doesn't print anything,
+        // but that doesn't work as `Output` for our boxed `Future`, so
+        // default to `()` (the same type as the function output).
+        syn::ReturnType::Default => quote! { () },
+        syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
+    };
+
+    input.stmts.last_mut().map(|stmt| {
+        *stmt = quote! {
+            let _lst_stmt: #output_type = {
+                #stmt
+            };
+            _lst_stmt
+        };
+    });
+
     let body = input.body();
 
     // For test functions pin the body to the stack and use `Pin<&mut dyn
@@ -392,13 +409,6 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
     // We don't do this for the main function as it should only be used once so
     // there will be no benefit.
     let body = if is_test {
-        let output_type = match &input.sig.output {
-            // For functions with no return value syn doesn't print anything,
-            // but that doesn't work as `Output` for our boxed `Future`, so
-            // default to `()` (the same type as the function output).
-            syn::ReturnType::Default => quote! { () },
-            syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
-        };
         quote! {
             let body = async #body;
             #crate_path::pin!(body);