CancellationTokenRegistration.Dispose in Async Task

Question

I am making a SocketExtender class which will provide async/await Task extension methods to the Socket class. In my extension methods I am adding is the ability to cancel a Socket operation, such as ConnectAsync, ReceiveAsync, or SendAsync via a CancellationToken parameter.

Doing some reading on the optimal way to do this, I have decided to wrap the calls in TaskCompletionSource<T> and pass back a Task for the operation.

For example, I want to wrap the BeginReceive and EndReceive APM methods into the following Task-based method:

public static Task<int> ReceiveAsync(this Socket socket, byte[] buffer, int offset, int count, SocketFlags flags, CancellationToken token) 
{
  var tcs = new TaskCompletionSource<int>();

  socket.BeginReceive(buffer, offset, count, flags, result => 
  {
    try
    {
      var bytesReceived = socket.EndReceive(result);
      tcs.SetResult(bytesReceived);
    }
    catch(Exception ex)
    {
      if(token.IsCancellationRequested)
        tcs.SetCanceled();  // <- Yes, this is a typo in the .NET framework! :)
      else
        tcs.SetException(ex);
    }

  });

  return tcs.Task;
}

NOTE: The above method does not actually implement cancellation! Yes, it will handle it, but it will not cause the Socket to actually cancel the operation. So I need the CancellationToken to trigger a Socket.Close() call, via token.Register(). No biggie, but the question becomes how do I properly dispose of the CancellationTokenRegistration object returned?

Here is my first thought:

public static Task<int> ReceiveAsync(this Socket socket, byte[] buffer, int offset, int count, SocketFlags flags, CancellationToken token) 
{
  var tcs = new TaskCompletionSource<int>();

  using(token.Register(socket.Close))  // <- Here we are ensuring disposal of the CTR
  {
    socket.BeginReceive(buffer, offset, count, flags, result => 
    {
      try
      {
        var bytesReceived = socket.EndReceive(result);
        tcs.SetResult(bytesReceived);
      }
      catch(Exception ex)
      {
        if(token.IsCancellationRequested)
          tcs.SetCanceled();  // <- Yes, this is a typo in the .NET framework! :)
        else
          tcs.SetException(ex);
      }
    });
  }

  return tcs.Task;
}

However, my question is that will the CancellationTokenRegistration be disposed of via using before the asynchronous call completes? My first instinct wants to say yes, because the socket.BeginReceive call will not block and return immediately, causing the using statement to clean up once the method returns tcs.Task.

But then again, maybe the C# compiler 'understands' what I am doing and will see there's an anonymous method executing in a child scope and some black magic will take place to make it work the way I want.

Do I even need to care about disposing of the CancellationTokenRegistration here? Should I keep a reference to it and call Dispose() in the finally block instead? Or will this work as I want it to as is?

Answer 1

tcs.SetCanceled(); // <- Yes, this is a typo in the .NET framework! :)

Actually, it's an oddity in the U.S. version of the English language. British (and other) Englishes would use Cancelled, but U.S. English uses Canceled. Note that Cancellation always uses two l's in every version of the English language.

So I need the CancellationToken to trigger a Socket.Close() call, via token.Register()

I disagree. The convention for cancellation is that only that operation is cancelled. So if you had a Send operation that took a cancellation token, I would be surprised to find the semantics of send cancellation also faults any Receive operations and further makes the entire socket unusable.

will the CancellationTokenRegistration be disposed of via using before the asynchronous call completes?

Yes. There's no magic understanding on the compiler's part here.

Do I even need to care about disposing of the CancellationTokenRegistration here? Should I keep a reference to it and call Dispose() in the finally block instead? Or will this work as I want it to as is?

I would recommend not using a CancellationToken in the individual operation calls at all. It is logically an object-level cancellation, so creating a separate class and passing it into the constructor is a more proper design IMO. But if you insist, the finally block sounds like a good approach.