Catching exceptions immediately does not work with Task.WhenAll

Question

I have two instances of a class which creates a UDP socket to receive data from UDP clients. If one of the instances throws an exception I want to handle it immediately in a higher layer. In my program they're started with await Task.WhenAll(recv1.StartAsync(), recv2.StartAsync). This however waits for all tasks to finish before the first exception is thrown. Any ideas on how to resolve this problem?

static async Task Main(string[] args)
{
  var udpReceiver1 = new UdpReceiver(localEndpoint1);
  var udpReceiver2 = new UdpReceiver(localEndpoint2);

  var cts = new CancellationTokenSource();

  try
  {
    await Task.WhenAll(udpReceiver1.StartAsync(cts.Token), udpReceiver2.StartAsync(cts.Token));
  }
  catch (Exception e)
  {
    // Handle Exception...

    cts.Cancel();
  }
}

class UdpReceiver
{
  public UdpReceiver(IPEndPoint endpoint)
  {
    udpClient = new UdpClient(endpoint);
  }

  public async Task StartAsync(CancellationToken cancellationToken)
  {
    try
    {
      while (!cancellationToken.IsCancellationRequested)
      {
        var result = await ReceiveAsync(cancellationToken);
        var message = Encoding.UTF8.GetString(result.Buffer);
        Trace.WriteLine($"UdpClient1 received message:{Encoding.UTF8.GetString(result.Buffer)}");
      
        // throw new Exception("UdpClient1 raising exception");
      }
    }
  }

  private async Task<UdpReceiveResult> ReceiveAsync(CancellationToken cancellationToken)
  {
    var tcs = new TaskCompletionSource<UdpReceiveResult>();
    using (cancellationToken.Register(() => tcs.TrySetCanceled(), false))
    {
      var task = udpClient.ReceiveAsync();

      var completedTask = await Task.WhenAny(task, tcs.Task);

      var result = await completedTask.ConfigureAwait(false);

      return result;
    }
  }

  private UdpClient udpClient;
}

Update 1: Task.WhenAny would be a viable solution. Thanks @CamiloTerevinto

try
{
  await await Task.WhenAny(udpReceiver1.StartAsync(cts.Token), udpReceiver2.StartAsync(cts.Token));
}
catch (Exception e)
{
  // Handle Exception...

  cts.Cancel();
}

Update 2: For a more fine grained exception handling of all tasks I'd go with my own adapted implementation of Task.WhenAll proposed by @Servy.

Answer 1

The behavior is different enough from the framework WhenAll implementation that you're probably best just writing your own adapted version, fortunately it's not particularly hard to implement. Just attach a continuation to every single task, if any is cancelled or faulted, the resulting task does the same, if it succeeds store the result, and if the last task was the one to succeed, complete the task with all of the stored results.

public static Task<IEnumerable<TResult>> WhenAll<TResult>(IEnumerable<Task<TResult>> tasks)
{
    var listOfTasks = tasks.ToList();
    if (listOfTasks.Count == 0)
    {
        return Task.FromResult(Enumerable.Empty<TResult>());
    }
    var tcs = new TaskCompletionSource<IEnumerable<TResult>>();
    var results = new TResult[listOfTasks.Count];
    int completedTasks = 0;
    for (int i = 0; i < listOfTasks.Count; i++)
    {
        int taskIndex = i;
        Task<TResult> task = listOfTasks[i];
        task.ContinueWith(_ =>
        {
            if (task.IsCanceled)
                tcs.TrySetCanceled();
            else if (task.IsFaulted)
                tcs.TrySetException(task.Exception.InnerExceptions);
            else
            {
                results[taskIndex] = task.Result;
                if (Interlocked.Increment(ref completedTasks) == listOfTasks.Count)
                {
                    tcs.TrySetResult(results);
                }
            }
        });
    }
    return tcs.Task;
}

As with so many task based generic operations, you also need a version without the results, and if you don't want to deal with a notable overhead, you really need to just copy-paste the result-based approach but with all of the results ripped out, which isn't hard, just inelegant. Turning all of these tasks into tasks with a result would also work, but for an operation like this the overhead is likely problematic.

public static Task WhenAll(IEnumerable<Task> tasks)
{
    var listOfTasks = tasks.ToList();
    if (listOfTasks.Count == 0)
    {
        return Task.CompletedTask;
    }
    var tcs = new TaskCompletionSource<bool>();
    int completedTasks = 0;
    for (int i = 0; i < listOfTasks.Count; i++)
    {
        int taskIndex = i;
        Task task = listOfTasks[i];
        task.ContinueWith(_ =>
        {
            if (task.IsCanceled)
                tcs.TrySetCanceled();
            else if (task.IsFaulted)
                tcs.TrySetException(task.Exception.InnerExceptions);
            else
            {
                if (Interlocked.Increment(ref completedTasks) == listOfTasks.Count)
                {
                    tcs.TrySetResult(true);
                }
            }
        });
    }
    return tcs.Task;
}

Answer 2

There's probably a way to do it but I can't think of one without making your code very messy. It'd be better to handle the exception in the actual task. If you need to handle it with common code, use a handler delegate.

static async Task Main(string[] args)
{
    var cts = new CancellationTokenSource();

    //This is our common error handler
    void HandleException(Exception ex)
    {
        Log("Exception!" + ex.Message);
        cts.Cancel();
    }

    var udpReceiver1 = new UdpReceiver(localEndpoint1);
    var udpReceiver2 = new UdpReceiver(localEndpoint1);

    //We pass the handler as one of the arguments
    await Task.WhenAll(udpReceiver1.StartAsync(cts.Token, HandleException), udpReceiver2.StartAsync(cts.Token, HandleException));
}

class UdpReceiver
{
  public async Task StartAsync(CancellationToken cancellationToken, Action<Exception> errorHandler)
  {
      try
      {
          while (!cancellationToken.IsCancellationRequested)
          {
              //Main logic goes here
          }
      }
      catch(Exception ex)
      {
          errorHandler(ex);  //Call common error handling code
      }
  }

Answer 3

You could await the tasks in two steps. In the first step await any of them to complete, and in case of failure initiate the cancellation. Don't handle the exception in this step. Delay the exception handling until the second step, after awaiting all of the tasks to complete. Both tasks may have failed, so you may want to handle the exception of each one separately.

Task task1 = udpReceiver1.StartAsync(cts.Token);
Task task2 = udpReceiver2.StartAsync(cts.Token);

// Await any task to complete
Task firstCompletedTask = await Task.WhenAny(task1, task2);
if (firstCompletedTask.IsFaulted) cts.Cancel();

try
{
    // Await them all to complete
    await Task.WhenAll(task1, task2);
}
catch
{
    if (task1.IsFaulted) HandleException(task1.Exception.InnerException);
    if (task2.IsFaulted) HandleException(task2.Exception.InnerException);
}