Fix AssetBundleRequestAll

src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AssetBundleRequestAllAssets.cs

@@ -24,24 +24,17 @@ namespace Cysharp.Threading.Tasks
             return AwaitForAllAssets(asyncOperation, null, PlayerLoopTiming.Update, cancellationToken: cancellationToken);
         }
 
-        public static UniTask<UnityEngine.Object[]> AwaitForAllAssets(this AssetBundleRequest asyncOperation, bool handleImmediately, CancellationToken cancellationToken)
+        public static UniTask<UnityEngine.Object[]> AwaitForAllAssets(this AssetBundleRequest asyncOperation, CancellationToken cancellationToken, bool cancelImmediately)
         {
-            if (handleImmediately)
+            return AwaitForAllAssets(asyncOperation, progress: null, cancellationToken: cancellationToken, cancelImmediately: cancelImmediately);
-            {
-                Error.ThrowArgumentNullException(asyncOperation, nameof(asyncOperation));
-                if (cancellationToken.IsCancellationRequested) return UniTask.FromCanceled<UnityEngine.Object[]>(cancellationToken);
-                if (asyncOperation.isDone) return UniTask.FromResult(asyncOperation.allAssets);
-                return new UniTask<UnityEngine.Object[]>(AssetBundleRequestAllAssetsCallbackHandlerSource.Create(asyncOperation, cancellationToken, out var token), token);
-            }
-            return AwaitForAllAssets(asyncOperation, progress: null, cancellationToken: cancellationToken);
         }
 
-        public static UniTask<UnityEngine.Object[]> AwaitForAllAssets(this AssetBundleRequest asyncOperation, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken))
+        public static UniTask<UnityEngine.Object[]> AwaitForAllAssets(this AssetBundleRequest asyncOperation, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
         {
             Error.ThrowArgumentNullException(asyncOperation, nameof(asyncOperation));
             if (cancellationToken.IsCancellationRequested) return UniTask.FromCanceled<UnityEngine.Object[]>(cancellationToken);
             if (asyncOperation.isDone) return UniTask.FromResult(asyncOperation.allAssets);
-            return new UniTask<UnityEngine.Object[]>(AssetBundleRequestAllAssetsConfiguredSource.Create(asyncOperation, timing, progress, cancellationToken, out var token), token);
+            return new UniTask<UnityEngine.Object[]>(AssetBundleRequestAllAssetsConfiguredSource.Create(asyncOperation, timing, progress, cancellationToken, cancelImmediately, out var token), token);
         }
 
         public struct AssetBundleRequestAllAssetsAwaiter : ICriticalNotifyCompletion
@@ -93,108 +86,6 @@ namespace Cysharp.Threading.Tasks
             }
         }
 
-        sealed class AssetBundleRequestAllAssetsCallbackHandlerSource : IUniTaskSource<UnityEngine.Object[]>, ITaskPoolNode<AssetBundleRequestAllAssetsCallbackHandlerSource>
-        {
-            static TaskPool<AssetBundleRequestAllAssetsCallbackHandlerSource> pool;
-            AssetBundleRequestAllAssetsCallbackHandlerSource nextNode;
-            public ref AssetBundleRequestAllAssetsCallbackHandlerSource NextNode => ref nextNode;
-
-            static AssetBundleRequestAllAssetsCallbackHandlerSource()
-            {
-                TaskPool.RegisterSizeGetter(typeof(AssetBundleRequestConfiguredSource), () => pool.Size);
-            }
-
-            AssetBundleRequest asyncOperation;
-            CancellationToken cancellationToken;
-            CancellationTokenRegistration cancellationTokenRegistration;
-
-            UniTaskCompletionSourceCore<UnityEngine.Object[]> core;
-
-            AssetBundleRequestAllAssetsCallbackHandlerSource()
-            {
-            }
-
-            public static IUniTaskSource<UnityEngine.Object[]> Create(AssetBundleRequest asyncOperation, CancellationToken cancellationToken, out short token)
-            {
-                if (cancellationToken.IsCancellationRequested)
-                {
-                    return AutoResetUniTaskCompletionSource<UnityEngine.Object[]>.CreateFromCanceled(cancellationToken, out token);
-                }
-
-                if (!pool.TryPop(out var result))
-                {
-                    result = new AssetBundleRequestAllAssetsCallbackHandlerSource();
-                }
-
-                result.asyncOperation = asyncOperation;
-                result.cancellationToken = cancellationToken;
-
-                asyncOperation.completed += result.AsyncOperationCompletedHandler;
-                
-                if (cancellationToken.CanBeCanceled)
-                {
-                    result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
-                    {
-                        var source = (AssetBundleRequestAllAssetsCallbackHandlerSource)state;
-                        source.core.TrySetCanceled(source.cancellationToken);
-                    }, result);
-                }
-
-                TaskTracker.TrackActiveTask(result, 3);
-
-                token = result.core.Version;
-                return result;
-            }
-
-            public UnityEngine.Object[] GetResult(short token)
-            {
-                try
-                {
-                    return core.GetResult(token);
-                }
-                finally
-                {
-                    TryReturn();
-                }
-            }
-
-            void IUniTaskSource.GetResult(short token)
-            {
-                GetResult(token);
-            }
-
-            public UniTaskStatus GetStatus(short token)
-            {
-                return core.GetStatus(token);
-            }
-
-            public UniTaskStatus UnsafeGetStatus()
-            {
-                return core.UnsafeGetStatus();
-            }
-
-            public void OnCompleted(Action<object> continuation, object state, short token)
-            {
-                core.OnCompleted(continuation, state, token);
-            }
-
-            bool TryReturn()
-            {
-                TaskTracker.RemoveTracking(this);
-                core.Reset();
-                asyncOperation.completed -= AsyncOperationCompletedHandler;
-                asyncOperation = default;
-                cancellationToken = default;
-                cancellationTokenRegistration.Dispose();
-                return pool.TryPush(this);
-            }
-            
-            void AsyncOperationCompletedHandler(AsyncOperation _)
-            {
-                core.TrySetResult(asyncOperation.allAssets);
-            }
-        }
-
         sealed class AssetBundleRequestAllAssetsConfiguredSource : IUniTaskSource<UnityEngine.Object[]>, IPlayerLoopItem, ITaskPoolNode<AssetBundleRequestAllAssetsConfiguredSource>
         {
             static TaskPool<AssetBundleRequestAllAssetsConfiguredSource> pool;
@@ -209,14 +100,19 @@ namespace Cysharp.Threading.Tasks
             AssetBundleRequest asyncOperation;
             IProgress<float> progress;
             CancellationToken cancellationToken;
+            CancellationTokenRegistration cancellationTokenRegistration;
+            bool completed;
 
             UniTaskCompletionSourceCore<UnityEngine.Object[]> core;
 
+            Action<AsyncOperation> continuationAction;
+
             AssetBundleRequestAllAssetsConfiguredSource()
             {
+                continuationAction = Continuation;
             }
 
-            public static IUniTaskSource<UnityEngine.Object[]> Create(AssetBundleRequest asyncOperation, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, out short token)
+            public static IUniTaskSource<UnityEngine.Object[]> Create(AssetBundleRequest asyncOperation, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, bool cancelImmediately, out short token)
             {
                 if (cancellationToken.IsCancellationRequested)
                 {
@@ -231,6 +127,18 @@ namespace Cysharp.Threading.Tasks
                 result.asyncOperation = asyncOperation;
                 result.progress = progress;
                 result.cancellationToken = cancellationToken;
+                result.completed = false;
+
+                asyncOperation.completed += result.continuationAction;
+
+                if (cancelImmediately && cancellationToken.CanBeCanceled)
+                {
+                    result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
+                    {
+                        var source = (AssetBundleRequestAllAssetsConfiguredSource)state;
+                        source.core.TrySetCanceled(source.cancellationToken);
+                    }, result);
+                }
 
                 TaskTracker.TrackActiveTask(result, 3);
 
@@ -274,6 +182,12 @@ namespace Cysharp.Threading.Tasks
 
             public bool MoveNext()
             {
+                // Already completed 
+                if (completed || asyncOperation == null)
+                {
+                    return false;
+                }
+                
                 if (cancellationToken.IsCancellationRequested)
                 {
                     core.TrySetCanceled(cancellationToken);
@@ -301,8 +215,29 @@ namespace Cysharp.Threading.Tasks
                 asyncOperation = default;
                 progress = default;
                 cancellationToken = default;
+                cancellationTokenRegistration.Dispose();
                 return pool.TryPush(this);
             }
+            
+            void Continuation(AsyncOperation _)
+            {
+                if (completed)
+                {
+                    TryReturn();
+                }
+                else
+                {
+                    completed = true;
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        core.TrySetCanceled(cancellationToken);
+                    }
+                    else
+                    {
+                        core.TrySetResult(asyncOperation.allAssets);
+                    }
+                }
+            }
         }
     }
 }

src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs

@@ -171,14 +171,15 @@ namespace Cysharp.Threading.Tasks
 
             public bool MoveNext()
             {
-                if (cancellationToken.IsCancellationRequested)
+                // Already completed
+                if (completed || asyncOperation == null)
                 {
-                    core.TrySetCanceled(cancellationToken);
                     return false;
                 }
 
-                if (asyncOperation == null)
+                if (cancellationToken.IsCancellationRequested)
                 {
+                    core.TrySetCanceled(cancellationToken);
                     return false;
                 }
 
@@ -397,14 +398,15 @@ namespace Cysharp.Threading.Tasks
 
             public bool MoveNext()
             {
-                if (cancellationToken.IsCancellationRequested)
+                // Already completed
+                if (completed || asyncOperation == null)
                 {
-                    core.TrySetCanceled(cancellationToken);
                     return false;
                 }
 
-                if (asyncOperation == null)
+                if (cancellationToken.IsCancellationRequested)
                 {
+                    core.TrySetCanceled(cancellationToken);
                     return false;
                 }
 
@@ -624,14 +626,15 @@ namespace Cysharp.Threading.Tasks
 
             public bool MoveNext()
             {
-                if (cancellationToken.IsCancellationRequested)
+                // Already completed
+                if (completed || asyncOperation == null)
                 {
-                    core.TrySetCanceled(cancellationToken);
                     return false;
                 }
 
-                if (asyncOperation == null)
+                if (cancellationToken.IsCancellationRequested)
                 {
+                    core.TrySetCanceled(cancellationToken);
                     return false;
                 }
 
@@ -852,14 +855,15 @@ namespace Cysharp.Threading.Tasks
 
             public bool MoveNext()
             {
-                if (cancellationToken.IsCancellationRequested)
+                // Already completed
+                if (completed || asyncOperation == null)
                 {
-                    core.TrySetCanceled(cancellationToken);
                     return false;
                 }
 
-                if (asyncOperation == null)
+                if (cancellationToken.IsCancellationRequested)
                 {
+                    core.TrySetCanceled(cancellationToken);
                     return false;
                 }
 
@@ -1096,6 +1100,12 @@ namespace Cysharp.Threading.Tasks
 
             public bool MoveNext()
             {
+                // Already completed
+                if (completed || asyncOperation == null)
+                {
+                    return false;
+                }
+
                 if (cancellationToken.IsCancellationRequested)
                 {
                     asyncOperation.webRequest.Abort();
@@ -1103,11 +1113,6 @@ namespace Cysharp.Threading.Tasks
                     return false;
                 }
 
-                if (asyncOperation == null)
-                {
-                    return false;
-                }
-
                 if (progress != null)
                 {
                     progress.Report(asyncOperation.progress);

src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt

@@ -247,6 +247,12 @@ namespace Cysharp.Threading.Tasks
 
             public bool MoveNext()
             {
+                // Already completed
+                if (completed || asyncOperation == null)
+                {
+                    return false;
+                }
+
                 if (cancellationToken.IsCancellationRequested)
                 {
 <# if(IsUnityWebRequest(t)) { #>
@@ -256,11 +262,6 @@ namespace Cysharp.Threading.Tasks
                     return false;
                 }
 
-                if (asyncOperation == null)
-                {
-                    return false;
-                }
-
                 if (progress != null)
                 {
                     progress.Report(asyncOperation.progress);