WithCancellation

This commit is contained in:
neuecc
2020-05-08 03:48:46 +09:00
parent 61b798b6e9
commit 856a049dd0
4 changed files with 64 additions and 5 deletions

View File

@@ -1,4 +1,5 @@
using System.Threading;
using System.Runtime.InteropServices;
using System.Threading;
namespace Cysharp.Threading.Tasks
{
@@ -17,4 +18,55 @@ namespace Cysharp.Threading.Tasks
{
UniTask DisposeAsync();
}
public static class UniTaskAsyncEnumerableExtensions
{
public static UniTaskCancelableAsyncEnumerable<T> WithCancellation<T>(this IUniTaskAsyncEnumerable<T> source, CancellationToken cancellationToken)
{
return new UniTaskCancelableAsyncEnumerable<T>(source, cancellationToken);
}
}
[StructLayout(LayoutKind.Auto)]
public readonly struct UniTaskCancelableAsyncEnumerable<T>
{
private readonly IUniTaskAsyncEnumerable<T> enumerable;
private readonly CancellationToken cancellationToken;
internal UniTaskCancelableAsyncEnumerable(IUniTaskAsyncEnumerable<T> enumerable, CancellationToken cancellationToken)
{
this.enumerable = enumerable;
this.cancellationToken = cancellationToken;
}
public Enumerator GetAsyncEnumerator()
{
cancellationToken.ThrowIfCancellationRequested();
return new Enumerator(enumerable.GetAsyncEnumerator(cancellationToken));
}
[StructLayout(LayoutKind.Auto)]
public readonly struct Enumerator
{
private readonly IUniTaskAsyncEnumerator<T> enumerator;
internal Enumerator(IUniTaskAsyncEnumerator<T> enumerator)
{
this.enumerator = enumerator;
}
public T Current => enumerator.Current;
public UniTask<bool> MoveNextAsync()
{
return enumerator.MoveNextAsync();
}
public UniTask DisposeAsync()
{
return enumerator.DisposeAsync();
}
}
}
}

View File

@@ -55,7 +55,7 @@ namespace Cysharp.Threading.Tasks.Linq
public UniTask<bool> MoveNextAsync()
{
if (cancellationToken.IsCancellationRequested) return CompletedTasks.False;
cancellationToken.ThrowIfCancellationRequested();
current++;

View File

@@ -50,7 +50,7 @@ namespace Cysharp.Threading.Tasks.Linq
public UniTask<bool> MoveNextAsync()
{
if (cancellationToken.IsCancellationRequested) return CompletedTasks.False;
cancellationToken.ThrowIfCancellationRequested();
if (remaining-- != 0)
{