Fix async GlobalSetup/GlobalCleanup not being awaited with InProcessEmit toolchain. (#2109)

Author: timcassell

src/BenchmarkDotNet/Toolchains/InProcess/Emit/Implementation/Emitters/RunnableEmitter.cs

@@ -247,6 +247,10 @@ private static void EmitNoArgsMethodCallPopReturn(
         private TypeBuilder runnableBuilder;
         private ConsumableTypeInfo consumableInfo;
         private ConsumeEmitter consumeEmitter;
+        private ConsumableTypeInfo globalSetupReturnInfo;
+        private ConsumableTypeInfo globalCleanupReturnInfo;
+        private ConsumableTypeInfo iterationSetupReturnInfo;
+        private ConsumableTypeInfo iterationCleanupReturnInfo;
 
         private FieldBuilder globalSetupActionField;
         private FieldBuilder globalCleanupActionField;
@@ -358,13 +362,22 @@ private void InitForEmitRunnable(BenchmarkBuildInfo newBenchmark)
 
             consumableInfo = new ConsumableTypeInfo(benchmark.BenchmarkCase.Descriptor.WorkloadMethod.ReturnType);
             consumeEmitter = ConsumeEmitter.GetConsumeEmitter(consumableInfo);
+            globalSetupReturnInfo = GetConsumableTypeInfo(benchmark.BenchmarkCase.Descriptor.GlobalSetupMethod?.ReturnType);
+            globalCleanupReturnInfo = GetConsumableTypeInfo(benchmark.BenchmarkCase.Descriptor.GlobalCleanupMethod?.ReturnType);
+            iterationSetupReturnInfo = GetConsumableTypeInfo(benchmark.BenchmarkCase.Descriptor.IterationSetupMethod?.ReturnType);
+            iterationCleanupReturnInfo = GetConsumableTypeInfo(benchmark.BenchmarkCase.Descriptor.IterationCleanupMethod?.ReturnType);
 
             // Init types
             runnableBuilder = DefineRunnableTypeBuilder(benchmark, moduleBuilder);
             overheadDelegateType = EmitOverheadDelegateType();
             workloadDelegateType = EmitWorkloadDelegateType();
         }
 
+        private static ConsumableTypeInfo GetConsumableTypeInfo(Type methodReturnType)
+        {
+            return methodReturnType == null ? null : new ConsumableTypeInfo(methodReturnType);
+        }
+
         private Type EmitOverheadDelegateType()
         {
             // .class public auto ansi sealed BenchmarkDotNet.Autogenerated.Runnable_0OverheadDelegate
@@ -890,34 +903,84 @@ private void EmitSetupCleanupMethods()
         {
             // Emit Setup/Cleanup methods
             // We emit empty method instead of EmptyAction = "() => { }"
-            globalSetupMethod = EmitWrapperMethod(
-                GlobalSetupMethodName,
-                Descriptor.GlobalSetupMethod);
-            globalCleanupMethod = EmitWrapperMethod(
-                GlobalCleanupMethodName,
-                Descriptor.GlobalCleanupMethod);
-            iterationSetupMethod = EmitWrapperMethod(
-                IterationSetupMethodName,
-                Descriptor.IterationSetupMethod);
-            iterationCleanupMethod = EmitWrapperMethod(
-                IterationCleanupMethodName,
-                Descriptor.IterationCleanupMethod);
+            globalSetupMethod = EmitWrapperMethod(GlobalSetupMethodName, Descriptor.GlobalSetupMethod, globalSetupReturnInfo);
+            globalCleanupMethod = EmitWrapperMethod(GlobalCleanupMethodName, Descriptor.GlobalCleanupMethod, globalCleanupReturnInfo);
+            iterationSetupMethod = EmitWrapperMethod(IterationSetupMethodName, Descriptor.IterationSetupMethod, iterationSetupReturnInfo);
+            iterationCleanupMethod = EmitWrapperMethod(IterationCleanupMethodName, Descriptor.IterationCleanupMethod, iterationCleanupReturnInfo);
         }
 
-        private MethodBuilder EmitWrapperMethod(string methodName, MethodInfo optionalTargetMethod)
+        private MethodBuilder EmitWrapperMethod(string methodName, MethodInfo optionalTargetMethod, ConsumableTypeInfo returnTypeInfo)
         {
             var methodBuilder = runnableBuilder.DefinePrivateVoidInstanceMethod(methodName);
 
             var ilBuilder = methodBuilder.GetILGenerator();
 
             if (optionalTargetMethod != null)
-                EmitNoArgsMethodCallPopReturn(methodBuilder, optionalTargetMethod, ilBuilder, forceDirectCall: true);
+            {
+                if (returnTypeInfo?.IsAwaitable == true)
+                {
+                    EmitAwaitableSetupTeardown(methodBuilder, optionalTargetMethod, ilBuilder, returnTypeInfo);
+                }
+                else
+                {
+                    EmitNoArgsMethodCallPopReturn(methodBuilder, optionalTargetMethod, ilBuilder, forceDirectCall: true);
+                }
+            }
 
             ilBuilder.EmitVoidReturn(methodBuilder);
 
             return methodBuilder;
         }
 
+        private void EmitAwaitableSetupTeardown(
+            MethodBuilder methodBuilder,
+            MethodInfo targetMethod,
+            ILGenerator ilBuilder,
+            ConsumableTypeInfo returnTypeInfo)
+        {
+            if (targetMethod == null)
+                throw new ArgumentNullException(nameof(targetMethod));
+
+            if (returnTypeInfo.WorkloadMethodReturnType == typeof(void))
+            {
+                ilBuilder.Emit(OpCodes.Ldarg_0);
+            }
+            /*
+                // call for instance
+                // GlobalSetup();
+                IL_0006: ldarg.0
+                IL_0007: call instance void [BenchmarkDotNet]BenchmarkDotNet.Samples.SampleBenchmark::GlobalSetup()
+            */
+            /*
+                // call for static
+                // GlobalSetup();
+                IL_0006: call string [BenchmarkDotNet]BenchmarkDotNet.Samples.SampleBenchmark::GlobalCleanup()
+            */
+            if (targetMethod.IsStatic)
+            {
+                ilBuilder.Emit(OpCodes.Call, targetMethod);
+
+            }
+            else if (methodBuilder.IsStatic)
+            {
+                throw new InvalidOperationException(
+                    $"[BUG] Static method {methodBuilder.Name} tries to call instance member {targetMethod.Name}");
+            }
+            else
+            {
+                ilBuilder.Emit(OpCodes.Ldarg_0);
+                ilBuilder.Emit(OpCodes.Call, targetMethod);
+            }
+
+            /*
+                // BenchmarkDotNet.Helpers.AwaitHelper.GetResult(...);
+                IL_000e: call !!0 BenchmarkDotNet.Helpers.AwaitHelper::GetResult<int32>(valuetype [System.Runtime]System.Threading.Tasks.ValueTask`1<!!0>)
+            */
+
+            ilBuilder.Emit(OpCodes.Call, returnTypeInfo.GetResultMethod);
+            ilBuilder.Emit(OpCodes.Pop);
+        }
+
         private void EmitCtorBody()
         {
             var ilBuilder = ctorMethod.GetILGenerator();

tests/BenchmarkDotNet.IntegrationTests/InProcessEmitTest.cs

@@ -239,35 +239,125 @@ public static ValueTask<decimal> InvokeOnceStaticValueTaskOfT()
             }
         }
 
-        [Fact]
-        public void InProcessEmitToolchainSupportsIterationSetupAndCleanup()
+        [Theory]
+        [InlineData(typeof(IterationSetupCleanup))]
+        [InlineData(typeof(GlobalSetupCleanupTask))]
+        [InlineData(typeof(GlobalSetupCleanupValueTask))]
+        [InlineData(typeof(GlobalSetupCleanupValueTaskSource))]
+        public void InProcessEmitToolchainSupportsSetupAndCleanup(Type benchmarkType)
         {
             var logger = new OutputLogger(Output);
             var config = CreateInProcessConfig(logger);
 
-            WithIterationSetupAndCleanup.SetupCounter = 0;
-            WithIterationSetupAndCleanup.BenchmarkCounter = 0;
-            WithIterationSetupAndCleanup.CleanupCounter = 0;
+            Counters.SetupCounter = 0;
+            Counters.BenchmarkCounter = 0;
+            Counters.CleanupCounter = 0;
 
-            var summary = CanExecute<WithIterationSetupAndCleanup>(config);
+            var summary = CanExecute(benchmarkType, config);
 
-            Assert.Equal(1, WithIterationSetupAndCleanup.SetupCounter);
-            Assert.Equal(16, WithIterationSetupAndCleanup.BenchmarkCounter);
-            Assert.Equal(1, WithIterationSetupAndCleanup.CleanupCounter);
+            Assert.Equal(1, Counters.SetupCounter);
+            Assert.Equal(16, Counters.BenchmarkCounter);
+            Assert.Equal(1, Counters.CleanupCounter);
         }
 
-        public class WithIterationSetupAndCleanup
+        private static class Counters
         {
             public static int SetupCounter, BenchmarkCounter, CleanupCounter;
+        }
 
+        public class IterationSetupCleanup
+        {
             [IterationSetup]
-            public void Setup() => Interlocked.Increment(ref SetupCounter);
+            public void Setup() => Interlocked.Increment(ref Counters.SetupCounter);
 
             [Benchmark]
-            public void Benchmark() => Interlocked.Increment(ref BenchmarkCounter);
+            public void Benchmark() => Interlocked.Increment(ref Counters.BenchmarkCounter);
 
             [IterationCleanup]
-            public void Cleanup() => Interlocked.Increment(ref CleanupCounter);
+            public void Cleanup() => Interlocked.Increment(ref Counters.CleanupCounter);
+        }
+
+        public class GlobalSetupCleanupTask
+        {
+            [GlobalSetup]
+            public static async Task GlobalSetup()
+            {
+                await Task.Yield();
+                Interlocked.Increment(ref Counters.SetupCounter);
+            }
+
+            [GlobalCleanup]
+            public async Task<int> GlobalCleanup()
+            {
+                await Task.Yield();
+                Interlocked.Increment(ref Counters.CleanupCounter);
+                return 42;
+            }
+
+            [Benchmark]
+            public void InvokeOnceVoid()
+            {
+                Interlocked.Increment(ref Counters.BenchmarkCounter);
+            }
+        }
+
+        public class GlobalSetupCleanupValueTask
+        {
+            [GlobalSetup]
+            public static async ValueTask GlobalSetup()
+            {
+                await Task.Yield();
+                Interlocked.Increment(ref Counters.SetupCounter);
+            }
+
+            [GlobalCleanup]
+            public async ValueTask<int> GlobalCleanup()
+            {
+                await Task.Yield();
+                Interlocked.Increment(ref Counters.CleanupCounter);
+                return 42;
+            }
+
+            [Benchmark]
+            public void InvokeOnceVoid()
+            {
+                Interlocked.Increment(ref Counters.BenchmarkCounter);
+            }
+        }
+
+        public class GlobalSetupCleanupValueTaskSource
+        {
+            private readonly static ValueTaskSource<int> valueTaskSource = new ();
+
+            [GlobalSetup]
+            public static ValueTask GlobalSetup()
+            {
+                valueTaskSource.Reset();
+                Task.Delay(1).ContinueWith(_ =>
+                {
+                    Interlocked.Increment(ref Counters.SetupCounter);
+                    valueTaskSource.SetResult(42);
+                });
+                return new ValueTask(valueTaskSource, valueTaskSource.Token);
+            }
+
+            [GlobalCleanup]
+            public ValueTask<int> GlobalCleanup()
+            {
+                valueTaskSource.Reset();
+                Task.Delay(1).ContinueWith(_ =>
+                {
+                    Interlocked.Increment(ref Counters.CleanupCounter);
+                    valueTaskSource.SetResult(42);
+                });
+                return new ValueTask<int>(valueTaskSource, valueTaskSource.Token);
+            }
+
+            [Benchmark]
+            public void InvokeOnceVoid()
+            {
+                Interlocked.Increment(ref Counters.BenchmarkCounter);
+            }
         }
     }
 }