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