C# 5 Awaitable WaitHandle

WaitHandlerAwaitable.cs

public static class AwaitExtensions
{
    /// <summary>
    /// Provides await functionality for ordinary <see cref="WaitHandle"/>s.
    /// </summary>
    /// <param name="handle">The handle to wait on.</param>
    /// <returns>The awaiter.</returns>
    public static TaskAwaiter GetAwaiter(this WaitHandle handle)
    {
        Contract.Requires<ArgumentNullException>(handle != null);
        return handle.ToTask().GetAwaiter();
    }

    /// <summary>
    /// Creates a TPL Task that is marked as completed when a <see cref="WaitHandle"/> is signaled.
    /// </summary>
    /// <param name="handle">The handle whose signal triggers the task to be completed.</param>
    /// <returns>A Task that is completed after the handle is signaled.</returns>
    /// <remarks>
    /// There is a (brief) time delay between when the handle is signaled and when the task is marked as completed.
    /// </remarks>
    public static Task ToTask(this WaitHandle handle)
    {
        Contract.Requires<ArgumentNullException>(handle != null);
        Contract.Ensures(Contract.Result<Task>() != null);

        var tcs = new TaskCompletionSource<object>();
        var localVariableInitLock = new object();
        lock (localVariableInitLock)
        {
            RegisteredWaitHandle callbackHandle = null;
            callbackHandle = ThreadPool.RegisterWaitForSingleObject(
                handle,
                (state, timedOut) =>
                {
                    tcs.SetResult(null);

                    // We take a lock here to make sure the outer method has completed setting the local variable callbackHandle.
                    lock (localVariableInitLock)
                    {
                        callbackHandle.Unregister(null);
                    }
                },
                state: null,
                millisecondsTimeOutInterval: Timeout.Infinite,
                executeOnlyOnce: true);
        }

        return tcs.Task;
    }
}

Thanks for sharing this. There is a potential problem which could result in a NullReferenceException if RegisterWaitForSingleObject could call the callback directly in the same thread (it may be just a theoretical problem). Here is a slightly modified example using Interlocked.CompareExchange and a while loop (taken from http://stackoverflow.com/questions/10741669/c-using-registerwaitforsingleobject-if-operation-completes-first):

    public static Task ToTask(this WaitHandle handle)
    {
        if (handle == null) throw new ArgumentNullException("handle");

        var tcs = new TaskCompletionSource<object>();
        RegisteredWaitHandle shared = null;
        RegisteredWaitHandle produced = ThreadPool.RegisterWaitForSingleObject(
            handle,
            (state, timedOut) =>
            {
                tcs.SetResult(null);

                while (true)
                {
                    RegisteredWaitHandle consumed = Interlocked.CompareExchange(ref shared, null, null);
                    if (consumed != null)
                    {
                        consumed.Unregister(null);
                        break;
                    }
                }
            },
            state: null,
            millisecondsTimeOutInterval: Timeout.Infinite,
            executeOnlyOnce: true);

        // Publish the RegisteredWaitHandle so that the callback can see it.
        Interlocked.CompareExchange(ref shared, produced, null);

        return tcs.Task;
    }

Official documentation – From Wait Handles to TAP – does not implement any same thread execution avoidance. It is likely it is not needed. Or am I wrong, @utopius?

RegisterWaitForSingleObject will never execute the callback inline on the same thread. Per the docs:

The RegisterWaitForSingleObject method queues the specified delegate to the thread pool.

What may happen however is that there is a race. Consider that the WaitHandle is signaled and the callback is scheduled to the threadpool. That callback may happen before RegisterWaitForSingleObject returns and the assignment to callbackHandle is made. But I mitigate that threat using the lock. There's no need for a while loop.

However doc does not tell if it is guaranteed that such thread pool thread will be different from that which scheduled delegate for execution. On the other hand I cannot find functional difference between that spinning and simple locking. So how can spinning prevent anything more?

It doesn't matter whether the threadpool thread is the same or different. What would matter is whether the callback could be executed directly by the RegisterWaitForSingleObject method before it returns (which would be on the same thread). And if so, spinning would deadlock because the assignment cannot be completed while my lock would throw NRE. So both would be broken.

But since the documentation is that it's queued to the threadpool, my lock is perfectly adequate. The spin approach would work too, perhaps, but is more complicated and may spin the CPU more.