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

Fix "Release handle when cancellation is requested"
pull/541/head
Yoshifumi Kawai 2024-01-25 20:26:13 +09:00 committed by GitHub
commit 4b0bd3b509
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 57 additions and 31 deletions

View File

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

View File

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