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

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

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

Awaiter:

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

        public bool IsCompleted => _isCompleted;

        public bool GetResult() => true;

        public ManualResetEventAwaiter GetAwaiter() => this;

        public void Continue()
        {
            Volatile.Write(ref _isCompleted, true);

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

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

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

В Continue будет вызываться это продолжение работы. Я обернул его в Task для того что бы «ожидающие таски» так же продолжились асинхронно.

ManualResetEventAsync:

    public class ManualResetEventAsync
    {
        ConcurrentQueue<ManualResetEventAwaiter> awaitQueue = null;

        public ManualResetEventAsync(bool initialState)
        {
            if (!initialState)
                Reset();
        }

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

            var queue = Volatile.Read(ref awaitQueue);

            if (queue == null)
                awaitable.Continue();
            else
            {
                queue.Enqueue(awaitable);

                //check queue
                var upd_queue = Volatile.Read(ref awaitQueue);               

                if (!ReferenceEquals(queue, upd_queue))
                    awaitable.Continue();
            }

            return awaitable;
        }

        public void Set()
        {
            var queue = Interlocked.Exchange(ref awaitQueue, null);

            if (queue != null)
                while (queue.TryDequeue(out var awaitable))
                    awaitable.Continue();
        }

        public void Reset()
        {
            Interlocked.CompareExchange(
                ref awaitQueue,
                new ConcurrentQueue<ManualResetEventAwaiter>(),
                null);
        }
    }

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

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

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

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);
    }
}