Update projects/RabbitMQ.Client/client/TaskExtensions.cs

Co-authored-by: Paulo Morgado <470455+paulomorgado@users.noreply.github.com>
rabbitmq · Nov 19, 2023 · 90b72b9 · 90b72b9
1 parent 0d13a2c
commit 90b72b9
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/projects/RabbitMQ.Client/client/TaskExtensions.cs b/projects/RabbitMQ.Client/client/TaskExtensions.cs
@@ -75,15 +75,21 @@ static async Task TimeoutAfterImplement(Task task, TimeSpan timeout)
 #endif
         }
 
-        public static async ValueTask TimeoutAfter(this ValueTask task, TimeSpan timeout)
+        public static ValueTask TimeoutAfter(this ValueTask task, TimeSpan timeout)
         {
-            if (false == task.IsCompletedSuccessfully)
+            if (task.IsCompletedSuccessfully)
             {
-                Task actualTask = task.AsTask();
+                return task;
+            }
+
 #if NET6_0_OR_GREATER
-                await actualTask.WaitAsync(timeout)
-                    .ConfigureAwait(false);
+            return actualTask.WaitAsync(timeout);
 #else
+            return TimeoutAfterImplement(task, timeout);
+
+            assync static ValueTask TimeoutAfterImpl(ValueTask task, TimeSpan timeout)
+            {
+                Task actualTask = task.AsTask();
                 if (actualTask == await Task.WhenAny(actualTask, Task.Delay(timeout)).ConfigureAwait(false))
                 {
                     await actualTask.ConfigureAwait(false);
@@ -98,6 +104,7 @@ await actualTask.WaitAsync(timeout)
                         scheduler: TaskScheduler.Default);
                     throw new TimeoutException();
                 }
+            }
 #endif
             }
         }