Add UniTask.WithCancellation

master
neuecc 2020-07-30 08:10:16 +09:00
parent 44af123b6c
commit 1999d94b33
3 changed files with 209 additions and 0 deletions

View File

@ -38,6 +38,7 @@
..\UniTask\Assets\Plugins\UniTask\Runtime\Internal\ContinuationQueue.cs;
..\UniTask\Assets\Plugins\UniTask\Runtime\Internal\UnityWebRequestExtensions.cs;
..\UniTask\Assets\Plugins\UniTask\Runtime\UniTaskSynchronizationContext.cs;
..\UniTask\Assets\Plugins\UniTask\Runtime\CancellationTokenSourceExtensions.cs;
..\UniTask\Assets\Plugins\UniTask\Runtime\EnumeratorAsyncExtensions.cs;
..\UniTask\Assets\Plugins\UniTask\Runtime\PlayerLoopHelper.cs;

View File

@ -0,0 +1,40 @@
using Cysharp.Threading.Tasks;
using FluentAssertions;
using System;
using System.Collections.Generic;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Xunit;
namespace NetCoreTests
{
public class WithCancellationTest
{
[Fact]
public async Task Standard()
{
CancellationTokenSource cts = new CancellationTokenSource();
var v = await UniTask.Run(() => 10).WithCancellation(cts.Token);
v.Should().Be(10);
}
[Fact]
public async Task Cancel()
{
CancellationTokenSource cts = new CancellationTokenSource();
var t = UniTask.Create(async () =>
{
await Task.Delay(TimeSpan.FromSeconds(1));
return 10;
}).WithCancellation(cts.Token);
cts.Cancel();
(await Assert.ThrowsAsync<OperationCanceledException>(async () => await t)).CancellationToken.Should().Be(cts.Token);
}
}
}

View File

@ -189,6 +189,174 @@ namespace Cysharp.Threading.Tasks
return new AsyncLazy<T>(task);
}
/// <summary>
/// Ignore task result when cancel raised first.
/// </summary>
public static UniTask WithCancellation(this UniTask task, CancellationToken cancellationToken)
{
if (!cancellationToken.CanBeCanceled)
{
return task;
}
if (cancellationToken.IsCancellationRequested)
{
return UniTask.FromCanceled(cancellationToken);
}
if (task.Status.IsCompleted())
{
return task;
}
return new UniTask(new WithCancellationSource(task, cancellationToken), 0);
}
/// <summary>
/// Ignore task result when cancel raised first.
/// </summary>
public static UniTask<T> WithCancellation<T>(this UniTask<T> task, CancellationToken cancellationToken)
{
if (!cancellationToken.CanBeCanceled)
{
return task;
}
if (cancellationToken.IsCancellationRequested)
{
return UniTask.FromCanceled<T>(cancellationToken);
}
if (task.Status.IsCompleted())
{
return task;
}
return new UniTask<T>(new WithCancellationSource<T>(task, cancellationToken), 0);
}
sealed class WithCancellationSource : IUniTaskSource
{
static readonly Action<object> cancellationCallbackDelegate = CancellationCallback;
CancellationToken cancellationToken;
CancellationTokenRegistration tokenRegistration;
UniTaskCompletionSourceCore<AsyncUnit> core;
public WithCancellationSource(UniTask task, CancellationToken cancellationToken)
{
this.cancellationToken = cancellationToken;
this.tokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(cancellationCallbackDelegate, this);
RunTask(task).Forget();
}
async UniTaskVoid RunTask(UniTask task)
{
try
{
await task;
core.TrySetResult(AsyncUnit.Default);
}
catch (Exception ex)
{
core.TrySetException(ex);
}
finally
{
tokenRegistration.Dispose();
}
}
static void CancellationCallback(object state)
{
var self = (WithCancellationSource)state;
self.core.TrySetCanceled(self.cancellationToken);
}
public void GetResult(short token)
{
core.GetResult(token);
}
public UniTaskStatus GetStatus(short token)
{
return core.GetStatus(token);
}
public void OnCompleted(Action<object> continuation, object state, short token)
{
core.OnCompleted(continuation, state, token);
}
public UniTaskStatus UnsafeGetStatus()
{
return core.UnsafeGetStatus();
}
}
sealed class WithCancellationSource<T> : IUniTaskSource<T>
{
static readonly Action<object> cancellationCallbackDelegate = CancellationCallback;
CancellationToken cancellationToken;
CancellationTokenRegistration tokenRegistration;
UniTaskCompletionSourceCore<T> core;
public WithCancellationSource(UniTask<T> task, CancellationToken cancellationToken)
{
this.cancellationToken = cancellationToken;
this.tokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(cancellationCallbackDelegate, this);
RunTask(task).Forget();
}
async UniTaskVoid RunTask(UniTask<T> task)
{
try
{
core.TrySetResult(await task);
}
catch (Exception ex)
{
core.TrySetException(ex);
}
finally
{
tokenRegistration.Dispose();
}
}
static void CancellationCallback(object state)
{
var self = (WithCancellationSource<T>)state;
self.core.TrySetCanceled(self.cancellationToken);
}
void IUniTaskSource.GetResult(short token)
{
core.GetResult(token);
}
public T GetResult(short token)
{
return core.GetResult(token);
}
public UniTaskStatus GetStatus(short token)
{
return core.GetStatus(token);
}
public void OnCompleted(Action<object> continuation, object state, short token)
{
core.OnCompleted(continuation, state, token);
}
public UniTaskStatus UnsafeGetStatus()
{
return core.UnsafeGetStatus();
}
}
#if UNITY_2018_3_OR_NEWER
public static IEnumerator ToCoroutine<T>(this UniTask<T> task, Action<T> resultHandler = null, Action<Exception> exceptionHandler = null)