Adjust AwaitHelper to allow multiple threads to use it concurrently.

Author: timcassell

src/BenchmarkDotNet/Helpers/AwaitHelper.cs

@@ -8,77 +8,95 @@ namespace BenchmarkDotNet.Helpers
 {
     public class AwaitHelper
     {
-        private readonly object awaiterLock = new object();
-        private readonly Action awaiterCallback;
-        private bool awaiterCompleted;
-
-        public AwaitHelper()
+        private class ValueTaskWaiter
         {
-            awaiterCallback = AwaiterCallback;
-        }
+            private readonly Action awaiterCallback;
+            private bool awaiterCompleted;
 
-        private void AwaiterCallback()
-        {
-            lock (awaiterLock)
+            internal ValueTaskWaiter()
             {
-                awaiterCompleted = true;
-                System.Threading.Monitor.Pulse(awaiterLock);
+                awaiterCallback = AwaiterCallback;
             }
-        }
 
-        // 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 void GetResult(Task task)
-        {
-            task.GetAwaiter().GetResult();
-        }
+            private void AwaiterCallback()
+            {
+                lock (this)
+                {
+                    awaiterCompleted = true;
+                    System.Threading.Monitor.Pulse(this);
+                }
+            }
 
-        public T GetResult<T>(Task<T> task)
-        {
-            return task.GetAwaiter().GetResult();
-        }
+            // Hook up a callback instead of converting to Task to prevent extra allocations on each benchmark run.
+            internal void GetResult(ValueTask task)
+            {
+                // Don't continue on the captured context, as that may result in a deadlock if the user runs this in-process.
+                var awaiter = task.ConfigureAwait(false).GetAwaiter();
+                if (!awaiter.IsCompleted)
+                {
+                    lock (this)
+                    {
+                        awaiterCompleted = false;
+                        awaiter.UnsafeOnCompleted(awaiterCallback);
+                        // Check if the callback executed synchronously before blocking.
+                        if (!awaiterCompleted)
+                        {
+                            System.Threading.Monitor.Wait(this);
+                        }
+                    }
+                }
+                awaiter.GetResult();
+            }
 
-        // It is illegal to call GetResult from an uncomplete ValueTask, so we must hook up a callback.
-        public void GetResult(ValueTask task)
-        {
-            // Don't continue on the captured context, as that may result in a deadlock if the user runs this in-process.
-            var awaiter = task.ConfigureAwait(false).GetAwaiter();
-            if (!awaiter.IsCompleted)
+            internal T GetResult<T>(ValueTask<T> task)
             {
-                lock (awaiterLock)
+                // Don't continue on the captured context, as that may result in a deadlock if the user runs this in-process.
+                var awaiter = task.ConfigureAwait(false).GetAwaiter();
+                if (!awaiter.IsCompleted)
                 {
-                    awaiterCompleted = false;
-                    awaiter.UnsafeOnCompleted(awaiterCallback);
-                    // Check if the callback executed synchronously before blocking.
-                    if (!awaiterCompleted)
+                    lock (this)
                     {
-                        System.Threading.Monitor.Wait(awaiterLock);
+                        awaiterCompleted = false;
+                        awaiter.UnsafeOnCompleted(awaiterCallback);
+                        // Check if the callback executed synchronously before blocking.
+                        if (!awaiterCompleted)
+                        {
+                            System.Threading.Monitor.Wait(this);
+                        }
                     }
                 }
+                return awaiter.GetResult();
             }
-            awaiter.GetResult();
         }
 
-        public T GetResult<T>(ValueTask<T> task)
+        // We use thread static field so that multiple threads can use individual lock object and callback.
+        [ThreadStatic]
+        private static ValueTaskWaiter ts_valueTaskWaiter;
+
+        private ValueTaskWaiter CurrentValueTaskWaiter
         {
-            // Don't continue on the captured context, as that may result in a deadlock if the user runs this in-process.
-            var awaiter = task.ConfigureAwait(false).GetAwaiter();
-            if (!awaiter.IsCompleted)
+            get
             {
-                lock (awaiterLock)
+                if (ts_valueTaskWaiter == null)
                 {
-                    awaiterCompleted = false;
-                    awaiter.UnsafeOnCompleted(awaiterCallback);
-                    // Check if the callback executed synchronously before blocking.
-                    if (!awaiterCompleted)
-                    {
-                        System.Threading.Monitor.Wait(awaiterLock);
-                    }
+                    ts_valueTaskWaiter = new ValueTaskWaiter();
                 }
+                return ts_valueTaskWaiter;
             }
-            return awaiter.GetResult();
         }
 
+        // 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 void GetResult(Task task) => task.GetAwaiter().GetResult();
+
+        public T GetResult<T>(Task<T> task) => task.GetAwaiter().GetResult();
+
+        // ValueTask can be backed by an IValueTaskSource that only supports asynchronous awaits, so we have to hook up a callback instead of calling .GetAwaiter().GetResult() like we do for Task.
+        // The alternative is to convert it to Task using .AsTask(), but that causes allocations which we must avoid for memory diagnoser.
+        public void GetResult(ValueTask task) => CurrentValueTaskWaiter.GetResult(task);
+
+        public T GetResult<T>(ValueTask<T> task) => CurrentValueTaskWaiter.GetResult(task);
+
         internal static MethodInfo GetGetResultMethod(Type taskType)
         {
             if (!taskType.IsGenericType)