Reduce async overhead

src/BenchmarkDotNet/Code/CodeGenerator.cs

@@ -35,6 +35,8 @@ internal static string Generate(BuildPartition buildPartition)
 
                 var provider = GetDeclarationsProvider(benchmark.Descriptor);
 
+                provider.OverrideUnrollFactor(benchmark);
+
                 string passArguments = GetPassArguments(benchmark);
 
                 string compilationId = $"{provider.ReturnsDefinition}_{buildInfo.Id}";
@@ -49,6 +51,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 +158,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 $"() => {method.Name}().GetAwaiter().GetResult()";
+                return $"() => BenchmarkDotNet.Helpers.AwaitHelper.ToValueTaskVoid({method.Name}())";
             }
 
-            return method.Name;
+            return $"() => {{ {method.Name}(); return new System.Threading.Tasks.ValueTask(); }}";
         }
     }
 
@@ -145,34 +146,18 @@ public ByReadOnlyRefDeclarationsProvider(Descriptor descriptor) : base(descripto
         public override string WorkloadMethodReturnTypeModifiers => "ref readonly";
     }
 
-    internal class TaskDeclarationsProvider : VoidDeclarationsProvider
+    internal class TaskDeclarationsProvider : DeclarationsProvider
     {
         public TaskDeclarationsProvider(Descriptor descriptor) : base(descriptor) { }
 
-        // we use GetAwaiter().GetResult() because it's fastest way to obtain the result in blocking way,
-        // and will eventually throw actual exception, not aggregated one
-        public override string WorkloadMethodDelegate(string passArguments)
-            => $"({passArguments}) => {{ {Descriptor.WorkloadMethod.Name}({passArguments}).GetAwaiter().GetResult(); }}";
-
-        public override string GetWorkloadMethodCall(string passArguments) => $"{Descriptor.WorkloadMethod.Name}({passArguments}).GetAwaiter().GetResult()";
+        public override string ReturnsDefinition => "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()});";
 
-        // we use GetAwaiter().GetResult() because it's fastest way to obtain the result in blocking way,
-        // and will eventually throw actual exception, not aggregated one
-        public override string WorkloadMethodDelegate(string passArguments)
-            => $"({passArguments}) => {{ return {Descriptor.WorkloadMethod.Name}({passArguments}).GetAwaiter().GetResult(); }}";
+        protected override Type OverheadMethodReturnType => WorkloadMethodReturnType;
 
-        public override string GetWorkloadMethodCall(string passArguments) => $"{Descriptor.WorkloadMethod.Name}({passArguments}).GetAwaiter().GetResult()";
+        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 @@ internal Engine(
             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)
             {
@@ -155,7 +158,7 @@ public Measurement RunIteration(IterationData data)
             var action = isOverhead ? OverheadAction : WorkloadAction;
 
             if (!isOverhead)
-                IterationSetupAction();
+                awaitHelper.GetResult(IterationSetupAction());
 
             GcCollect();
 
@@ -165,15 +168,14 @@ 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);
 
             if (!isOverhead)
-                IterationCleanupAction();
+                awaitHelper.GetResult(IterationCleanupAction());
 
             if (randomizeMemory)
                 RandomizeManagedHeapMemory();
@@ -196,17 +198,18 @@ public Measurement RunIteration(IterationData data)
             // it does not matter, because we have already obtained the results!
             EnableMonitoring();
 
-            IterationSetupAction(); // we run iteration setup first, so even if it allocates, it is not included in the results
+            awaitHelper.GetResult(IterationSetupAction()); // we run iteration setup first, so even if it allocates, it is not included in the results
 
             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();
 
-            IterationCleanupAction(); // we run iteration cleanup after collecting GC stats
+            awaitHelper.GetResult(IterationCleanupAction()); // we run iteration cleanup after collecting GC stats
 
             GcStats gcStats = (finalGcStats - initialGcStats).WithTotalOperations(data.InvokeCount * OperationsPerInvoke);
             ThreadingStats threadingStats = (finalThreadingStats - initialThreadingStats).WithTotalOperations(data.InvokeCount * OperationsPerInvoke);
@@ -220,14 +223,14 @@ private void Consume(in Span<byte> _) { }
         private void RandomizeManagedHeapMemory()
         {
             // invoke global cleanup before global setup
-            GlobalCleanupAction?.Invoke();
+            awaitHelper.GetResult(GlobalCleanupAction.Invoke());
 
             var gen0object = new byte[random.Next(32)];
             var lohObject = new byte[85 * 1024 + random.Next(32)];
 
             // we expect the key allocations to happen in global setup (not ctor)
             // so we call it while keeping the random-size objects alive
-            GlobalSetupAction?.Invoke();
+            awaitHelper.GetResult(GlobalSetupAction.Invoke());
 
             GC.KeepAlive(gen0object);
             GC.KeepAlive(lohObject);

src/BenchmarkDotNet/Engines/EngineFactory.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading.Tasks;
 using BenchmarkDotNet.Jobs;
 using Perfolizer.Horology;
 
@@ -25,7 +26,7 @@ public IEngine CreateReadyToRun(EngineParameters engineParameters)
             if (engineParameters.TargetJob == null)
                 throw new ArgumentNullException(nameof(engineParameters.TargetJob));
 
-            engineParameters.GlobalSetupAction?.Invoke(); // whatever the settings are, we MUST call global setup here, the global cleanup is part of Engine's Dispose
+            engineParameters.GlobalSetupAction.Invoke().AsTask().GetAwaiter().GetResult(); // whatever the settings are, we MUST call global setup here, the global cleanup is part of Engine's Dispose
 
             if (!engineParameters.NeedsJitting) // just create the engine, do NOT jit
                 return CreateMultiActionEngine(engineParameters);
@@ -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,9 +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 Perfolizer.Horology;
 using NotNullAttribute = JetBrains.Annotations.NotNullAttribute;
 
 namespace BenchmarkDotNet.Engines
@@ -24,16 +26,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; }