TaskExtensions
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
namespace System.ClientModel.Internal
{
[System.Runtime.CompilerServices.NullableContext(1)]
[System.Runtime.CompilerServices.Nullable(0)]
internal static class TaskExtensions
{
[System.Runtime.CompilerServices.Nullable(0)]
public readonly struct Enumerable<[System.Runtime.CompilerServices.Nullable(2)] T> : IEnumerable<T>, IEnumerable
{
private readonly IAsyncEnumerable<T> _asyncEnumerable;
public Enumerable(IAsyncEnumerable<T> asyncEnumerable)
{
_asyncEnumerable = asyncEnumerable;
}
[return: System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
public Enumerator<T> GetEnumerator()
{
return new Enumerator<T>(_asyncEnumerable.GetAsyncEnumerator(default(CancellationToken)));
}
IEnumerator<T> IEnumerable<T>.GetEnumerator()
{
return new Enumerator<T>(_asyncEnumerable.GetAsyncEnumerator(default(CancellationToken)));
}
IEnumerator IEnumerable.GetEnumerator()
{
return GetEnumerator();
}
}
[System.Runtime.CompilerServices.Nullable(0)]
public readonly struct Enumerator<[System.Runtime.CompilerServices.Nullable(2)] T> : IEnumerator<T>, IEnumerator, IDisposable
{
private readonly IAsyncEnumerator<T> _asyncEnumerator;
public T Current => _asyncEnumerator.Current;
[System.Runtime.CompilerServices.Nullable(2)]
object IEnumerator.Current {
[System.Runtime.CompilerServices.NullableContext(2)]
get {
return Current;
}
}
public Enumerator(IAsyncEnumerator<T> asyncEnumerator)
{
_asyncEnumerator = asyncEnumerator;
}
public bool MoveNext()
{
return TaskExtensions.EnsureCompleted<bool>(_asyncEnumerator.MoveNextAsync());
}
public void Reset()
{
throw new NotSupportedException($"{GetType()}""{_asyncEnumerator.GetType()}""");
}
public void Dispose()
{
_asyncEnumerator.DisposeAsync().EnsureCompleted();
}
}
[System.Runtime.CompilerServices.NullableContext(0)]
public readonly struct WithCancellationTaskAwaitable
{
private readonly CancellationToken _cancellationToken;
private readonly ConfiguredTaskAwaitable _awaitable;
[System.Runtime.CompilerServices.NullableContext(1)]
public WithCancellationTaskAwaitable(Task task, CancellationToken cancellationToken)
{
_awaitable = task.ConfigureAwait(false);
_cancellationToken = cancellationToken;
}
public WithCancellationTaskAwaiter GetAwaiter()
{
return new WithCancellationTaskAwaiter(_awaitable.GetAwaiter(), _cancellationToken);
}
}
[System.Runtime.CompilerServices.NullableContext(0)]
public readonly struct WithCancellationTaskAwaitable<[System.Runtime.CompilerServices.Nullable(2)] T>
{
private readonly CancellationToken _cancellationToken;
[System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
private readonly ConfiguredTaskAwaitable<T> _awaitable;
[System.Runtime.CompilerServices.NullableContext(1)]
public WithCancellationTaskAwaitable(Task<T> task, CancellationToken cancellationToken)
{
_awaitable = task.ConfigureAwait(false);
_cancellationToken = cancellationToken;
}
[return: System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
public WithCancellationTaskAwaiter<T> GetAwaiter()
{
return new WithCancellationTaskAwaiter<T>(_awaitable.GetAwaiter(), _cancellationToken);
}
}
[System.Runtime.CompilerServices.NullableContext(0)]
public readonly struct WithCancellationValueTaskAwaitable<[System.Runtime.CompilerServices.Nullable(2)] T>
{
private readonly CancellationToken _cancellationToken;
[System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
private readonly ConfiguredValueTaskAwaitable<T> _awaitable;
public WithCancellationValueTaskAwaitable([System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})] ValueTask<T> task, CancellationToken cancellationToken)
{
_awaitable = task.ConfigureAwait(false);
_cancellationToken = cancellationToken;
}
[return: System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
public WithCancellationValueTaskAwaiter<T> GetAwaiter()
{
return new WithCancellationValueTaskAwaiter<T>(_awaitable.GetAwaiter(), _cancellationToken);
}
}
[System.Runtime.CompilerServices.NullableContext(0)]
public readonly struct WithCancellationTaskAwaiter : ICriticalNotifyCompletion, INotifyCompletion
{
private readonly CancellationToken _cancellationToken;
private readonly ConfiguredTaskAwaitable.ConfiguredTaskAwaiter _taskAwaiter;
public bool IsCompleted {
get {
if (!_taskAwaiter.IsCompleted)
return _cancellationToken.IsCancellationRequested;
return true;
}
}
public WithCancellationTaskAwaiter(ConfiguredTaskAwaitable.ConfiguredTaskAwaiter awaiter, CancellationToken cancellationToken)
{
_taskAwaiter = awaiter;
_cancellationToken = cancellationToken;
}
[System.Runtime.CompilerServices.NullableContext(1)]
public void OnCompleted(Action continuation)
{
_taskAwaiter.OnCompleted(WrapContinuation(ref continuation));
}
[System.Runtime.CompilerServices.NullableContext(1)]
public void UnsafeOnCompleted(Action continuation)
{
_taskAwaiter.UnsafeOnCompleted(WrapContinuation(ref continuation));
}
public void GetResult()
{
ConfiguredTaskAwaitable.ConfiguredTaskAwaiter taskAwaiter = _taskAwaiter;
if (!taskAwaiter.IsCompleted)
_cancellationToken.ThrowIfCancellationRequested();
taskAwaiter = _taskAwaiter;
taskAwaiter.GetResult();
}
[System.Runtime.CompilerServices.NullableContext(1)]
private Action WrapContinuation([In] [System.Runtime.CompilerServices.IsReadOnly] ref Action originalContinuation)
{
if (!_cancellationToken.CanBeCanceled)
return originalContinuation;
return new WithCancellationContinuationWrapper(originalContinuation, _cancellationToken).Continuation;
}
}
[System.Runtime.CompilerServices.Nullable(0)]
public readonly struct WithCancellationTaskAwaiter<[System.Runtime.CompilerServices.Nullable(2)] T> : ICriticalNotifyCompletion, INotifyCompletion
{
private readonly CancellationToken _cancellationToken;
[System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
private readonly ConfiguredTaskAwaitable<T>.ConfiguredTaskAwaiter _taskAwaiter;
public bool IsCompleted {
get {
if (!_taskAwaiter.IsCompleted)
return _cancellationToken.IsCancellationRequested;
return true;
}
}
public WithCancellationTaskAwaiter([System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})] ConfiguredTaskAwaitable<T>.ConfiguredTaskAwaiter awaiter, CancellationToken cancellationToken)
{
_taskAwaiter = awaiter;
_cancellationToken = cancellationToken;
}
public void OnCompleted(Action continuation)
{
_taskAwaiter.OnCompleted(WrapContinuation(ref continuation));
}
public void UnsafeOnCompleted(Action continuation)
{
_taskAwaiter.UnsafeOnCompleted(WrapContinuation(ref continuation));
}
public T GetResult()
{
ConfiguredTaskAwaitable<T>.ConfiguredTaskAwaiter taskAwaiter = _taskAwaiter;
if (!taskAwaiter.IsCompleted)
_cancellationToken.ThrowIfCancellationRequested();
taskAwaiter = _taskAwaiter;
return taskAwaiter.GetResult();
}
private Action WrapContinuation([In] [System.Runtime.CompilerServices.IsReadOnly] ref Action originalContinuation)
{
if (!_cancellationToken.CanBeCanceled)
return originalContinuation;
return new WithCancellationContinuationWrapper(originalContinuation, _cancellationToken).Continuation;
}
}
[System.Runtime.CompilerServices.Nullable(0)]
public readonly struct WithCancellationValueTaskAwaiter<[System.Runtime.CompilerServices.Nullable(2)] T> : ICriticalNotifyCompletion, INotifyCompletion
{
private readonly CancellationToken _cancellationToken;
[System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
private readonly ConfiguredValueTaskAwaitable<T>.ConfiguredValueTaskAwaiter _taskAwaiter;
public bool IsCompleted {
get {
if (!_taskAwaiter.IsCompleted)
return _cancellationToken.IsCancellationRequested;
return true;
}
}
public WithCancellationValueTaskAwaiter([System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})] ConfiguredValueTaskAwaitable<T>.ConfiguredValueTaskAwaiter awaiter, CancellationToken cancellationToken)
{
_taskAwaiter = awaiter;
_cancellationToken = cancellationToken;
}
public void OnCompleted(Action continuation)
{
_taskAwaiter.OnCompleted(WrapContinuation(ref continuation));
}
public void UnsafeOnCompleted(Action continuation)
{
_taskAwaiter.UnsafeOnCompleted(WrapContinuation(ref continuation));
}
public T GetResult()
{
if (!_taskAwaiter.IsCompleted)
_cancellationToken.ThrowIfCancellationRequested();
return _taskAwaiter.GetResult();
}
private Action WrapContinuation([In] [System.Runtime.CompilerServices.IsReadOnly] ref Action originalContinuation)
{
if (!_cancellationToken.CanBeCanceled)
return originalContinuation;
return new WithCancellationContinuationWrapper(originalContinuation, _cancellationToken).Continuation;
}
}
[System.Runtime.CompilerServices.Nullable(0)]
private class WithCancellationContinuationWrapper
{
[System.Runtime.CompilerServices.Nullable(2)]
private Action _originalContinuation;
private readonly CancellationTokenRegistration _registration;
public Action Continuation { get; }
public WithCancellationContinuationWrapper(Action originalContinuation, CancellationToken cancellationToken)
{
Action callback = ContinuationImplementation;
_originalContinuation = originalContinuation;
_registration = cancellationToken.Register(callback);
Continuation = callback;
}
private void ContinuationImplementation()
{
Action action = Interlocked.Exchange(ref _originalContinuation, null);
if (action != null) {
_registration.Dispose();
action();
}
}
}
public static WithCancellationTaskAwaitable AwaitWithCancellation(this Task task, CancellationToken cancellationToken)
{
return new WithCancellationTaskAwaitable(task, cancellationToken);
}
[return: System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
public static WithCancellationTaskAwaitable<T> AwaitWithCancellation<[System.Runtime.CompilerServices.Nullable(2)] T>(this Task<T> task, CancellationToken cancellationToken)
{
return new WithCancellationTaskAwaitable<T>(task, cancellationToken);
}
[System.Runtime.CompilerServices.NullableContext(2)]
[return: System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
public static WithCancellationValueTaskAwaitable<T> AwaitWithCancellation<T>([System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})] this ValueTask<T> task, CancellationToken cancellationToken)
{
return new WithCancellationValueTaskAwaitable<T>(task, cancellationToken);
}
public static T EnsureCompleted<[System.Runtime.CompilerServices.Nullable(2)] T>(this Task<T> task)
{
return task.GetAwaiter().GetResult();
}
public static void EnsureCompleted(this Task task)
{
task.GetAwaiter().GetResult();
}
public static T EnsureCompleted<[System.Runtime.CompilerServices.Nullable(2)] T>([System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})] this ValueTask<T> task)
{
return task.GetAwaiter().GetResult();
}
public static void EnsureCompleted(this ValueTask task)
{
task.GetAwaiter().GetResult();
}
[return: System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
public static Enumerable<T> EnsureSyncEnumerable<[System.Runtime.CompilerServices.Nullable(2)] T>(this IAsyncEnumerable<T> asyncEnumerable)
{
return new Enumerable<T>(asyncEnumerable);
}
[System.Runtime.CompilerServices.NullableContext(2)]
[return: System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})]
public static ConfiguredValueTaskAwaitable<T> EnsureCompleted<T>([System.Runtime.CompilerServices.Nullable(new byte[] {
0,
1
})] this ConfiguredValueTaskAwaitable<T> awaitable, bool async)
{
return awaitable;
}
public static ConfiguredValueTaskAwaitable EnsureCompleted(this ConfiguredValueTaskAwaitable awaitable, bool async)
{
return awaitable;
}
[Conditional("DEBUG")]
private static void VerifyTaskCompleted(bool isCompleted)
{
if (!isCompleted) {
if (Debugger.IsAttached)
Debugger.Break();
throw new InvalidOperationException("Task is not completed");
}
}
}
}