WIP

Refactored delegates to pass in IClock and return ValueTask<ClockSpan>.

TODO:
Fix compile errors from BenchmarkType.txt in NetFramework and Mono.
Update InProcessEmitBuilder for the new behavior.

src/BenchmarkDotNet/Code/CodeGenerator.cs

@@ -49,6 +49,7 @@ internal static string Generate(BuildPartition buildPartition)
                     .Replace("$WorkloadMethodReturnType$", provider.WorkloadMethodReturnTypeName)
                     .Replace("$WorkloadMethodReturnTypeModifiers$", provider.WorkloadMethodReturnTypeModifiers)
                     .Replace("$OverheadMethodReturnTypeName$", provider.OverheadMethodReturnTypeName)
+                    .Replace("$AwaiterTypeName$", provider.AwaiterTypeName)
                     .Replace("$GlobalSetupMethodName$", provider.GlobalSetupMethodName)
                     .Replace("$GlobalCleanupMethodName$", provider.GlobalCleanupMethodName)
                     .Replace("$IterationSetupMethodName$", provider.IterationSetupMethodName)
@@ -155,15 +156,12 @@ private static DeclarationsProvider GetDeclarationsProvider(Descriptor descripto
         {
             var method = descriptor.WorkloadMethod;
 
-            if (method.ReturnType == typeof(Task) || method.ReturnType == typeof(ValueTask))
-            {
-                return new TaskDeclarationsProvider(descriptor);
-            }
-            if (method.ReturnType.GetTypeInfo().IsGenericType
-                && (method.ReturnType.GetTypeInfo().GetGenericTypeDefinition() == typeof(Task<>)
+            if (method.ReturnType == typeof(Task) || method.ReturnType == typeof(ValueTask)
+                || method.ReturnType.GetTypeInfo().IsGenericType
+                    && (method.ReturnType.GetTypeInfo().GetGenericTypeDefinition() == typeof(Task<>)
                     || method.ReturnType.GetTypeInfo().GetGenericTypeDefinition() == typeof(ValueTask<>)))
             {
-                return new GenericTaskDeclarationsProvider(descriptor);
+                return new TaskDeclarationsProvider(descriptor);
             }
 
             if (method.ReturnType == typeof(void))

src/BenchmarkDotNet/Code/DeclarationsProvider.cs

@@ -1,5 +1,4 @@
 using System;
-using System.Linq;
 using System.Reflection;
 using System.Threading.Tasks;
 using BenchmarkDotNet.Engines;
@@ -11,9 +10,6 @@ namespace BenchmarkDotNet.Code
 {
     internal abstract class DeclarationsProvider
     {
-        // "GlobalSetup" or "GlobalCleanup" methods are optional, so default to an empty delegate, so there is always something that can be invoked
-        private const string EmptyAction = "() => { }";
-
         protected readonly Descriptor Descriptor;
 
         internal DeclarationsProvider(Descriptor descriptor) => Descriptor = descriptor;
@@ -26,9 +22,9 @@ internal abstract class DeclarationsProvider
 
         public string GlobalCleanupMethodName => GetMethodName(Descriptor.GlobalCleanupMethod);
 
-        public string IterationSetupMethodName => Descriptor.IterationSetupMethod?.Name ?? EmptyAction;
+        public string IterationSetupMethodName => GetMethodName(Descriptor.IterationSetupMethod);
 
-        public string IterationCleanupMethodName => Descriptor.IterationCleanupMethod?.Name ?? EmptyAction;
+        public string IterationCleanupMethodName => GetMethodName(Descriptor.IterationCleanupMethod);
 
         public abstract string ReturnsDefinition { get; }
 
@@ -48,13 +44,18 @@ internal abstract class DeclarationsProvider
 
         public string OverheadMethodReturnTypeName => OverheadMethodReturnType.GetCorrectCSharpTypeName();
 
+        public virtual string AwaiterTypeName => string.Empty;
+
+        public virtual void OverrideUnrollFactor(BenchmarkCase benchmarkCase) { }
+
         public abstract string OverheadImplementation { get; }
 
         private string GetMethodName(MethodInfo method)
         {
+            // "Setup" or "Cleanup" methods are optional, so default to a simple delegate, so there is always something that can be invoked
             if (method == null)
             {
-                return EmptyAction;
+                return "() => new System.Threading.Tasks.ValueTask()";
             }
 
             if (method.ReturnType == typeof(Task) ||
@@ -63,10 +64,10 @@ private string GetMethodName(MethodInfo method)
                     (method.ReturnType.GetGenericTypeDefinition() == typeof(Task<>) ||
                      method.ReturnType.GetGenericTypeDefinition() == typeof(ValueTask<>))))
             {
-                return $"() => awaitHelper.GetResult({method.Name}())";
+                return $"() => BenchmarkDotNet.Helpers.AwaitHelper.ToValueTaskVoid({method.Name}())";
             }
 
-            return method.Name;
+            return $"() => {{ {method.Name}(); return new System.Threading.Tasks.ValueTask(); }}";
         }
     }
 
@@ -145,30 +146,18 @@ internal class ByReadOnlyRefDeclarationsProvider : ByRefDeclarationsProvider
         public override string WorkloadMethodReturnTypeModifiers => "ref readonly";
     }
 
-    internal class TaskDeclarationsProvider : VoidDeclarationsProvider
+    internal class TaskDeclarationsProvider : DeclarationsProvider
     {
         public TaskDeclarationsProvider(Descriptor descriptor) : base(descriptor) { }
 
-        public override string WorkloadMethodDelegate(string passArguments)
-            => $"({passArguments}) => {{ awaitHelper.GetResult({Descriptor.WorkloadMethod.Name}({passArguments})); }}";
-
-        public override string GetWorkloadMethodCall(string passArguments) => $"awaitHelper.GetResult({Descriptor.WorkloadMethod.Name}({passArguments}))";
+        public override string ExtraDefines => "#define RETURNS_AWAITABLE";
 
-        protected override Type WorkloadMethodReturnType => typeof(void);
-    }
-
-    /// <summary>
-    /// declarations provider for <see cref="Task{TResult}" /> and <see cref="ValueTask{TResult}" />
-    /// </summary>
-    internal class GenericTaskDeclarationsProvider : NonVoidDeclarationsProvider
-    {
-        public GenericTaskDeclarationsProvider(Descriptor descriptor) : base(descriptor) { }
+        public override string AwaiterTypeName => WorkloadMethodReturnType.GetMethod(nameof(Task.GetAwaiter), BindingFlags.Public | BindingFlags.Instance).ReturnType.GetCorrectCSharpTypeName();
 
-        protected override Type WorkloadMethodReturnType => Descriptor.WorkloadMethod.ReturnType.GetTypeInfo().GetGenericArguments().Single();
+        public override string OverheadImplementation => $"return default({OverheadMethodReturnType.GetCorrectCSharpTypeName()});";
 
-        public override string WorkloadMethodDelegate(string passArguments)
-            => $"({passArguments}) => {{ return awaitHelper.GetResult({Descriptor.WorkloadMethod.Name}({passArguments})); }}";
+        protected override Type OverheadMethodReturnType => WorkloadMethodReturnType;
 
-        public override string GetWorkloadMethodCall(string passArguments) => $"awaitHelper.GetResult({Descriptor.WorkloadMethod.Name}({passArguments}))";
+        public override void OverrideUnrollFactor(BenchmarkCase benchmarkCase) => benchmarkCase.ForceUnrollFactorForAsync();
     }
 }

src/BenchmarkDotNet/Engines/Engine.cs

@@ -4,6 +4,7 @@
 using System.Globalization;
 using System.Linq;
 using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
 using BenchmarkDotNet.Characteristics;
 using BenchmarkDotNet.Jobs;
 using BenchmarkDotNet.Portability;
@@ -19,17 +20,17 @@ public class Engine : IEngine
         public const int MinInvokeCount = 4;
 
         [PublicAPI] public IHost Host { get; }
-        [PublicAPI] public Action<long> WorkloadAction { get; }
+        [PublicAPI] public Func<long, IClock, ValueTask<ClockSpan>> WorkloadAction { get; }
         [PublicAPI] public Action Dummy1Action { get; }
         [PublicAPI] public Action Dummy2Action { get; }
         [PublicAPI] public Action Dummy3Action { get; }
-        [PublicAPI] public Action<long> OverheadAction { get; }
+        [PublicAPI] public Func<long, IClock, ValueTask<ClockSpan>> OverheadAction { get; }
         [PublicAPI] public Job TargetJob { get; }
         [PublicAPI] public long OperationsPerInvoke { get; }
-        [PublicAPI] public Action GlobalSetupAction { get; }
-        [PublicAPI] public Action GlobalCleanupAction { get; }
-        [PublicAPI] public Action IterationSetupAction { get; }
-        [PublicAPI] public Action IterationCleanupAction { get; }
+        [PublicAPI] public Func<ValueTask> GlobalSetupAction { get; }
+        [PublicAPI] public Func<ValueTask> GlobalCleanupAction { get; }
+        [PublicAPI] public Func<ValueTask> IterationSetupAction { get; }
+        [PublicAPI] public Func<ValueTask> IterationCleanupAction { get; }
         [PublicAPI] public IResolver Resolver { get; }
         [PublicAPI] public CultureInfo CultureInfo { get; }
         [PublicAPI] public string BenchmarkName { get; }
@@ -46,13 +47,14 @@ public class Engine : IEngine
         private readonly EngineActualStage actualStage;
         private readonly bool includeExtraStats;
         private readonly Random random;
+        private readonly Helpers.AwaitHelper awaitHelper;
 
         internal Engine(
             IHost host,
             IResolver resolver,
-            Action dummy1Action, Action dummy2Action, Action dummy3Action, Action<long> overheadAction, Action<long> workloadAction, Job targetJob,
-            Action globalSetupAction, Action globalCleanupAction, Action iterationSetupAction, Action iterationCleanupAction, long operationsPerInvoke,
-            bool includeExtraStats, string benchmarkName)
+            Action dummy1Action, Action dummy2Action, Action dummy3Action, Func<long, IClock, ValueTask<ClockSpan>> overheadAction, Func<long, IClock, ValueTask<ClockSpan>> workloadAction,
+            Job targetJob, Func<ValueTask> globalSetupAction, Func<ValueTask> globalCleanupAction, Func<ValueTask> iterationSetupAction, Func<ValueTask> iterationCleanupAction,
+            long operationsPerInvoke, bool includeExtraStats, string benchmarkName)
         {
 
             Host = host;
@@ -84,13 +86,14 @@ public class Engine : IEngine
             actualStage = new EngineActualStage(this);
 
             random = new Random(12345); // we are using constant seed to try to get repeatable results
+            awaitHelper = new Helpers.AwaitHelper();
         }
 
         public void Dispose()
         {
             try
             {
-                GlobalCleanupAction?.Invoke();
+                awaitHelper.GetResult(GlobalCleanupAction.Invoke());
             }
             catch (Exception e)
             {
@@ -165,9 +168,8 @@ public Measurement RunIteration(IterationData data)
             Span<byte> stackMemory = randomizeMemory ? stackalloc byte[random.Next(32)] : Span<byte>.Empty;
 
             // Measure
-            var clock = Clock.Start();
-            action(invokeCount / unrollFactor);
-            var clockSpan = clock.GetElapsed();
+            var op = action(invokeCount / unrollFactor, Clock);
+            var clockSpan = awaitHelper.GetResult(op);
 
             if (EngineEventSource.Log.IsEnabled())
                 EngineEventSource.Log.IterationStop(data.IterationMode, data.IterationStage, totalOperations);
@@ -201,7 +203,8 @@ public Measurement RunIteration(IterationData data)
             var initialThreadingStats = ThreadingStats.ReadInitial(); // this method might allocate
             var initialGcStats = GcStats.ReadInitial();
 
-            WorkloadAction(data.InvokeCount / data.UnrollFactor);
+            var op = WorkloadAction(data.InvokeCount / data.UnrollFactor, Clock);
+            awaitHelper.GetResult(op);
 
             var finalGcStats = GcStats.ReadFinal();
             var finalThreadingStats = ThreadingStats.ReadFinal();

src/BenchmarkDotNet/Engines/EngineFactory.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading.Tasks;
 using BenchmarkDotNet.Jobs;
 using Perfolizer.Horology;
 
@@ -109,7 +110,7 @@ private static Engine CreateSingleActionEngine(EngineParameters engineParameters
                 engineParameters.OverheadActionNoUnroll,
                 engineParameters.WorkloadActionNoUnroll);
 
-        private static Engine CreateEngine(EngineParameters engineParameters, Job job, Action<long> idle, Action<long> main)
+        private static Engine CreateEngine(EngineParameters engineParameters, Job job, Func<long, IClock, ValueTask<ClockSpan>> idle, Func<long, IClock, ValueTask<ClockSpan>> main)
             => new Engine(
                 engineParameters.Host,
                 EngineParameters.DefaultResolver,

src/BenchmarkDotNet/Engines/EngineParameters.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading.Tasks;
 using BenchmarkDotNet.Characteristics;
 using BenchmarkDotNet.Jobs;
 using BenchmarkDotNet.Running;
@@ -12,19 +13,19 @@ public class EngineParameters
         public static readonly IResolver DefaultResolver = new CompositeResolver(BenchmarkRunnerClean.DefaultResolver, EngineResolver.Instance);
 
         public IHost Host { get; set; }
-        public Action<long> WorkloadActionNoUnroll { get; set; }
-        public Action<long> WorkloadActionUnroll { get; set; }
+        public Func<long, IClock, ValueTask<ClockSpan>> WorkloadActionNoUnroll { get; set; }
+        public Func<long, IClock, ValueTask<ClockSpan>> WorkloadActionUnroll { get; set; }
         public Action Dummy1Action { get; set; }
         public Action Dummy2Action { get; set; }
         public Action Dummy3Action { get; set; }
-        public Action<long> OverheadActionNoUnroll { get; set; }
-        public Action<long> OverheadActionUnroll { get; set; }
+        public Func<long, IClock, ValueTask<ClockSpan>> OverheadActionNoUnroll { get; set; }
+        public Func<long, IClock, ValueTask<ClockSpan>> OverheadActionUnroll { get; set; }
         public Job TargetJob { get; set; } = Job.Default;
         public long OperationsPerInvoke { get; set; } = 1;
-        public Action GlobalSetupAction { get; set; }
-        public Action GlobalCleanupAction { get; set; }
-        public Action IterationSetupAction { get; set; }
-        public Action IterationCleanupAction { get; set; }
+        public Func<ValueTask> GlobalSetupAction { get; set; }
+        public Func<ValueTask> GlobalCleanupAction { get; set; }
+        public Func<ValueTask> IterationSetupAction { get; set; }
+        public Func<ValueTask> IterationCleanupAction { get; set; }
         public bool MeasureExtraStats { get; set; }
 
         [PublicAPI] public string BenchmarkName { get;  set; }

src/BenchmarkDotNet/Engines/IEngine.cs

@@ -1,10 +1,11 @@
 using System;
 using System.Diagnostics.CodeAnalysis;
+using System.Threading.Tasks;
 using BenchmarkDotNet.Characteristics;
 using BenchmarkDotNet.Jobs;
 using BenchmarkDotNet.Reports;
 using JetBrains.Annotations;
-using NotNullAttribute = JetBrains.Annotations.NotNullAttribute;
+using Perfolizer.Horology;
 
 namespace BenchmarkDotNet.Engines
 {
@@ -24,16 +25,16 @@ public interface IEngine : IDisposable
         long OperationsPerInvoke { get; }
 
         [CanBeNull]
-        Action GlobalSetupAction { get; }
+        Func<ValueTask> GlobalSetupAction { get; }
 
         [CanBeNull]
-        Action GlobalCleanupAction { get; }
+        Func<ValueTask> GlobalCleanupAction { get; }
 
         [NotNull]
-        Action<long> WorkloadAction { get; }
+        Func<long, IClock, ValueTask<ClockSpan>> WorkloadAction { get; }
 
         [NotNull]
-        Action<long> OverheadAction { get; }
+        Func<long, IClock, ValueTask<ClockSpan>> OverheadAction { get; }
 
         IResolver Resolver { get; }
 

src/BenchmarkDotNet/Helpers/AwaitHelper.cs

@@ -75,5 +75,28 @@ public T GetResult<T>(ValueTask<T> task)
             }
             return awaiter.GetResult();
         }
+
+        public static ValueTask ToValueTaskVoid(Task task)
+        {
+            return new ValueTask(task);
+        }
+
+        public static ValueTask ToValueTaskVoid<T>(Task<T> task)
+        {
+            return new ValueTask(task);
+        }
+
+        public static ValueTask ToValueTaskVoid(ValueTask task)
+        {
+            return task;
+        }
+
+        // ValueTask<T> unfortunately can't be converted to a ValueTask for free, so we must create a state machine.
+        // It's not a big deal though, as this is only used for Setup/Cleanup where allocations aren't measured.
+        // And in practice, this should never be used, as (Value)Task<T> Setup/Cleanup methods have no utility.
+        public static async ValueTask ToValueTaskVoid<T>(ValueTask<T> task)
+        {
+            _ = await task.ConfigureAwait(false);
+        }
     }
 }