Using RegisterWaitForSingleObject if operation completes first

Question

I'm doing some asynchronous network I/O using Begin/End style methods. (It's actually a query against Azure Table Storage, but I don't think that matters.) I've implemented a client side timeout using the ThreadPool.RegisterWaitForSingleObject(). This is working fine as far as I can tell.

Because ThreadPool.RegisterWaitForSingleObject() takes a WaitHandle as an argument, I have to begin the I/O operation, then execute ThreadPool.RegisterWaitForSingleObject(). It seems like this introduces the possibility that the I/O completes before I even register the wait.

A simplified code sample:

private void RunQuery(QueryState queryState)
{
    //Start I/O operation
    IAsyncResult asyncResult = queryState.Query.BeginExecuteSegmented(NoopAsyncCallback, queryState);

    //What if the I/O operation completes here? 

    queryState.TimeoutWaitHandle = ThreadPool.RegisterWaitForSingleObject(asyncResult.AsyncWaitHandle, QuerySegmentCompleted, asyncResult, queryTimeout, true);
}

private void QuerySegmentCompleted(object opState, bool timedOut){
    IAsyncResult asyncResult = opState as IAsyncResult;
    QueryState state = asyncResult.AsyncState as QueryState;

    //If the I/O completed quickly, could TimeoutWaitHandle could be null here?
    //If so, what do I do about that?
    state.TimeoutWaitHandle.Unregister(asyncResult.AsyncWaitHandle);
}

What's the proper way to handle this? Do I still need to worry about Unregister()'ing the AsyncWaitHandle? If so, is there a fairly easy way to wait for it to be set?

Answer 1

Yep, you and everyone else has this problem. And it does not matter if the IO completed synchronously or not. There is still a race between the callback and the assignment. Microsoft should have provided the RegisteredWaitHandle to that callback function automatically. That would have solved everything. Oh well, hindsight is always 20-20 as they say.

What you need to do is keep reading the RegisteredWaitHandle variable until it is no longer null. It is okay to do this in a tight loop because the race is subtle enough that the loop will not be spinning around very many times.

private void RunQuery(QueryState queryState)
{
  // Start the operation.
  var asyncResult = queryState.Query.BeginExecuteSegmented(NoopAsyncCallback, queryState);

  // Register a callback.
  RegisteredWaitHandle shared = null;
  RegisteredWaitHandle produced = ThreadPool.RegisterWaitForSingleObject(asyncResult.AsyncWaitHandle,
    (state, timedout) =>
    {
      var asyncResult = opState as IAsyncResult;
      var state = asyncResult.AsyncState as QueryState;
      while (true)
      {
        // Keep reading until the value is no longer null.
        RegisteredWaitHandle consumed = Interlocked.CompareExchange(ref shared, null, null);
        if (consumed != null)
        {
          consumed.Unregister(asyncResult.AsyncWaitHandle);
          break;
        }
      }
    }, asyncResult, queryTimeout, true);

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

Answer 2

You do not need to Unregister if the I/O completed before the timeout as it was the completion that signalled your callback. In fact upon reading the docs of the Unregister method it seems totally unnecessary to call it as you are executing only once and you are not Unregistering in an unrelated method.

http://msdn.microsoft.com/en-us/library/system.threading.registeredwaithandle.unregister.aspx

If a callback method is in progress when Unregister executes, waitObject is not signaled until the callback method completes. In particular, if a callback method executes Unregister, waitObject is not signaled until that callback method completes.