Support custom awaitable types

src/BenchmarkDotNet/Code/CodeGenerator.cs

@@ -7,6 +7,7 @@
 using System.Text;
 using System.Threading.Tasks;
 using BenchmarkDotNet.Characteristics;
+using BenchmarkDotNet.Configs;
 using BenchmarkDotNet.Disassemblers;
 using BenchmarkDotNet.Environments;
 using BenchmarkDotNet.Extensions;
@@ -33,7 +34,9 @@ internal static string Generate(BuildPartition buildPartition)
             {
                 var benchmark = buildInfo.BenchmarkCase;
 
-                var provider = GetDeclarationsProvider(benchmark.Descriptor);
+                var provider = GetDeclarationsProvider(benchmark.Descriptor, benchmark.Config);
+
+                provider.OverrideUnrollFactor(benchmark);
 
                 string passArguments = GetPassArguments(benchmark);
 
@@ -49,6 +52,7 @@ internal static string Generate(BuildPartition buildPartition)
                     .Replace("$WorkloadMethodReturnType$", provider.WorkloadMethodReturnTypeName)
                     .Replace("$WorkloadMethodReturnTypeModifiers$", provider.WorkloadMethodReturnTypeModifiers)
                     .Replace("$OverheadMethodReturnTypeName$", provider.OverheadMethodReturnTypeName)
+                    .Replace("$InitializeAsyncBenchmarkRunnerFields$", provider.GetInitializeAsyncBenchmarkRunnerFields(buildInfo.Id))
                     .Replace("$GlobalSetupMethodName$", provider.GlobalSetupMethodName)
                     .Replace("$GlobalCleanupMethodName$", provider.GlobalCleanupMethodName)
                     .Replace("$IterationSetupMethodName$", provider.IterationSetupMethodName)
@@ -59,8 +63,10 @@ internal static string Generate(BuildPartition buildPartition)
                     .Replace("$ParamsContent$", GetParamsContent(benchmark))
                     .Replace("$ArgumentsDefinition$", GetArgumentsDefinition(benchmark))
                     .Replace("$DeclareArgumentFields$", GetDeclareArgumentFields(benchmark))
-                    .Replace("$InitializeArgumentFields$", GetInitializeArgumentFields(benchmark)).Replace("$LoadArguments$", GetLoadArguments(benchmark))
+                    .Replace("$InitializeArgumentFields$", GetInitializeArgumentFields(benchmark))
+                    .Replace("$LoadArguments$", GetLoadArguments(benchmark))
                     .Replace("$PassArguments$", passArguments)
+                    .Replace("$PassArgumentsDirect$", GetPassArgumentsDirect(benchmark))
                     .Replace("$EngineFactoryType$", GetEngineFactoryTypeName(benchmark))
                     .Replace("$MeasureExtraStats$", buildInfo.Config.HasExtraStatsDiagnoser() ? "true" : "false")
                     .Replace("$DisassemblerEntryMethodName$", DisassemblerConstants.DisassemblerEntryMethodName)
@@ -148,21 +154,10 @@ private static string GetJobsSetDefinition(BenchmarkCase benchmarkCase)
                 Replace("; ", ";\n                ");
         }
 
-        private static DeclarationsProvider GetDeclarationsProvider(Descriptor descriptor)
+        private static DeclarationsProvider GetDeclarationsProvider(Descriptor descriptor, IConfig config)
         {
             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<>)
-                    || method.ReturnType.GetTypeInfo().GetGenericTypeDefinition() == typeof(ValueTask<>)))
-            {
-                return new GenericTaskDeclarationsProvider(descriptor);
-            }
-
             if (method.ReturnType == typeof(void))
             {
                 bool isUsingAsyncKeyword = method.HasAttribute<AsyncStateMachineAttribute>();
@@ -183,6 +178,11 @@ private static DeclarationsProvider GetDeclarationsProvider(Descriptor descripto
                     return new ByRefDeclarationsProvider(descriptor);
             }
 
+            if (config.GetIsAwaitable(method.ReturnType, out var adapter))
+            {
+                return new AwaitableDeclarationsProvider(descriptor, adapter);
+            }
+
             return new NonVoidDeclarationsProvider(descriptor);
         }
 
@@ -224,6 +224,12 @@ private static string GetPassArguments(BenchmarkCase benchmarkCase)
                 benchmarkCase.Descriptor.WorkloadMethod.GetParameters()
                     .Select((parameter, index) => $"{GetParameterModifier(parameter)} arg{index}"));
 
+        private static string GetPassArgumentsDirect(BenchmarkCase benchmarkCase)
+            => string.Join(
+                ", ",
+                benchmarkCase.Descriptor.WorkloadMethod.GetParameters()
+                    .Select((parameter, index) => $"{GetParameterModifier(parameter)} __argField{index}"));
+
         private static string GetExtraAttributes(Descriptor descriptor)
             => descriptor.WorkloadMethod.GetCustomAttributes(false).OfType<STAThreadAttribute>().Any() ? "[System.STAThreadAttribute]" : string.Empty;
 

src/BenchmarkDotNet/Code/DeclarationsProvider.cs

@@ -2,6 +2,7 @@
 using System.Linq;
 using System.Reflection;
 using System.Threading.Tasks;
+using BenchmarkDotNet.Configs;
 using BenchmarkDotNet.Engines;
 using BenchmarkDotNet.Extensions;
 using BenchmarkDotNet.Helpers;
@@ -11,9 +12,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 +24,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 +46,18 @@ internal abstract class DeclarationsProvider
 
         public string OverheadMethodReturnTypeName => OverheadMethodReturnType.GetCorrectCSharpTypeName();
 
+        public virtual string GetInitializeAsyncBenchmarkRunnerFields(BenchmarkId id) => null;
+
+        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 +66,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 +148,42 @@ internal class ByReadOnlyRefDeclarationsProvider : ByRefDeclarationsProvider
         public override string WorkloadMethodReturnTypeModifiers => "ref readonly";
     }
 
-    internal class TaskDeclarationsProvider : VoidDeclarationsProvider
+    internal class AwaitableDeclarationsProvider : DeclarationsProvider
     {
-        public TaskDeclarationsProvider(Descriptor descriptor) : base(descriptor) { }
+        private readonly ConcreteAsyncAdapter adapter;
 
-        // 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 AwaitableDeclarationsProvider(Descriptor descriptor, ConcreteAsyncAdapter adapter) : base(descriptor) => this.adapter = adapter;
 
-        public override string GetWorkloadMethodCall(string passArguments) => $"{Descriptor.WorkloadMethod.Name}({passArguments}).GetAwaiter().GetResult()";
+        public override string ReturnsDefinition => "RETURNS_AWAITABLE";
 
-        protected override Type WorkloadMethodReturnType => typeof(void);
-    }
+        public override string OverheadImplementation => $"return default({OverheadMethodReturnType.GetCorrectCSharpTypeName()});";
 
-    /// <summary>
-    /// declarations provider for <see cref="Task{TResult}" /> and <see cref="ValueTask{TResult}" />
-    /// </summary>
-    internal class GenericTaskDeclarationsProvider : NonVoidDeclarationsProvider
-    {
-        public GenericTaskDeclarationsProvider(Descriptor descriptor) : base(descriptor) { }
+        protected override Type OverheadMethodReturnType => WorkloadMethodReturnType.IsValueType && !WorkloadMethodReturnType.IsPrimitive
+            ? typeof(EmptyAwaiter) // we return this simple type so we don't include the cost of a large struct in the overhead
+            : WorkloadMethodReturnType;
 
-        protected override Type WorkloadMethodReturnType => Descriptor.WorkloadMethod.ReturnType.GetTypeInfo().GetGenericArguments().Single();
+        private string GetRunnableName(BenchmarkId id) => $"BenchmarkDotNet.Autogenerated.Runnable_{id}";
 
-        // 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(); }}";
+        public override string GetInitializeAsyncBenchmarkRunnerFields(BenchmarkId id)
+        {
+            string awaitableAdapterTypeName = adapter.awaitableAdapterType.GetCorrectCSharpTypeName();
+            string asyncMethodBuilderAdapterTypeName = adapter.asyncMethodBuilderAdapterType.GetCorrectCSharpTypeName();
+
+            string awaiterTypeName = adapter.awaiterType.GetCorrectCSharpTypeName();
+            string overheadAwaiterTypeName = adapter.awaiterType.IsValueType
+                ? typeof(EmptyAwaiter).GetCorrectCSharpTypeName() // we use this simple type so we don't include the cost of a large struct in the overhead
+                : awaiterTypeName;
+            string appendResultType = adapter.resultType == null ? string.Empty : $", {adapter.resultType.GetCorrectCSharpTypeName()}";
+
+            string runnableName = GetRunnableName(id);
+
+            var workloadRunnerTypeName = $"BenchmarkDotNet.Engines.AsyncWorkloadRunner<{runnableName}.WorkloadFunc, {asyncMethodBuilderAdapterTypeName}, {awaitableAdapterTypeName}, {WorkloadMethodReturnTypeName}, {awaiterTypeName}{appendResultType}>";
+            var overheadRunnerTypeName = $"BenchmarkDotNet.Engines.AsyncOverheadRunner<{runnableName}.OverheadFunc, {asyncMethodBuilderAdapterTypeName}, {OverheadMethodReturnTypeName}, {overheadAwaiterTypeName}{appendResultType}>";
+
+            return $"__asyncWorkloadRunner = new {workloadRunnerTypeName}(new {runnableName}.WorkloadFunc(this));" + Environment.NewLine
+     + $"            __asyncOverheadRunner = new {overheadRunnerTypeName}(new {runnableName}.OverheadFunc(this));";
+        }
 
-        public override string GetWorkloadMethodCall(string passArguments) => $"{Descriptor.WorkloadMethod.Name}({passArguments}).GetAwaiter().GetResult()";
+        public override void OverrideUnrollFactor(BenchmarkCase benchmarkCase) => benchmarkCase.ForceUnrollFactorForAsync();
     }
 }