c# - How to to make UdpClient.ReceiveAsync() cancelable?

Question

I have an interface INetwork with a method:

Task<bool> SendAsync(string messageToSend, CancellationToken ct)

One implementation of the interface has code like this:

public async Task<bool> SendAsync(string messageToSend, CancellationToken ct)
{
  var udpClient = new UdpClient();
  var data = Encoding.UTF8.GetBytes (messageToSend);
  var sentBytes = await udpClient.SendAsync(data);
  return sentBytes == data.Length; 
}

Unfortunately, SendAsync() of the UdpClient class does not accept a CancellationToken.

So I started changing it to:

public Task<bool> SendAsync(string messageToSend, CancellationToken ct)
{
  var udpClient = new UdpClient();
  var data = Encoding.UTF8.GetBytes (messageToSend);
  var sendTask = udpClient.SendAsync(data);
  sendTask.Wait(ct);

  if(sendTask.Status == RanToCompletion)
  {
    return sendTask.Result == data.Length;
  }
}

Obviously this won't work because there is no Task being returned. However if I return the Task, the signatures don't match anymore. SendAsync() returns a Task<int>, but I need a Task<bool>.

And now I'm confused. :-) How to resolve this?

Answer 1

I know this is a little late, but I just recently had to make a UdpClient ReceiveAsync/SendAsync cancellable.

Your first code block is sending without a cancel (your title says receive by the way...).

Your second code block is defintely not the way to do it. You are calling *Async, and then Task.Wait, which blocks until the call is complete. This makes the call effectively synchronous and there's no point in calling the *Async version. The best solution is to use Async as follows:

...
var sendTask = udpClient.SendAsync(data);
var tcs = new TaskCompletionSource<bool>();
using( ct.Register( s => tcs.TrySetResult(true), null) )
{
    if( sendTask != await Task.WhenAny( task, tcs.Task) )
        // ct.Cancel() called
    else
        // sendTask completed first, so .Result will not block
}
...

There's no built-in way to cancel on UdpClient (none of the functions accept a CancellationToken), but you can take advantage of the ability to await multiple tasks with Task.WhenAny. This will return with the first task that completes (this is also an easy way to use Task.Delay() to implement timeouts). We then just need to create a Task that will complete when the CancellationToken is canceled, which we can do by creating a TaskCompletionSource and setting it with the CancellationToken's callback.

Once canceled, we can close the socket to actually "cancel" the underlying read/write.

The original idea for this came from another SO answer dealing with file handles, but it works with sockets too. I generally wrap it up in an extension method like so:

public static class AsyncExtensions
{
    public static async Task<T> WithCancellation<T>( this Task<T> task, CancellationToken cancellationToken )
    {
        var tcs = new TaskCompletionSource<bool>();
        using( cancellationToken.Register( s => ( (TaskCompletionSource<bool>)s ).TrySetResult( true ), tcs ) )
        {
            if( task != await Task.WhenAny( task, tcs.Task ) )
            {
                throw new OperationCanceledException( cancellationToken );
            }
        }

        return task.Result;
    }
}

Then use it like so:

try
{
    var data = await client.ReceiveAsync().WithCancellation(cts.Token);
    await client.SendAsync(data.Buffer, data.Buffer.Length, toep).WithCancellation(cts.Token);
}
catch(OperationCanceledException)
{
    client.Close();
}