Fix workload continuer (#3079)

* Fix workload continuer for synchronous continuations.

* Fix test.

Author: timcassell

src/BenchmarkDotNet/Code/DeclarationsProvider.cs

@@ -1,4 +1,4 @@
-using BenchmarkDotNet.Attributes;
+using BenchmarkDotNet.Attributes;
 using BenchmarkDotNet.Engines;
 using BenchmarkDotNet.Environments;
 using BenchmarkDotNet.Extensions;
@@ -177,7 +177,7 @@ internal class AsyncDeclarationsProvider(BenchmarkCase benchmark) : Declarations
     {
         public override string[] GetExtraFields() =>
         [
-            $"public {typeof(WorkloadContinuerAndValueTaskSource).GetCorrectCSharpTypeName()} workloadContinuerAndValueTaskSource;",
+            $"public {typeof(WorkloadValueTaskSource).GetCorrectCSharpTypeName()} workloadContinuerAndValueTaskSource;",
             $"public {typeof(IClock).GetCorrectCSharpTypeName()} clock;",
             "public long invokeCount;"
         ];
@@ -224,7 +224,7 @@ protected override SmartStringBuilder ReplaceCore(SmartStringBuilder smartString
                         this.__fieldsContainer.clock = clock;
                         if (this.__fieldsContainer.workloadContinuerAndValueTaskSource == null)
                         {
-                            this.__fieldsContainer.workloadContinuerAndValueTaskSource = new {{typeof(WorkloadContinuerAndValueTaskSource).GetCorrectCSharpTypeName()}}();
+                            this.__fieldsContainer.workloadContinuerAndValueTaskSource = new {{typeof(WorkloadValueTaskSource).GetCorrectCSharpTypeName()}}();
                             this.__StartWorkload();
                         }
                         return this.__fieldsContainer.workloadContinuerAndValueTaskSource.Continue();
@@ -240,14 +240,12 @@ private async void __StartWorkload()
                     {
                         try
                         {
+                            if (await this.__fieldsContainer.workloadContinuerAndValueTaskSource.GetIsComplete())
+                            {
+                                {{finalReturn}}
+                            }
                             while (true)
                             {
-                                await this.__fieldsContainer.workloadContinuerAndValueTaskSource;
-                                if (this.__fieldsContainer.workloadContinuerAndValueTaskSource.IsCompleted)
-                                {
-                                    {{finalReturn}}
-                                }
-
                                 {{typeof(StartedClock).GetCorrectCSharpTypeName()}} startedClock = {{typeof(ClockExtensions).GetCorrectCSharpTypeName()}}.Start(this.__fieldsContainer.clock);
                                 while (--this.__fieldsContainer.invokeCount >= 0)
                                 {
@@ -256,7 +254,10 @@ private async void __StartWorkload()
                                     unsafe { awaitable = {{workloadMethodCall}} }
                                     await awaitable;
                                 }
-                                this.__fieldsContainer.workloadContinuerAndValueTaskSource.SetResult(startedClock.GetElapsed());
+                                if (await this.__fieldsContainer.workloadContinuerAndValueTaskSource.SetResultAndGetIsComplete(startedClock.GetElapsed()))
+                                {
+                                    {{finalReturn}}
+                                }
                             }
                         }
                         catch (global::System.Exception e)

src/BenchmarkDotNet/Engines/WorkloadContinuerAndValueTaskSource.cs

@@ -0,0 +1,57 @@
+using JetBrains.Annotations;
+using Perfolizer.Horology;
+using System.ComponentModel;
+using System.Threading.Tasks.Sources;
+
+namespace BenchmarkDotNet.Engines;
+
+// This is used to prevent allocating a new async state machine on every benchmark iteration.
+[UsedImplicitly]
+[EditorBrowsable(EditorBrowsableState.Never)]
+public sealed class WorkloadValueTaskSource : IValueTaskSource<ClockSpan>, IValueTaskSource<bool>
+{
+    private ManualResetValueTaskSourceCore<bool> continuerSource;
+    private ManualResetValueTaskSourceCore<ClockSpan> clockSpanSource;
+
+    public ValueTask<ClockSpan> Continue()
+    {
+        clockSpanSource.Reset();
+        continuerSource.SetResult(false);
+        return new(this, clockSpanSource.Version);
+    }
+
+    public ValueTask<bool> GetIsComplete()
+        => new(this, continuerSource.Version);
+
+    public void Complete()
+        => continuerSource.SetResult(true);
+
+    public ValueTask<bool> SetResultAndGetIsComplete(ClockSpan result)
+    {
+        continuerSource.Reset();
+        clockSpanSource.SetResult(result);
+        return GetIsComplete();
+    }
+
+    public void SetException(Exception exception)
+        => clockSpanSource.SetException(exception);
+
+    ValueTaskSourceStatus IValueTaskSource<bool>.GetStatus(short token)
+        => continuerSource.GetStatus(token);
+
+    void IValueTaskSource<bool>.OnCompleted(Action<object?> continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags)
+        // Strip UseSchedulingContext to ensure Continue() and Complete() synchronously continue __WorkloadCore.
+        => continuerSource.OnCompleted(continuation, state, token, flags & ~ValueTaskSourceOnCompletedFlags.UseSchedulingContext);
+
+    bool IValueTaskSource<bool>.GetResult(short token)
+        => continuerSource.GetResult(token);
+
+    ClockSpan IValueTaskSource<ClockSpan>.GetResult(short token)
+        => clockSpanSource.GetResult(token);
+
+    ValueTaskSourceStatus IValueTaskSource<ClockSpan>.GetStatus(short token)
+        => clockSpanSource.GetStatus(token);
+
+    void IValueTaskSource<ClockSpan>.OnCompleted(Action<object?> continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags)
+        => clockSpanSource.OnCompleted(continuation, state, token, flags);
+}