H ManualResetEventAsync в черновиках

Возникла сегодня идея написать асинхронную версию ManualResetEvent, которая в задаче будет «ожидать» через await и при этом не занимать никакой поток.

В теории все просто, для стейтмашины нужен объект, который имеет метод GetAwaiter, который вернет awaiter, в котором реализован INotifyCompletion с OnCompleted, поле IsCompleted и метод GetResult.

Awaiter:

    public class ManualResetEventAwaiter : INotifyCompletion
    {
        Action _continuation;
        public void OnCompleted(Action continuation)
        {
            Volatile.Write(ref _continuation, continuation);
        }

        public bool GetResult()
        {
            return true;
        }

        public bool IsCompleted { get; private set; } = false;

        public ManualResetEventAwaiter GetAwaiter() => this;

        public void Set()
        {
            IsCompleted = true;

            Action continuation = Interlocked.Exchange(ref _continuation, null);

            if (continuation != null)
                Task.Run(continuation);
        }
    }

В OnCompleted стейтмашина передаст действие которое продолжит нашу логику выполнения. Для нас это воспринимается как логика после await.

В Set будет вызываться это продолжение работы. Я обернул его в Task для того что бы «ожидающие таски»(это условность) так же продолжились асинхронно. Хотя, возможно, continuation(там у нас IAsyncStateMachine.MoveNext) итак это сделает, но я до конца не уверен. Если кто в курсе напишите.

ManualResetEventAsync:

    public class ManualResetEventAsync
    {
        ConcurrentQueue<ManualResetEventAwaiter> awaiteQueue
            = new ConcurrentQueue<ManualResetEventAwaiter>();

        public bool State { get; private set; }

        public ManualResetEventAsync(bool initialState)
        {
            State = initialState;
        }

        public ManualResetEventAwaiter WaitOneAsync()
        {
            var awaitable = new ManualResetEventAwaiter();

            if (State)
                awaitable.Set();
            else
                awaiteQueue.Enqueue(awaitable);

            return awaitable;
        }

        public void Set()
        {
            State = true;

            while (State && awaiteQueue.TryDequeue(out var awaitable))
                awaitable.Set();
        }

        public void Reset()
        {
            State = false;
        }
    }

Метод WaitOneAsync будет использоваться с await.

Set вызовет все продолжения и установит сигнальное состояние.

Reset установит не сигнальное состояние.

И свойство State позволит узнать текущее состояние. Его можно использовать в условии до WaitOneAsync:

if (!mre.State)
   await mre.WaitOneAsync();

Но и без условия работа не нарушится.

UPD: Упростил немного код.

Ссылка на проект

var tsc = new TaskCompletionSource<bool>();
using (cancellationToken.Register(() => tsc.SetResult(false))
    await tsc.Task;

var tsc = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);

var source = new CancellationTokenSource();
var cancellationToken = source.Token;
source.CancelAfter(2 * 1000);
bool cleaned = false;

cancellationToken.Register(() =>
{
	//some work
	cleaned = true;
});

var tsc = new TaskCompletionSource<bool>();
using (cancellationToken.Register(() => tsc.SetResult(false)))
	await tsc.Task;

while (true)
{
	if (!cleaned)
	{
		Thread.Sleep(100); //some synchronous work emulation
	}
	else break;
}

Console.WriteLine("Finished"); //never reach this

cancellationToken.Register(() =>
{
    while (!cleaned)
    {
        Thread.Sleep(100); //some synchronous work emulation
    }
});

public class ManualResetEventAsync
{
    private TaskCompletionSource<bool> tcs = null;

    public Task WaitOneAsync() 
    {
        var tcs = Volatile.Read(ref this.tcs);
        if (tcs == null)
            return Task.CompletedTask;
        else
            return tcs.Task;
    }

    public void Set()
    {
        var old_tcs = Interlocked.Exchange(ref this.tcs, null);
        old_tcs?.SetResult(false);
    }

    public void Reset()
    {
        var new_tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
        Interlocked.CompareExchange(ref this.tcs, new_tcs, null);
    }
}