Merge pull request #557 from Cysharp/hadashiA/fix-cancel-immediate

Fix unintended returning to pool with cancelImmediately
Cysharp · Mar 28, 2024 · b724a2a · b724a2a
2 parents 01c8fad + 938ddd5
commit b724a2a
Show file tree

Hide file tree

Showing 7 changed files with 257 additions and 146 deletions.
diff --git a/...iTask/Assets/Plugins/UniTask/Runtime/External/Addressables/AddressablesAsyncExtensions.cs b/...iTask/Assets/Plugins/UniTask/Runtime/External/Addressables/AddressablesAsyncExtensions.cs
@@ -110,6 +110,7 @@ static AsyncOperationHandleConfiguredSource()
             CancellationTokenRegistration cancellationTokenRegistration;
             IProgress<float> progress;
             bool autoReleaseWhenCanceled;
+            bool cancelImmediately;
             bool completed;
 
             UniTaskCompletionSourceCore<AsyncUnit> core;
@@ -134,8 +135,9 @@ public static IUniTaskSource Create(AsyncOperationHandle handle, PlayerLoopTimin
                 result.handle = handle;
                 result.progress = progress;
                 result.cancellationToken = cancellationToken;
-                result.completed = false;
+                result.cancelImmediately = cancelImmediately;
                 result.autoReleaseWhenCanceled = autoReleaseWhenCanceled;
+                result.completed = false;
 
                 if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
@@ -169,33 +171,41 @@ void HandleCompleted(AsyncOperationHandle _)
 
                 if (completed)
                 {
-                    TryReturn();
+                    return;
                 }
-                else
+
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
                 {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        if (autoReleaseWhenCanceled && handle.IsValid())
-                        {
-                            Addressables.Release(handle);
-                        }
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else if (handle.Status == AsyncOperationStatus.Failed)
-                    {
-                        core.TrySetException(handle.OperationException);
-                    }
-                    else
+                    if (autoReleaseWhenCanceled && handle.IsValid())
                     {
-                        core.TrySetResult(AsyncUnit.Default);
+                        Addressables.Release(handle);
                     }
+                    core.TrySetCanceled(cancellationToken);
+                }
+                else if (handle.Status == AsyncOperationStatus.Failed)
+                {
+                    core.TrySetException(handle.OperationException);
+                }
+                else
+                {
+                    core.TrySetResult(AsyncUnit.Default);
                 }
             }
 
             public void GetResult(short token)
             {
-                core.GetResult(token);
+                try
+                {
+                    core.GetResult(token);
+                }
+                finally
+                {
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
+                }
             }
 
             public UniTaskStatus GetStatus(short token)
@@ -217,7 +227,6 @@ public bool MoveNext()
             {
                 if (completed)
                 {
-                    TryReturn();
                     return false;
                 }
 
@@ -304,6 +313,7 @@ static AsyncOperationHandleConfiguredSource()
             CancellationTokenRegistration cancellationTokenRegistration;
             IProgress<float> progress;
             bool autoReleaseWhenCanceled;
+            bool cancelImmediately;
             bool completed;
 
             UniTaskCompletionSourceCore<T> core;
@@ -330,6 +340,7 @@ public static IUniTaskSource<T> Create(AsyncOperationHandle<T> handle, PlayerLoo
                 result.completed = false;
                 result.progress = progress;
                 result.autoReleaseWhenCanceled = autoReleaseWhenCanceled;
+                result.cancelImmediately = cancelImmediately;
 
                 if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
@@ -363,33 +374,38 @@ void HandleCompleted(AsyncOperationHandle<T> argHandle)
 
                 if (completed)
                 {
-                    TryReturn();
+                    return;
                 }
-                else
+                completed = true;
+                if (cancellationToken.IsCancellationRequested)
                 {
-                    completed = true;
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        if (autoReleaseWhenCanceled && handle.IsValid())
-                        {
-                            Addressables.Release(handle);
-                        }
-                        core.TrySetCanceled(cancellationToken);
-                    }
-                    else if (argHandle.Status == AsyncOperationStatus.Failed)
-                    {
-                        core.TrySetException(argHandle.OperationException);
-                    }
-                    else
+                    if (autoReleaseWhenCanceled && handle.IsValid())
                     {
-                        core.TrySetResult(argHandle.Result);
+                        Addressables.Release(handle);
                     }
+                    core.TrySetCanceled(cancellationToken);
+                }
+                else if (argHandle.Status == AsyncOperationStatus.Failed)
+                {
+                    core.TrySetException(argHandle.OperationException);
+                }
+                else
+                {
+                    core.TrySetResult(argHandle.Result);
                 }
             }
 
             public T GetResult(short token)
             {
-                return core.GetResult(token);
+                try
+                {
+                    return core.GetResult(token);
+                }
+                finally
+                {
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                        TryReturn();
+                }
             }
 
             void IUniTaskSource.GetResult(short token)
@@ -416,7 +432,6 @@ public bool MoveNext()
             {
                 if (completed)
                 {
-                    TryReturn();
                     return false;
                 }