fix race condition in GetResult/SetResult

.gitignore

@@ -1,3 +1,4 @@
 .vs/
 bin/
-obj/
+obj/
+BenchmarkDotNet.Artifacts/

global.json

@@ -1,5 +1,5 @@
 {
   "sdk": {
-    "version": "3.0.100-preview6-012264"
+    "version": "3.0.100-preview6"
   }
 }

src/PooledAwait/Internal/PooledState.cs

@@ -30,29 +30,21 @@ public PooledValueTask PooledValueTask
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         void IValueTaskSource.GetResult(short token)
         {
-            // we only support getting the result once; doing this recycles
-            // the source and advances the token
-            try
-            {
-                if (_source.GetStatus(token) == ValueTaskSourceStatus.Pending) WaitForResult(token);
-                _source.GetResult(token);
-            }
-            finally
-            {
-                _source.Reset();
-                Pool<PooledState>.TryPut(this);
-                Counters.PooledStateRecycled.Increment();
-            }
-        }
+            // we only support getting the result once; doing this recycles the source and advances the token
 
-        [MethodImpl(MethodImplOptions.NoInlining)]
-        private void WaitForResult(short token)
-        {
-            lock (SyncLock)
+            lock (SyncLock) // we need to be really paranoid about cross-threading over changing the token
             {
-                if (token == _source.Version && _source.GetStatus(token) == ValueTaskSourceStatus.Pending)
+                var status = _source.GetStatus(token); // do this *outside* the try/finally
+                try // so that we don't increment the counter if someone asks for the wrong value
+                {
+                    if (status == ValueTaskSourceStatus.Pending) Monitor.Wait(SyncLock);
+                    _source.GetResult(token);
+                }
+                finally
                 {
-                    Monitor.Wait(SyncLock);
+                    _source.Reset();
+                    Pool<PooledState>.TryPut(this);
+                    Counters.PooledStateRecycled.Increment();
                 }
             }
         }

src/PooledAwait/Internal/PooledStateT.cs

@@ -30,29 +30,21 @@ public PooledValueTask<T> PooledValueTask
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         T IValueTaskSource<T>.GetResult(short token)
         {
-            // we only support getting the result once; doing this recycles
-            // the source and advances the token
-            try
-            {
-                if (_source.GetStatus(token) == ValueTaskSourceStatus.Pending) WaitForResult(token);
-                return _source.GetResult(token);
-            }
-            finally
-            {
-                _source.Reset();
-                Pool<PooledState<T>>.TryPut(this);
-                Counters.PooledStateRecycled.Increment();
-            }
-        }
+            // we only support getting the result once; doing this recycles the source and advances the token
 
-        [MethodImpl(MethodImplOptions.NoInlining)]
-        private void WaitForResult(short token)
-        {
-            lock (SyncLock)
+            lock (SyncLock) // we need to be really paranoid about cross-threading over changing the token
             {
-                if (token == _source.Version && _source.GetStatus(token) == ValueTaskSourceStatus.Pending)
+                var status = _source.GetStatus(token); // do this *outside* the try/finally
+                try // so that we don't increment the counter if someone asks for the wrong value
+                {
+                    if (status == ValueTaskSourceStatus.Pending) Monitor.Wait(SyncLock);
+                    return _source.GetResult(token);
+                }
+                finally
                 {
-                    Monitor.Wait(SyncLock);
+                    _source.Reset();
+                    Pool<PooledState<T>>.TryPut(this);
+                    Counters.PooledStateRecycled.Increment();
                 }
             }
         }

tests/Benchmark/ComparisonBenchmarks.cs

@@ -0,0 +1,160 @@
+using BenchmarkDotNet.Attributes;
+using BenchmarkDotNet.Configs;
+using PooledAwait;
+using System.Threading.Tasks;
+
+namespace Benchmark
+{
+    [MemoryDiagnoser, ShortRunJob]
+    [GroupBenchmarksBy(BenchmarkLogicalGroupRule.ByCategory)]
+    [CategoriesColumn]
+    public class ComparisonBenchmarks
+    {
+        private const int InnerOps = 1000;
+
+        [Benchmark(OperationsPerInvoke = InnerOps, Description = nameof(Task), Baseline = true)]
+        [BenchmarkCategory("int")]
+        public async Task<int> ViaTaskT()
+        {
+            int sum = 0;
+            for (int i = 0; i < InnerOps; i++)
+                sum += await Inner(1, 2).ConfigureAwait(false);
+            return sum;
+            static async Task<int> Inner(int x, int y)
+            {
+                int i = x;
+                await Task.Yield();
+                i *= y;
+                await Task.Yield();
+                return 5 * i;
+            }
+        }
+
+        [Benchmark(OperationsPerInvoke = InnerOps, Description = nameof(Task), Baseline = true)]
+        [BenchmarkCategory("void")]
+        public async Task ViaTask()
+        {
+            for (int i = 0; i < InnerOps; i++)
+                await Inner().ConfigureAwait(false);
+
+            static async Task Inner()
+            {
+                await Task.Yield();
+                await Task.Yield();
+            }
+        }
+
+        [Benchmark(OperationsPerInvoke = InnerOps, Description = nameof(ValueTask))]
+        [BenchmarkCategory("int")]
+        public async ValueTask<int> ViaValueTaskT()
+        {
+            int sum = 0;
+            for (int i = 0; i < InnerOps; i++)
+                sum += await Inner(1, 2).ConfigureAwait(false);
+            return sum;
+            static async ValueTask<int> Inner(int x, int y)
+            {
+                int i = x;
+                await Task.Yield();
+                i *= y;
+                await Task.Yield();
+                return 5 * i;
+            }
+        }
+
+        [Benchmark(OperationsPerInvoke = InnerOps, Description = nameof(ValueTask))]
+        [BenchmarkCategory("void")]
+        public async ValueTask ViaValueTask()
+        {
+            for (int i = 0; i < InnerOps; i++)
+                await Inner().ConfigureAwait(false);
+            static async ValueTask Inner()
+            {
+                await Task.Yield();
+                await Task.Yield();
+            }
+        }
+
+        [Benchmark(OperationsPerInvoke = InnerOps, Description = nameof(PooledValueTask))]
+        [BenchmarkCategory("int")]
+        public ValueTask<int> ViaPooledValueTaskT()
+        {
+            return Impl(); // thunks the type back to ValueTaskT
+
+            static async PooledValueTask<int> Impl()
+            {
+                int sum = 0;
+                for (int i = 0; i < InnerOps; i++)
+                    sum += await Inner(1, 2).ConfigureAwait(false);
+                return sum;
+                static async PooledValueTask<int> Inner(int x, int y)
+                {
+                    int i = x;
+                    await Task.Yield();
+                    i *= y;
+                    await Task.Yield();
+                    return 5 * i;
+                }
+            }
+        }
+
+        [Benchmark(OperationsPerInvoke = InnerOps, Description = nameof(PooledValueTask))]
+        [BenchmarkCategory("void")]
+        public ValueTask ViaPooledValueTask()
+        {
+            return Impl(); // thunks the type back to ValueTaskT
+
+            static async PooledValueTask Impl()
+            {
+                for (int i = 0; i < InnerOps; i++)
+                    await Inner().ConfigureAwait(false);
+                static async PooledValueTask Inner()
+                {
+                    await Task.Yield();
+                    await Task.Yield();
+                }
+            }
+        }
+
+        [Benchmark(OperationsPerInvoke = InnerOps, Description = nameof(PooledTask))]
+        [BenchmarkCategory("void")]
+        public Task ViaPooledTask()
+        {
+            return Impl(); // thunks the type back to ValueTaskT
+
+            static async PooledTask Impl()
+            {
+                for (int i = 0; i < InnerOps; i++)
+                    await Inner().ConfigureAwait(false);
+                static async PooledTask Inner()
+                {
+                    await Task.Yield();
+                    await Task.Yield();
+                }
+            }
+        }
+
+        [Benchmark(OperationsPerInvoke = InnerOps, Description = nameof(PooledTask))]
+        [BenchmarkCategory("int")]
+        public Task ViaPooledTaskT()
+        {
+            return Impl(); // thunks the type back to ValueTaskT
+
+            static async PooledTask<int> Impl()
+            {
+                int sum = 0;
+                for (int i = 0; i < InnerOps; i++)
+                    sum += await Inner(1, 2).ConfigureAwait(false);
+                return sum;
+                static async PooledTask<int> Inner(int x, int y)
+                {
+                    int i = x;
+                    await Task.Yield();
+                    i *= y;
+                    await Task.Yield();
+                    return 5 * i;
+                }
+            }
+        }
+    }
+}