Async refactor

Author: timcassell

build/BenchmarkDotNet.Build/Runners/BuildRunner.cs

@@ -8,7 +8,6 @@
 using Cake.Common.Tools.DotNet.Workload.Install;
 using Cake.Core;
 using Cake.Core.IO;
-using System;
 using System.IO;
 using System.Linq;
 
@@ -51,7 +50,6 @@ public void PackWeaver()
         {
             MSBuildSettings = context.MsBuildSettingsRestore,
         };
-        MaybeAppendArgument(restoreSettings);
         context.DotNetRestore(weaverPath.GetDirectory().FullPath, restoreSettings);
 
         context.Information("BuildSystemProvider: " + context.BuildSystem().Provider);
@@ -63,7 +61,6 @@ public void PackWeaver()
             Configuration = context.BuildConfiguration,
             Verbosity = context.BuildVerbosity
         };
-        MaybeAppendArgument(buildSettings);
         context.DotNetBuild(weaverPath.FullPath, buildSettings);
 
         var packSettings = new DotNetPackSettings
@@ -72,7 +69,6 @@ public void PackWeaver()
             MSBuildSettings = context.MsBuildSettingsPack,
             Configuration = context.BuildConfiguration
         };
-        MaybeAppendArgument(packSettings);
         context.DotNetPack(weaverPath.FullPath, packSettings);
     }
 

build/common.props

@@ -60,7 +60,7 @@
 
   <PropertyGroup>
     <!-- Increment this when the BenchmarkDotNet.Weaver package needs to be re-packed. -->
-    <WeaverVersionSuffix>-1</WeaverVersionSuffix>
+    <WeaverVersionSuffix>-2</WeaverVersionSuffix>
   </PropertyGroup>
 
   <ItemGroup>

src/BenchmarkDotNet.Annotations/Attributes/AsyncCallerTypeAttribute.cs

@@ -0,0 +1,15 @@
+using System;
+
+namespace BenchmarkDotNet.Attributes;
+
+/// <summary>
+/// When applied to an async benchmark method, overrides the return type of the async method that calls the benchmark method.
+/// </summary>
+[AttributeUsage(AttributeTargets.Method)]
+public sealed class AsyncCallerTypeAttribute(Type asyncCallerType) : Attribute
+{
+    /// <summary>
+    /// The return type of the async method that calls the benchmark method.
+    /// </summary>
+    public Type AsyncCallerType { get; private set; } = asyncCallerType;
+}

src/BenchmarkDotNet.Annotations/Attributes/CompilerServices/AggressivelyOptimizeMethods.cs

@@ -0,0 +1,19 @@
+using System;
+using System.Runtime.CompilerServices;
+
+namespace BenchmarkDotNet.Attributes.CompilerServices;
+
+// MethodImplOptions.AggressiveOptimization is applied to all methods to force them to go straight to tier1 JIT,
+// eliminating tiered JIT as a potential variable in measurements.
+// This is necessary because C# does not support any way to apply attributes to compiler-generated state machine methods.
+// This is applied both to the core Engine and auto-generated classes.
+#pragma warning disable CS1574
+/// <summary>
+/// Instructs the BenchmarkDotNet assembly weaver to apply <see cref="MethodImplOptions.AggressiveOptimization"/> to all declared
+/// methods in the annotated type and nested types that are not already annotated with <see cref="MethodImplOptions.NoOptimization"/>.
+/// </summary>
+#pragma warning restore CS1574
+[AttributeUsage(AttributeTargets.Class | AttributeTargets.Struct)]
+public sealed class AggressivelyOptimizeMethodsAttribute : Attribute
+{
+}

src/BenchmarkDotNet.Weaver/buildTransitive/netstandard2.0/BenchmarkDotNet.Weaver.targets

@@ -17,7 +17,7 @@
       Inputs="$(BenchmarkDotNetWeaveAssemblyPath)"
       Outputs="$(BenchmarkDotNetWeaveAssembliesStampFile)">
 
-    <WeaveAssemblyTask TargetAssembly="$(BenchmarkDotNetWeaveAssemblyPath)" />
+    <WeaveAssemblyTask TargetAssembly="$(BenchmarkDotNetWeaveAssemblyPath)" TreatWarningsAsErrors="$(TreatWarningsAsErrors)" />
 
     <!-- Create stamp file for incrementality -->
     <Touch Files="$(BenchmarkDotNetWeaveAssembliesStampFile)" AlwaysCreate="true" />

src/BenchmarkDotNet.Weaver/packages/BenchmarkDotNet.Weaver.0.16.0-develop-1.nupkg

@@ -24,45 +24,76 @@ public sealed class WeaveAssemblyTask : Task
     [Required]
     public string TargetAssembly { get; set; }
 
+    /// <summary>
+    /// Whether to treat warnings as errors.
+    /// </summary>
+    public bool TreatWarningsAsErrors { get; set; }
+
     /// <summary>
     /// Runs the weave assembly task.
     /// </summary>
     /// <returns><see langword="true"/> if successful; <see langword="false"/> otherwise.</returns>
     public override bool Execute()
-    {
+    {
         if (!File.Exists(TargetAssembly))
         {
             Log.LogError($"Assembly not found: {TargetAssembly}");
             return false;
         }
 
-
         bool benchmarkMethodsImplAdjusted = false;
         try
         {
             var module = ModuleDefinition.FromFile(TargetAssembly);
 
+            bool anyAdjustments = false;
             foreach (var type in module.GetAllTypes())
             {
-                // We can skip non-public types as they are not valid for benchmarks.
-                if (type.IsNotPublic)
+                if (type.CustomAttributes.Any(attr => attr.Constructor.DeclaringType.FullName == "BenchmarkDotNet.Attributes.CompilerServices.AggressivelyOptimizeMethodsAttribute"))
                 {
-                    continue;
+                    ApplyAggressiveOptimizationToMethods(type);
+
+                    void ApplyAggressiveOptimizationToMethods(TypeDefinition type)
+                    {
+                        // Apply AggressiveOptimization to all methods in the type and nested types that
+                        // aren't annotated with NoOptimization (this includes compiler-generated state machines).
+                        foreach (var method in type.Methods)
+                        {
+                            if ((method.ImplAttributes & MethodImplAttributes.NoOptimization) == 0)
+                            {
+                                var oldImpl = method.ImplAttributes;
+                                method.ImplAttributes |= MethodImplAttributes.AggressiveOptimization;
+                                anyAdjustments |= (oldImpl & MethodImplAttributes.AggressiveOptimization) == 0;
+                            }
+                        }
+
+                        // Recurse into nested types
+                        foreach (var nested in type.NestedTypes)
+                        {
+                            ApplyAggressiveOptimizationToMethods(nested);
+                        }
+                    }
                 }
 
-                foreach (var method in type.Methods)
+                // We can skip non-public types as they are not valid for benchmarks.
+                // !type.IsNotPublic handles nested types, while type.IsPublic does not.
+                if (!type.IsNotPublic)
                 {
-                    if (method.CustomAttributes.Any(IsBenchmarkAttribute))
+                    foreach (var method in type.Methods)
                     {
-                        var oldImpl = method.ImplAttributes;
-                        // Remove AggressiveInlining and add NoInlining.
-                        method.ImplAttributes = (oldImpl & ~MethodImplAttributes.AggressiveInlining) | MethodImplAttributes.NoInlining;
-                        benchmarkMethodsImplAdjusted |= (oldImpl & MethodImplAttributes.NoInlining) == 0;
+                        if (method.CustomAttributes.Any(IsBenchmarkAttribute))
+                        {
+                            var oldImpl = method.ImplAttributes;
+                            // Remove AggressiveInlining and add NoInlining.
+                            method.ImplAttributes = (oldImpl & ~MethodImplAttributes.AggressiveInlining) | MethodImplAttributes.NoInlining;
+                            benchmarkMethodsImplAdjusted |= (oldImpl & MethodImplAttributes.NoInlining) == 0;
+                            anyAdjustments |= benchmarkMethodsImplAdjusted;
+                        }
                     }
                 }
             }
 
-            if (benchmarkMethodsImplAdjusted)
+            if (anyAdjustments)
             {
                 // Write to a memory stream before overwriting the original file in case an exception occurs during the write (like unsupported platform).
                 // https://github.com/Washi1337/AsmResolver/issues/640
@@ -90,9 +121,17 @@ public override bool Execute()
         }
         catch (Exception e)
         {
-            Log.LogWarning($"Assembly weaving failed. Benchmark methods found requiring NoInlining: {benchmarkMethodsImplAdjusted}. Error:{Environment.NewLine}{e}");
+            if (TreatWarningsAsErrors)
+            {
+                Log.LogError($"Assembly weaving failed. Benchmark methods found requiring NoInlining: {benchmarkMethodsImplAdjusted}.");
+                Log.LogErrorFromException(e, true, true, null);
+            }
+            else
+            {
+                Log.LogWarning($"Assembly weaving failed. Benchmark methods found requiring NoInlining: {benchmarkMethodsImplAdjusted}. Error:{Environment.NewLine}{e}");
+            }
         }
-        return true;
+        return !Log.HasLoggedErrors;
     }
 
     private static bool IsBenchmarkAttribute(CustomAttribute attribute)

src/BenchmarkDotNet.Weaver/packages/BenchmarkDotNet.Weaver.0.16.0-develop-2.nupkg

@@ -1,5 +1,6 @@
 <Project Sdk="Microsoft.NET.Sdk">
   <Import Project="..\..\build\common.props" />
+  
   <PropertyGroup>
     <AssemblyTitle>BenchmarkDotNet</AssemblyTitle>
     <TargetFrameworks>netstandard2.0;net6.0;net8.0;net9.0;net10.0</TargetFrameworks>
@@ -49,5 +50,10 @@
       <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
     </PackageReference>
   </ItemGroup>
+  <!-- The transitive weaver reference is stripped during full pack, so we need to reference it directly. -->
+  <ItemGroup Condition="'$(IsFullPack)' == 'true'">
+    <PackageReference Include="BenchmarkDotNet.Weaver" Version="$(Version)$(WeaverVersionSuffix)" PrivateAssets="all" />
+  </ItemGroup>
+  
   <Import Project="..\..\build\common.targets" />
 </Project>

src/BenchmarkDotNet.Weaver/src/WeaveAssemblyTask.cs

@@ -5,7 +5,6 @@
 using System.Reflection;
 using System.Runtime.CompilerServices;
 using System.Text;
-using System.Threading.Tasks;
 using BenchmarkDotNet.Characteristics;
 using BenchmarkDotNet.Diagnosers;
 using BenchmarkDotNet.Disassemblers;
@@ -33,29 +32,17 @@ internal static string Generate(BuildPartition buildPartition)
             {
                 var benchmark = buildInfo.BenchmarkCase;
 
-                var provider = GetDeclarationsProvider(benchmark.Descriptor);
-
-                string passArguments = GetPassArguments(benchmark);
-
-                string benchmarkTypeCode = new SmartStringBuilder(ResourceHelper.LoadTemplate("BenchmarkType.txt"))
+                string benchmarkTypeCode = GetDeclarationsProvider(benchmark)
+                    .ReplaceTemplate(new SmartStringBuilder(ResourceHelper.LoadTemplate("BenchmarkType.txt")))
                     .Replace("$ID$", buildInfo.Id.ToString())
-                    .Replace("$OperationsPerInvoke$", provider.OperationsPerInvoke)
-                    .Replace("$WorkloadTypeName$", provider.WorkloadTypeName)
-                    .Replace("$GlobalSetupMethodName$", provider.GlobalSetupMethodName)
-                    .Replace("$GlobalCleanupMethodName$", provider.GlobalCleanupMethodName)
-                    .Replace("$IterationSetupMethodName$", provider.IterationSetupMethodName)
-                    .Replace("$IterationCleanupMethodName$", provider.IterationCleanupMethodName)
                     .Replace("$JobSetDefinition$", GetJobsSetDefinition(benchmark))
                     .Replace("$ParamsContent$", GetParamsContent(benchmark))
                     .Replace("$ArgumentsDefinition$", GetArgumentsDefinition(benchmark))
                     .Replace("$DeclareArgumentFields$", GetDeclareArgumentFields(benchmark))
                     .Replace("$InitializeArgumentFields$", GetInitializeArgumentFields(benchmark))
-                    .Replace("$LoadArguments$", GetLoadArguments(benchmark))
-                    .Replace("$PassArguments$", passArguments)
                     .Replace("$EngineFactoryType$", GetEngineFactoryTypeName(benchmark))
                     .Replace("$RunExtraIteration$", buildInfo.Config.HasExtraIterationDiagnoser(benchmark) ? "true" : "false")
                     .Replace("$DisassemblerEntryMethodName$", DisassemblerConstants.DisassemblerEntryMethodName)
-                    .Replace("$WorkloadMethodCall$", provider.GetWorkloadMethodCall(passArguments))
                     .Replace("$InProcessDiagnoserRouters$", GetInProcessDiagnoserRouters(buildInfo))
                     .ToString();
 
@@ -132,27 +119,21 @@ private static string GetJobsSetDefinition(BenchmarkCase benchmarkCase)
                 Replace("; ", ";\n                ");
         }
 
-        private static DeclarationsProvider GetDeclarationsProvider(Descriptor descriptor)
+        private static DeclarationsProvider GetDeclarationsProvider(BenchmarkCase benchmark)
         {
-            var method = descriptor.WorkloadMethod;
+            var method = benchmark.Descriptor.WorkloadMethod;
 
-            if (method.ReturnType == typeof(Task) || method.ReturnType == typeof(ValueTask))
-            {
-                return new AsyncDeclarationsProvider(descriptor);
-            }
-            if (method.ReturnType.GetTypeInfo().IsGenericType
-                && (method.ReturnType.GetTypeInfo().GetGenericTypeDefinition() == typeof(Task<>)
-                    || method.ReturnType.GetTypeInfo().GetGenericTypeDefinition() == typeof(ValueTask<>)))
+            if (method.ReturnType.IsAwaitable())
             {
-                return new AsyncDeclarationsProvider(descriptor);
+                return new AsyncDeclarationsProvider(benchmark);
             }
 
             if (method.ReturnType == typeof(void) && method.HasAttribute<AsyncStateMachineAttribute>())
             {
                 throw new NotSupportedException("async void is not supported by design");
             }
 
-            return new SyncDeclarationsProvider(descriptor);
+            return new SyncDeclarationsProvider(benchmark);
         }
 
         // internal for tests
@@ -168,31 +149,19 @@ private static string GetArgumentsDefinition(BenchmarkCase benchmarkCase)
             => string.Join(
                 ", ",
                 benchmarkCase.Descriptor.WorkloadMethod.GetParameters()
-                         .Select((parameter, index) => $"{GetParameterModifier(parameter)} {parameter.ParameterType.GetCorrectCSharpTypeName()} arg{index}"));
+                    .Select((parameter, index) => $"{GetParameterModifier(parameter)} {parameter.ParameterType.GetCorrectCSharpTypeName()} arg{index}"));
 
         private static string GetDeclareArgumentFields(BenchmarkCase benchmarkCase)
             => string.Join(
                 Environment.NewLine,
                 benchmarkCase.Descriptor.WorkloadMethod.GetParameters()
-                         .Select((parameter, index) => $"private {GetFieldType(parameter.ParameterType, benchmarkCase.Parameters.GetArgument(parameter.Name)).GetCorrectCSharpTypeName()} __argField{index};"));
+                    .Select((parameter, index) => $"public {GetFieldType(parameter.ParameterType, benchmarkCase.Parameters.GetArgument(parameter.Name)).GetCorrectCSharpTypeName()} __argField{index};"));
 
         private static string GetInitializeArgumentFields(BenchmarkCase benchmarkCase)
             => string.Join(
                 Environment.NewLine,
                 benchmarkCase.Descriptor.WorkloadMethod.GetParameters()
-                    .Select((parameter, index) => $"this.__argField{index} = {benchmarkCase.Parameters.GetArgument(parameter.Name).ToSourceCode()};")); // we init the fields in ctor to provoke all possible allocations and overhead of other type
-
-        private static string GetLoadArguments(BenchmarkCase benchmarkCase)
-            => string.Join(
-                Environment.NewLine,
-                benchmarkCase.Descriptor.WorkloadMethod.GetParameters()
-                    .Select((parameter, index) => $"{(parameter.ParameterType.IsByRef ? "ref" : string.Empty)} {parameter.ParameterType.GetCorrectCSharpTypeName()} arg{index} = {(parameter.ParameterType.IsByRef ? "ref" : string.Empty)} this.__argField{index};"));
-
-        private static string GetPassArguments(BenchmarkCase benchmarkCase)
-            => string.Join(
-                ", ",
-                benchmarkCase.Descriptor.WorkloadMethod.GetParameters()
-                    .Select((parameter, index) => $"{GetParameterModifier(parameter)} arg{index}"));
+                    .Select((parameter, index) => $"this.__fieldsContainer.__argField{index} = {benchmarkCase.Parameters.GetArgument(parameter.Name).ToSourceCode()};")); // we init the fields in ctor to provoke all possible allocations and overhead of other type
 
         private static string GetExtraAttributes(Descriptor descriptor)
             => descriptor.WorkloadMethod.GetCustomAttributes(false).OfType<STAThreadAttribute>().Any() ? "[System.STAThreadAttribute]" : string.Empty;
@@ -235,7 +204,7 @@ private static string GetInProcessDiagnoserRouters(BenchmarkBuildInfo buildInfo)
             }
         }
 
-        private static string GetParameterModifier(ParameterInfo parameterInfo)
+        internal static string GetParameterModifier(ParameterInfo parameterInfo)
         {
             if (!parameterInfo.ParameterType.IsByRef)
                 return string.Empty;
@@ -263,7 +232,7 @@ private static string GetNativeAotSwitch(BuildPartition buildPartition)
             @switch.AppendLine("switch (id) {");
 
             foreach (var buildInfo in buildPartition.Benchmarks)
-                @switch.AppendLine($"case {buildInfo.Id.Value}: BenchmarkDotNet.Autogenerated.Runnable_{buildInfo.Id.Value}.Run(host, benchmarkName, diagnoserRunMode); break;");
+                @switch.AppendLine($"case {buildInfo.Id.Value}: runTask = BenchmarkDotNet.Autogenerated.Runnable_{buildInfo.Id.Value}.Run(host, benchmarkName, diagnoserRunMode); break;");
 
             @switch.AppendLine("default: throw new System.NotSupportedException(\"invalid benchmark id\");");
             @switch.AppendLine("}");
@@ -279,28 +248,21 @@ private static Type GetFieldType(Type argumentType, ParameterInstance argument)
 
             return argumentType;
         }
+    }
 
-        private class SmartStringBuilder
-        {
-            private readonly string originalText;
-            private readonly StringBuilder builder;
-
-            public SmartStringBuilder(string text)
-            {
-                originalText = text;
-                builder = new StringBuilder(text);
-            }
-
-            public SmartStringBuilder Replace(string oldValue, string? newValue)
-            {
-                if (originalText.Contains(oldValue))
-                    builder.Replace(oldValue, newValue);
-                else
-                    builder.Append($"\n// '{oldValue}' not found");
-                return this;
-            }
+    internal class SmartStringBuilder(string text)
+    {
+        private readonly StringBuilder builder = new(text);
 
-            public override string ToString() => builder.ToString();
+        public SmartStringBuilder Replace(string oldValue, string? newValue)
+        {
+            if (text.Contains(oldValue))
+                builder.Replace(oldValue, newValue);
+            else
+                builder.Append($"\n// '{oldValue}' not found");
+            return this;
         }
+
+        public override string ToString() => builder.ToString();
     }
 }