From 370425578f40d5728078f560982546f51b3f8aec Mon Sep 17 00:00:00 2001 From: hadashiA Date: Thu, 2 Nov 2023 12:21:59 +0900 Subject: [PATCH] Fix AssetBundleRequestAll --- ...cExtensions.AssetBundleRequestAllAssets.cs | 163 ++++++------------ .../UnityAsyncExtensions.AsyncGPUReadback.cs | 2 +- .../UniTask/Runtime/UnityAsyncExtensions.cs | 39 +++-- .../UniTask/Runtime/UnityAsyncExtensions.tt | 11 +- 4 files changed, 78 insertions(+), 137 deletions(-) diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AssetBundleRequestAllAssets.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AssetBundleRequestAllAssets.cs index f6fccc0..5d34692 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AssetBundleRequestAllAssets.cs +++ b/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 AwaitForAllAssets(this AssetBundleRequest asyncOperation, bool handleImmediately, CancellationToken cancellationToken) + public static UniTask AwaitForAllAssets(this AssetBundleRequest asyncOperation, CancellationToken cancellationToken, bool cancelImmediately) { - if (handleImmediately) - { - Error.ThrowArgumentNullException(asyncOperation, nameof(asyncOperation)); - if (cancellationToken.IsCancellationRequested) return UniTask.FromCanceled(cancellationToken); - if (asyncOperation.isDone) return UniTask.FromResult(asyncOperation.allAssets); - return new UniTask(AssetBundleRequestAllAssetsCallbackHandlerSource.Create(asyncOperation, cancellationToken, out var token), token); - } - return AwaitForAllAssets(asyncOperation, progress: null, cancellationToken: cancellationToken); + return AwaitForAllAssets(asyncOperation, progress: null, cancellationToken: cancellationToken, cancelImmediately: cancelImmediately); } - public static UniTask AwaitForAllAssets(this AssetBundleRequest asyncOperation, IProgress progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken)) + public static UniTask AwaitForAllAssets(this AssetBundleRequest asyncOperation, IProgress progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false) { Error.ThrowArgumentNullException(asyncOperation, nameof(asyncOperation)); if (cancellationToken.IsCancellationRequested) return UniTask.FromCanceled(cancellationToken); if (asyncOperation.isDone) return UniTask.FromResult(asyncOperation.allAssets); - return new UniTask(AssetBundleRequestAllAssetsConfiguredSource.Create(asyncOperation, timing, progress, cancellationToken, out var token), token); + return new UniTask(AssetBundleRequestAllAssetsConfiguredSource.Create(asyncOperation, timing, progress, cancellationToken, cancelImmediately, out var token), token); } public struct AssetBundleRequestAllAssetsAwaiter : ICriticalNotifyCompletion @@ -92,108 +85,6 @@ namespace Cysharp.Threading.Tasks asyncOperation.completed += continuationAction; } } - - sealed class AssetBundleRequestAllAssetsCallbackHandlerSource : IUniTaskSource, ITaskPoolNode - { - static TaskPool pool; - AssetBundleRequestAllAssetsCallbackHandlerSource nextNode; - public ref AssetBundleRequestAllAssetsCallbackHandlerSource NextNode => ref nextNode; - - static AssetBundleRequestAllAssetsCallbackHandlerSource() - { - TaskPool.RegisterSizeGetter(typeof(AssetBundleRequestConfiguredSource), () => pool.Size); - } - - AssetBundleRequest asyncOperation; - CancellationToken cancellationToken; - CancellationTokenRegistration cancellationTokenRegistration; - - UniTaskCompletionSourceCore core; - - AssetBundleRequestAllAssetsCallbackHandlerSource() - { - } - - public static IUniTaskSource Create(AssetBundleRequest asyncOperation, CancellationToken cancellationToken, out short token) - { - if (cancellationToken.IsCancellationRequested) - { - return AutoResetUniTaskCompletionSource.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 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, IPlayerLoopItem, ITaskPoolNode { @@ -209,14 +100,19 @@ namespace Cysharp.Threading.Tasks AssetBundleRequest asyncOperation; IProgress progress; CancellationToken cancellationToken; + CancellationTokenRegistration cancellationTokenRegistration; + bool completed; UniTaskCompletionSourceCore core; + Action continuationAction; + AssetBundleRequestAllAssetsConfiguredSource() { + continuationAction = Continuation; } - public static IUniTaskSource Create(AssetBundleRequest asyncOperation, PlayerLoopTiming timing, IProgress progress, CancellationToken cancellationToken, out short token) + public static IUniTaskSource Create(AssetBundleRequest asyncOperation, PlayerLoopTiming timing, IProgress 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); + } + } + } } } } diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AsyncGPUReadback.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AsyncGPUReadback.cs index 0faf0be..5d73dc1 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AsyncGPUReadback.cs +++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.AsyncGPUReadback.cs @@ -30,7 +30,7 @@ namespace Cysharp.Threading.Tasks if (asyncOperation.done) return UniTask.FromResult(asyncOperation); return new UniTask(AsyncGPUReadbackRequestAwaiterConfiguredSource.Create(asyncOperation, timing, cancellationToken, cancelImmediately, out var token), token); } - + sealed class AsyncGPUReadbackRequestAwaiterConfiguredSource : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode { static TaskPool pool; diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs index fb84fca..7d86361 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.cs +++ b/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); diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt b/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt index 86dd72a..6c6ff91 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UnityAsyncExtensions.tt +++ b/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);