Async refactor WIP

Author: timcassell

BenchmarkDotNet.slnx

@@ -6,7 +6,7 @@
   <Folder Name="/samples/">
     <File Path="samples/Directory.Build.props" />
     <Project Path="samples/BenchmarkDotNet.Samples.FSharp/BenchmarkDotNet.Samples.FSharp.fsproj" />
-    <Project Path="samples/BenchmarkDotNet.Samples/BenchmarkDotNet.Samples.csproj" DefaultStartup="true" />
+    <Project Path="samples/BenchmarkDotNet.Samples/BenchmarkDotNet.Samples.csproj" />
   </Folder>
   <Folder Name="/src/">
     <Project Path="src/BenchmarkDotNet.Analyzers/BenchmarkDotNet.Analyzers.csproj" />

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>