From 4f6344a12fa6caf94e14347dd35065591c2bc28c Mon Sep 17 00:00:00 2001
From: hadashiA <dev@hadashikick.jp>
Date: Thu, 14 Sep 2023 10:22:04 +0900
Subject: [PATCH] Add check to AutoResetUniTaskSource already returned to the
 pool

---
 .../UniTaskCompletionSourceTest.cs            | 80 +++++++++++++++++++
 .../Runtime/UniTaskCompletionSource.cs        | 19 +++--
 2 files changed, 91 insertions(+), 8 deletions(-)
 create mode 100644 src/UniTask.NetCoreTests/UniTaskCompletionSourceTest.cs

diff --git a/src/UniTask.NetCoreTests/UniTaskCompletionSourceTest.cs b/src/UniTask.NetCoreTests/UniTaskCompletionSourceTest.cs
new file mode 100644
index 00000000..1c5b943b
--- /dev/null
+++ b/src/UniTask.NetCoreTests/UniTaskCompletionSourceTest.cs
@@ -0,0 +1,80 @@
+using System.Threading.Tasks;
+using Cysharp.Threading.Tasks;
+using FluentAssertions;
+using NetCoreTests.Linq;
+using Xunit;
+
+namespace NetCoreTests
+{
+    public class AutoResetUniTaskCompletionSourceTest
+    {
+        [Fact]
+        public async Task SetResultAfterReturn()
+        {
+            var source1 = AutoResetUniTaskCompletionSource.Create();
+            source1.TrySetResult();
+            await source1.Task;
+            
+            source1.TrySetResult().Should().BeFalse();
+            
+            var source2 = AutoResetUniTaskCompletionSource.Create();
+            source2.TrySetResult();
+            await source2.Task;
+
+            source2.TrySetResult().Should().BeFalse();
+        }
+
+        [Fact]
+        public async Task SetCancelAfterReturn()
+        {
+            var source = AutoResetUniTaskCompletionSource.Create();
+            source.TrySetResult();
+            await source.Task;
+            
+            source.TrySetCanceled().Should().BeFalse();
+        }
+        
+        [Fact]
+        public async Task SetExceptionAfterReturn()
+        {
+            var source = AutoResetUniTaskCompletionSource.Create();
+            source.TrySetResult();
+            await source.Task;
+            
+            source.TrySetException(new UniTaskTestException()).Should().BeFalse();
+        }
+        
+        [Fact]
+        public async Task SetResultWithValueAfterReturn()
+        {
+            var source1 = AutoResetUniTaskCompletionSource<int>.Create();
+            source1.TrySetResult(100);
+            (await source1.Task).Should().Be(100);
+            
+            source1.TrySetResult(100).Should().BeFalse();
+            
+            var source2 = AutoResetUniTaskCompletionSource.Create();
+            source2.TrySetResult();
+            await source2.Task;
+            source2.TrySetResult().Should().BeFalse();
+        }
+
+        [Fact]
+        public async Task SetCancelWithValueAfterReturn()
+        {
+            var source = AutoResetUniTaskCompletionSource<int>.Create();
+            source.TrySetResult(100);
+            (await source.Task).Should().Be(100);
+            source.TrySetCanceled().Should().BeFalse();
+        }
+        
+        [Fact]
+        public async Task SetExceptionWithValueAfterReturn()
+        {
+            var source = AutoResetUniTaskCompletionSource<int>.Create();
+            source.TrySetResult(100);
+            (await source.Task).Should().Be(100);
+            source.TrySetException(new UniTaskTestException()).Should().BeFalse();
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTaskCompletionSource.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTaskCompletionSource.cs
index 67b882ee..6e401947 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTaskCompletionSource.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTaskCompletionSource.cs
@@ -328,6 +328,7 @@ static AutoResetUniTaskCompletionSource()
         }
 
         UniTaskCompletionSourceCore<AsyncUnit> core;
+        short version;
 
         AutoResetUniTaskCompletionSource()
         {
@@ -340,6 +341,7 @@ public static AutoResetUniTaskCompletionSource Create()
             {
                 result = new AutoResetUniTaskCompletionSource();
             }
+            result.version = result.core.Version;
             TaskTracker.TrackActiveTask(result, 2);
             return result;
         }
@@ -383,19 +385,19 @@ public UniTask Task
         [DebuggerHidden]
         public bool TrySetResult()
         {
-            return core.TrySetResult(AsyncUnit.Default);
+            return version == core.Version && core.TrySetResult(AsyncUnit.Default);
         }
 
         [DebuggerHidden]
         public bool TrySetCanceled(CancellationToken cancellationToken = default)
         {
-            return core.TrySetCanceled(cancellationToken);
+            return version == core.Version && core.TrySetCanceled(cancellationToken);
         }
 
         [DebuggerHidden]
         public bool TrySetException(Exception exception)
         {
-            return core.TrySetException(exception);
+            return version == core.Version && core.TrySetException(exception);
         }
 
         [DebuggerHidden]
@@ -409,7 +411,6 @@ public void GetResult(short token)
             {
                 TryReturn();
             }
-
         }
 
         [DebuggerHidden]
@@ -451,6 +452,7 @@ static AutoResetUniTaskCompletionSource()
         }
 
         UniTaskCompletionSourceCore<T> core;
+        short version;
 
         AutoResetUniTaskCompletionSource()
         {
@@ -463,6 +465,7 @@ public static AutoResetUniTaskCompletionSource<T> Create()
             {
                 result = new AutoResetUniTaskCompletionSource<T>();
             }
+            result.version = result.core.Version;
             TaskTracker.TrackActiveTask(result, 2);
             return result;
         }
@@ -506,19 +509,19 @@ public UniTask<T> Task
         [DebuggerHidden]
         public bool TrySetResult(T result)
         {
-            return core.TrySetResult(result);
+            return version == core.Version && core.TrySetResult(result);
         }
 
         [DebuggerHidden]
         public bool TrySetCanceled(CancellationToken cancellationToken = default)
         {
-            return core.TrySetCanceled(cancellationToken);
+            return version == core.Version && core.TrySetCanceled(cancellationToken);
         }
 
         [DebuggerHidden]
         public bool TrySetException(Exception exception)
         {
-            return core.TrySetException(exception);
+            return version == core.Version && core.TrySetException(exception);
         }
 
         [DebuggerHidden]
@@ -937,5 +940,5 @@ bool TrySignalCompletion(UniTaskStatus status)
             }
             return false;
         }
-    }
+   }
 }
\ No newline at end of file