Merge pull request #536 from Cysharp/hadashiA/release-handle

Fix "Release handle when cancellation is requested"

src/UniTask/Assets/Plugins/UniTask/Runtime/External/Addressables/AddressablesAsyncExtensions.cs

@@ -7,6 +7,7 @@ using System;
 using System.Runtime.CompilerServices;
 using System.Runtime.ExceptionServices;
 using System.Threading;
+using UnityEngine.AddressableAssets;
 using UnityEngine.ResourceManagement.AsyncOperations;
 
 namespace Cysharp.Threading.Tasks
@@ -20,17 +21,12 @@ namespace Cysharp.Threading.Tasks
             return ToUniTask(handle).GetAwaiter();
         }
 
-        public static UniTask WithCancellation(this AsyncOperationHandle handle, CancellationToken cancellationToken)
+        public static UniTask WithCancellation(this AsyncOperationHandle handle, CancellationToken cancellationToken, bool cancelImmediately = false, bool autoReleaseWhenCancelled = false)
         {
-            return ToUniTask(handle, cancellationToken: cancellationToken);
+            return ToUniTask(handle, cancellationToken: cancellationToken, cancelImmediately: cancelImmediately, autoReleaseWhenCancelled: autoReleaseWhenCancelled);
         }
 
-        public static UniTask WithCancellation(this AsyncOperationHandle handle, CancellationToken cancellationToken, bool cancelImmediately)
+        public static UniTask ToUniTask(this AsyncOperationHandle handle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false, bool autoReleaseWhenCancelled = false)
-        {
-            return ToUniTask(handle, cancellationToken: cancellationToken, cancelImmediately: cancelImmediately);
-        }
-
-        public static UniTask ToUniTask(this AsyncOperationHandle handle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
         {
             if (cancellationToken.IsCancellationRequested) return UniTask.FromCanceled(cancellationToken);
 
@@ -49,7 +45,7 @@ namespace Cysharp.Threading.Tasks
                 return UniTask.CompletedTask;
             }
 
-            return new UniTask(AsyncOperationHandleConfiguredSource.Create(handle, timing, progress, cancellationToken, cancelImmediately, out var token), token);
+            return new UniTask(AsyncOperationHandleConfiguredSource.Create(handle, timing, progress, cancellationToken, cancelImmediately, autoReleaseWhenCancelled, out var token), token);
         }
 
         public struct AsyncOperationHandleAwaiter : ICriticalNotifyCompletion
@@ -108,21 +104,22 @@ namespace Cysharp.Threading.Tasks
                 TaskPool.RegisterSizeGetter(typeof(AsyncOperationHandleConfiguredSource), () => pool.Size);
             }
 
-            readonly Action<AsyncOperationHandle> continuationAction;
+            readonly Action<AsyncOperationHandle> completedCallback;
             AsyncOperationHandle handle;
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
             IProgress<float> progress;
+            bool autoReleaseWhenCancelled;
             bool completed;
 
             UniTaskCompletionSourceCore<AsyncUnit> core;
 
             AsyncOperationHandleConfiguredSource()
             {
-                continuationAction = Continuation;
+                completedCallback = HandleCompleted;
             }
 
-            public static IUniTaskSource Create(AsyncOperationHandle handle, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, bool cancelImmediately, out short token)
+            public static IUniTaskSource Create(AsyncOperationHandle handle, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, bool cancelImmediately, bool autoReleaseWhenCancelled, out short token)
             {
                 if (cancellationToken.IsCancellationRequested)
                 {
@@ -138,12 +135,17 @@ namespace Cysharp.Threading.Tasks
                 result.progress = progress;
                 result.cancellationToken = cancellationToken;
                 result.completed = false;
+                result.autoReleaseWhenCancelled = autoReleaseWhenCancelled;
                 
                 if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
                     result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
                     {
                         var promise = (AsyncOperationHandleConfiguredSource)state;
+                        if (promise.autoReleaseWhenCancelled && promise.handle.IsValid())
+                        {
+                            Addressables.Release(promise.handle);
+                        }
                         promise.core.TrySetCanceled(promise.cancellationToken);
                     }, result);
                 }
@@ -152,15 +154,18 @@ namespace Cysharp.Threading.Tasks
 
                 PlayerLoopHelper.AddAction(timing, result);
 
-                handle.Completed += result.continuationAction;
+                handle.Completed += result.completedCallback;
 
                 token = result.core.Version;
                 return result;
             }
 
-            void Continuation(AsyncOperationHandle _)
+            void HandleCompleted(AsyncOperationHandle _)
             {
-                handle.Completed -= continuationAction;
+                if (handle.IsValid())
+                {
+                    handle.Completed -= completedCallback;
+                }
 
                 if (completed)
                 {
@@ -171,6 +176,10 @@ namespace Cysharp.Threading.Tasks
                     completed = true;
                     if (cancellationToken.IsCancellationRequested)
                     {
+                        if (autoReleaseWhenCancelled && handle.IsValid())
+                        {
+                            Addressables.Release(handle);
+                        }
                         core.TrySetCanceled(cancellationToken);
                     }
                     else if (handle.Status == AsyncOperationStatus.Failed)
@@ -215,6 +224,10 @@ namespace Cysharp.Threading.Tasks
                 if (cancellationToken.IsCancellationRequested)
                 {
                     completed = true;
+                    if (autoReleaseWhenCancelled && handle.IsValid())
+                    {
+                        Addressables.Release(handle);
+                    }
                     core.TrySetCanceled(cancellationToken);
                     return false;
                 }
@@ -248,17 +261,12 @@ namespace Cysharp.Threading.Tasks
             return ToUniTask(handle).GetAwaiter();
         }
 
-        public static UniTask<T> WithCancellation<T>(this AsyncOperationHandle<T> handle, CancellationToken cancellationToken)
+        public static UniTask<T> WithCancellation<T>(this AsyncOperationHandle<T> handle, CancellationToken cancellationToken, bool cancelImmediately = false, bool autoReleaseWhenCancelled = false)
         {
-            return ToUniTask(handle, cancellationToken: cancellationToken);
+            return ToUniTask(handle, cancellationToken: cancellationToken, cancelImmediately: cancelImmediately, autoReleaseWhenCancelled: autoReleaseWhenCancelled);
         }
 
-        public static UniTask<T> WithCancellation<T>(this AsyncOperationHandle<T> handle, CancellationToken cancellationToken, bool cancelImmediately)
+        public static UniTask<T> ToUniTask<T>(this AsyncOperationHandle<T> handle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false, bool autoReleaseWhenCancelled = false)
-        {
-            return ToUniTask(handle, cancellationToken: cancellationToken, cancelImmediately: cancelImmediately);
-        }
-
-        public static UniTask<T> ToUniTask<T>(this AsyncOperationHandle<T> handle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
         {
             if (cancellationToken.IsCancellationRequested) return UniTask.FromCanceled<T>(cancellationToken);
 
@@ -276,7 +284,7 @@ namespace Cysharp.Threading.Tasks
                 return UniTask.FromResult(handle.Result);
             }
 
-            return new UniTask<T>(AsyncOperationHandleConfiguredSource<T>.Create(handle, timing, progress, cancellationToken, cancelImmediately, out var token), token);
+            return new UniTask<T>(AsyncOperationHandleConfiguredSource<T>.Create(handle, timing, progress, cancellationToken, cancelImmediately, autoReleaseWhenCancelled, out var token), token);
         }
 
         sealed class AsyncOperationHandleConfiguredSource<T> : IUniTaskSource<T>, IPlayerLoopItem, ITaskPoolNode<AsyncOperationHandleConfiguredSource<T>>
@@ -290,21 +298,22 @@ namespace Cysharp.Threading.Tasks
                 TaskPool.RegisterSizeGetter(typeof(AsyncOperationHandleConfiguredSource<T>), () => pool.Size);
             }
 
-            readonly Action<AsyncOperationHandle<T>> continuationAction;
+            readonly Action<AsyncOperationHandle<T>> completedCallback;
             AsyncOperationHandle<T> handle;
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
             IProgress<float> progress;
+            bool autoReleaseWhenCancelled;
             bool completed;
 
             UniTaskCompletionSourceCore<T> core;
 
             AsyncOperationHandleConfiguredSource()
             {
-                continuationAction = Continuation;
+                completedCallback = HandleCompleted;
             }
 
-            public static IUniTaskSource<T> Create(AsyncOperationHandle<T> handle, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, bool cancelImmediately, out short token)
+            public static IUniTaskSource<T> Create(AsyncOperationHandle<T> handle, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, bool cancelImmediately, bool autoReleaseWhenCancelled, out short token)
             {
                 if (cancellationToken.IsCancellationRequested)
                 {
@@ -320,12 +329,17 @@ namespace Cysharp.Threading.Tasks
                 result.cancellationToken = cancellationToken;
                 result.completed = false;
                 result.progress = progress;
+                result.autoReleaseWhenCancelled = autoReleaseWhenCancelled;
                 
                 if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
                     result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
                     {
                         var promise = (AsyncOperationHandleConfiguredSource<T>)state;
+                        if (promise.autoReleaseWhenCancelled && promise.handle.IsValid())
+                        {
+                            Addressables.Release(promise.handle);
+                        }
                         promise.core.TrySetCanceled(promise.cancellationToken);
                     }, result);
                 }
@@ -334,15 +348,18 @@ namespace Cysharp.Threading.Tasks
 
                 PlayerLoopHelper.AddAction(timing, result);
 
-                handle.Completed += result.continuationAction;
+                handle.Completed += result.completedCallback;
 
                 token = result.core.Version;
                 return result;
             }
 
-            void Continuation(AsyncOperationHandle<T> argHandle)
+            void HandleCompleted(AsyncOperationHandle<T> argHandle)
             {
-                handle.Completed -= continuationAction;
+                if (handle.IsValid())
+                {
+                    handle.Completed -= completedCallback;
+                }
 
                 if (completed)
                 {
@@ -353,6 +370,10 @@ namespace Cysharp.Threading.Tasks
                     completed = true;
                     if (cancellationToken.IsCancellationRequested)
                     {
+                        if (autoReleaseWhenCancelled && handle.IsValid())
+                        {
+                            Addressables.Release(handle);
+                        }
                         core.TrySetCanceled(cancellationToken);
                     }
                     else if (argHandle.Status == AsyncOperationStatus.Failed)
@@ -402,6 +423,10 @@ namespace Cysharp.Threading.Tasks
                 if (cancellationToken.IsCancellationRequested)
                 {
                     completed = true;
+                    if (autoReleaseWhenCancelled && handle.IsValid())
+                    {
+                        Addressables.Release(handle);
+                    }
                     core.TrySetCanceled(cancellationToken);
                     return false;
                 }

src/UniTask/Assets/Plugins/UniTask/Runtime/External/Addressables/UniTask.Addressables.asmdef

@@ -2,7 +2,8 @@
     "name": "UniTask.Addressables",
     "references": [
         "UniTask",
-        "Unity.ResourceManager"
+        "Unity.ResourceManager",
+        "Unity.Addressables"
     ],
     "includePlatforms": [],
     "excludePlatforms": [],