1
\$\begingroup\$

Async SelectMany takes an enumeration (length unknown) of tasks, each of which returns another enumeration (length unknown and likely different from one another), and returns results from the second dimension as soon as they are ready. I used the name FoldAsync to make it descriptive, but I'm open to suggestions.

Imagine a task, identified with a number, runs for that many seconds then returns. The 2d array (first # being the time for the first dimension task) of tasks could be

1a  : 6     11  16  
17  : 33    1b
22  : 3     4   5

Different ways of nesting Task.WhenAll could result in the main thread iterating through them like so:

    Depth first             Breadth first
time    :   task id     time    :   task id
1       :   1a          1       :   1a
7       :   6           17      :   17
12      :   11          22      :   22
17      :   16          23      :   1b
17      :   17          25      :   3
18      :   1b          26      :   4
50      :   33          27      :   5
50      :   22          28      :   6
50      :   3           33      :   11
50      :   4           38      :   16
50      :   5           50      :   33

Async SelectMany accomplishes:

time    :   task id
1       :   1a
7       :   6
12      :   11
17      :   (16, 17)
17      :   (16, 17)
18      :   1b
22      :   22
25      :   3
26      :   4
27      :   5
50      :   33

Task.WhenAny can obtain the same results, but you'd have to maintain a list of both Task<IEnumerable<Task<T>>>, and Task<T> which you're now casting, and spend O(n^2) iterating through it. It's workable, but I like cool code instead.

The following is my implementation (with 800% too many comments). The /// documentation comments are messy because I couldn't figure out how to display nested generics inside. I tried to manage cancellation and Exceptions in a sane manner. I'm looking for comments, critiques, and hopefully improvements.

using System;
using System.Collections;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;

namespace WhereIWork.Utilities.Async
{
    /// <summary>
    /// A helper class containing the <see cref="FoldAsync{TResult}(IEnumerable{Task{IEnumerable{Task{TResult}}}})"/> extension method that is an async equivalent of <see cref="System.Linq.Enumerable"/>.SelectMany.
    /// Also provides an overload taking a <see cref="CancellationToken"/>
    /// </summary>
    /// <remarks>
    /// If the size of your dimensions are known, use <see cref="InterleavingExtensions.Interleaved{TResult}"/> instead.
    /// <see cref="FoldAsync{TResult}(IEnumerable{Task{IEnumerable{Task{TResult}}}})"/> blocks, because it cannot know the Count of its results until the all tasks in the first dimension are complete.
    /// </remarks>
    public static class AsyncEnumerable
    {
        /// <summary>
        /// An asynchronous version of <see cref="System.Linq.Enumerable"/>.SelectMany.
        /// Will return results in the order they are completed and not the order from the two-dimensional <see cref="IEnumerable"/>.
        /// </summary>
        /// <typeparam name="TResult">The type of the expected results.</typeparam>
        /// <param name="tasks">The two dimensional <see cref="IEnumerable"/> of <see cref="Task{TResult}"/> that will be folded into one dimension.</param>
        /// <returns>An <see cref="IEnumerable{Task}"/> which will block as its results are iterated.</returns>
        public static IEnumerable<Task<TResult>> FoldAsync<TResult>(this IEnumerable<Task<IEnumerable<Task<TResult>>>> tasks)
        {
            return new TaskFolder<TResult>(tasks);
        }

        /// <summary>
        /// An asynchronous version of <see cref="System.Linq.Enumerable"/>.SelectMany.
        /// Will return results in the order they are completed and not the order from the two-dimensional <see cref="IEnumerable"/>
        /// </summary>
        /// <typeparam name="TResult">The type of the expected results.</typeparam>
        /// <param name="tasks">The two dimensional <see cref="IEnumerable"/> of <see cref="Task{TResult}"/> that will be folded into one dimension.</param>
        /// <param name="token">The cancellation token to be passed into every <see cref="Task"/>&lt;<see cref="IEnumerable"/>&lt;<see cref="Task{TResult}"/>&gt;&gt; and <see cref="Task{TResult}"/>.</param>
        /// <returns>An <see cref="IEnumerable{Task}"/> which will block as its results are iterated.</returns>
        public static IEnumerable<Task<TResult>> FoldAsync<TResult>(this IEnumerable<Task<IEnumerable<Task<TResult>>>> tasks, CancellationToken token)
        {
            return new TaskFolder<TResult>(tasks, token);
        }
    }

    /// <summary>
    /// Implements <see cref="System.Linq.Enumerable"/>.SelectMany(<see cref="IEnumerable"/>&lt;<see cref="Task"/>&lt;<see cref="IEnumerable"/>&lt;<see cref="Task{TResult}"/>&gt;&gt;&gt;,
    /// Func&lt;<see cref="Task"/>&lt;<see cref="IEnumerable"/>&lt;<see cref="Task{TResult}"/>&gt;&gt;,
    /// <see cref="IEnumerable"/>&lt;<see cref="Task{TResult}"/>&gt;&gt;)
    /// <br/>The advantage of this class is that we process results in the order they are ready, without waiting for tasks that are listed earlier in the initial enumeration.
    /// </summary>
    /// <remark>
    /// The constructors are internal to allow for testing, but this class should not be extended.
    /// </remark>
    /// <typeparam name="TResult">The common type of the result of all <see cref="Task{TResult}"/> being folded.</typeparam>
    // ReSharper disable once InheritdocConsiderUsage
    public class TaskFolder<TResult> : IEnumerable<Task<TResult>>, IDisposable
    {
        /// <summary>
        /// The collection to which we will post results as they are ready.
        /// The <see cref="IEnumerator{TResult}"/> returned by this class comes from this collection.
        /// </summary>
        protected readonly BlockingCollection<Task<TResult>> Collection = new BlockingCollection<Task<TResult>>();

        protected readonly CancellationToken Token;

        /// <summary>
        /// The number of active tasks.
        /// </summary>
        protected int ActiveTasks;

        /// <summary>
        /// The current state of this <see cref="TaskFolder{TResult}"/>
        /// </summary>
        protected int CurrentState = InitialState;

        protected const int InitialState = 0;
        protected const int TasksCompleted = 1;
        protected const int EnumerableRequested = 2;

        /// <summary>
        /// Creates a new TaskFolder and initiates the process of listening for completions.
        /// </summary>
        /// <param name="twoDimensionalWork">The two dimensional <see cref="IEnumerable"/> of <see cref="Task{TResult}"/> that will be folded into one dimension.</param>
        // ReSharper disable once InheritdocConsiderUsage
        internal TaskFolder(IEnumerable<Task<IEnumerable<Task<TResult>>>> twoDimensionalWork) : this(twoDimensionalWork, CancellationToken.None)
        {
        }

        /// <summary>
        /// Creates a new TaskFolder with a cancellation token, and initiates the process of listening for completions.
        /// </summary>
        /// <param name="twoDimensionalWork">The two dimensional <see cref="IEnumerable"/> of <see cref="Task{TResult}"/> that will be folded into one dimension.</param>
        /// <param name="cancellationToken">The cancellation token to be passed into every <see cref="Task"/>&lt;<see cref="IEnumerable"/>&lt;<see cref="Task{TResult}"/>&gt;&gt; and <see cref="Task{TResult}"/>.</param>
        internal TaskFolder(IEnumerable<Task<IEnumerable<Task<TResult>>>> twoDimensionalWork, CancellationToken cancellationToken)
        {
            // This initial increment represents the need to add the continuation to all first dimension tasks before we complete the BlockingCollection.
            Interlocked.Increment(ref ActiveTasks);

            cancellationToken.Register(() => AdvanceTaskCompletion(true));
            Token = cancellationToken;
            foreach (var outerTask in twoDimensionalWork)
            {
                if (IsCanceled())
                {
                    Console.WriteLine("woah");
                    return;
                }

                // Increment first, then any decrementing in OuterContinuation will necessarily not be premature.
                Interlocked.Increment(ref ActiveTasks);

                // As this is an example of Dynamic Task Parallelism, we do not await the result of ContinueWith
                // Because we manage completion using _taskCount, we do not need to track the continuation to guarantee its completion.
                // Do not use an overload without the TaskScheduler parameter. http://blog.stephencleary.com/2015/01/a-tour-of-task-part-7-continuations.html
                outerTask.ContinueWith(OuterContinuation, CancellationToken.None, TaskContinuationOptions.DenyChildAttach, TaskScheduler.Default);
            }

            DecrementTaskCounter();
        }

        /// <summary>
        /// Called when a first dimension <see cref="Task"/>&lt;<see cref="IEnumerable"/>&lt;<see cref="Task{TResult}"/>&gt;&gt; is complete.
        /// If there was an exception, it will be returned as a <see cref="Task{TResult}"/> in the enumeration.
        /// </summary>
        /// <remarks>
        /// Not async since we have another method of tracking the completion of newly created tasks.
        /// </remarks>
        /// <param name="task">The first dimension <see cref="Task"/>&lt;<see cref="IEnumerable"/>&lt;<see cref="Task{TResult}"/>&gt;&gt; that has completed.</param>
        protected void OuterContinuation(Task<IEnumerable<Task<TResult>>> task)
        {
            if (IsCanceled())
            {
                return;
            }

            // In both the faulted and canceled states, there are no second dimension Task<TResult>, and so we ignore _taskCount.
            // If the task has faulted, add an exception to the results.
            if (task.IsFaulted)
            {
                Debug.Assert(task.Exception != null, "If a task if faulted, it should have an exception.");
                Debug.Assert(task.Exception.InnerException != null, "It is aggregated, so go 1 deeper");
                Collection.Add(Task.FromException<TResult>(task.Exception.InnerException), Token);
                DecrementTaskCounter();
                return;
            }

            // If the task was cancelled or threw an exception, add a cancellation to the results. It won't have child tasks to deal with.
            if (task.IsCanceled)
            {
                Collection.Add(Task.FromCanceled<TResult>(Token), Token);
                DecrementTaskCounter();
                return;
            }

            // By exclusion, the task has completed and we can safely use the Result without worrying about blocking.
            Debug.Assert(task.Status == TaskStatus.RanToCompletion, "The continuation should not be called when the Task is not complete");
            try
            {
                foreach (var innerTask in task.Result)
                {
                    if (IsCanceled())
                    {
                        return;
                    }

                    // Increment first, then any decrementing in InnerContinuation won't cause _taskCount to go to 0
                    Interlocked.Increment(ref ActiveTasks);

                    // As this is an example of Dynamic Task Parallelism, we do not await the result of ContinueWith
                    // Because we manage completion using _taskCount, we do not need to track the continuation to guarantee its completion.
                    // Do not use an overload without the TaskScheduler parameter. http://blog.stephencleary.com/2015/01/a-tour-of-task-part-7-continuations.html
                    innerTask.ContinueWith(InnerContinuation, CancellationToken.None, TaskContinuationOptions.DenyChildAttach, TaskScheduler.Default);
                }
            }
            catch (Exception e)
            {
                // Any exceptions should be reported to the consumer.
                try
                {
                    Collection.Add(Task.FromException<TResult>(e), Token);
                }
                catch
                {
                    // Ignore, we can't add this exception to the collection, so there's no way to report it to the consumer.
                }
            }
            finally
            {
                // Decrement at the end, ensuring that all second dimension Task<TResult> have contributed to the counter.
                DecrementTaskCounter();
            }
        }

        /// <summary>
        /// Called when a result is ready to be added to <see cref="Collection"/> and returned as part of this <see cref="IEnumerable"/>&lt;<see cref="Task{TResult}"/>&gt;
        /// </summary>
        /// <remarks>
        /// Not async since we do not use the result of the <see cref="Task{TResult}"/> and simply return it still wrapped.
        /// </remarks>
        /// <param name="task">The next result to return to the consumer.</param>
        protected void InnerContinuation(Task<TResult> task)
        {
            if (IsCanceled())
            {
                return;
            }

            try
            {
                Collection.Add(task, Token);
                DecrementTaskCounter();
            }
            catch (OperationCanceledException)
            {
                // Ignore, we can't add this exception to the collection, so there's no way to report it to the consumer.
            }
        }

        /// <summary>
        /// Checks if <see cref="Token"/> has <see cref="CancellationToken.IsCancellationRequested"/> set to true. 
        /// If so, will call <see cref="AdvanceTaskCompletion"/> to stop processing further results.
        /// </summary>
        /// <returns>
        /// <see cref="Token"/>.IsCancellationRequested
        /// </returns>
        protected bool IsCanceled()
        {
            if (!Token.IsCancellationRequested)
            {
                return false;
            }

            // Pretend all the tasks are done, no further results will be added.
            ActiveTasks = 0;
            AdvanceTaskCompletion(true);

            return true;
        }

        /// <summary>
        /// Calls <see cref="BlockingCollection{T}.CompleteAdding"/> to tell it we have finished processing <see cref="Task{TResult}"/>.
        /// Sets the <see cref="TasksCompleted"/> flag on <see cref="CurrentState"/> to indicate all <see cref="Task{TResult}"/> have been completed.
        /// </summary>
        protected void AdvanceTaskCompletion(bool fromCancellation = false)
        {
            // If this has already been run, ignore it.
            if ((CurrentState & TasksCompleted) != 0)
            {
                return;
            }

            // Completed or Canceled both mean we won't be adding anything new to the collection
            // If we're the first to advance to TasksComplete
            // ReSharper disable once InvertIf
            if (Interlocked.CompareExchange(ref CurrentState, TasksCompleted, InitialState) == InitialState ||
                Interlocked.CompareExchange(ref CurrentState, TasksCompleted | EnumerableRequested, EnumerableRequested) == EnumerableRequested)
            {
                if (fromCancellation)
                {
                    Collection.Add(Task.FromCanceled<TResult>(Token), CancellationToken.None);
                }

                Collection.CompleteAdding();
            }
        }

        /// <summary>
        /// Decrements <see cref="ActiveTasks"/>, and if all the tasks have completed, will advance the internal state.
        /// </summary>
        protected void DecrementTaskCounter()
        {
            if (Interlocked.Decrement(ref ActiveTasks) > 0)
            {
                return;
            }

            AdvanceTaskCompletion();
        }

        /// <summary>
        /// Returns an <see cref="IEnumerator"/> which will provide <see cref="Task{TResult}"/> in the order they have completed, not the order they were supplied.
        /// </summary>
        /// <returns>A blocking <see cref="IEnumerator"/>&lt;<see cref="Task{TResult}"/>&gt;.</returns>
        // ReSharper disable once InheritdocConsiderUsage
        public IEnumerator<Task<TResult>> GetEnumerator()
        {
            IEnumerable<Task<TResult>> YieldCollectionContents()
            {
                using (Collection)
                {
                    foreach (var task in Collection.GetConsumingEnumerable())
                    {
                        yield return task;
                    }
                }
            }

            // Allow only one thread to advance into EnumerableRequested
            // ReSharper disable once InvertIf
            if (Interlocked.CompareExchange(ref CurrentState, EnumerableRequested, InitialState) == InitialState ||
                Interlocked.CompareExchange(ref CurrentState, TasksCompleted | EnumerableRequested, TasksCompleted) == TasksCompleted)
            {
                // GetEnumerator's created IEnumerator will notice the using(collection) and call BlockingCollection.Dispose in it's IEnumerator.Dispose.
                // foreach is aware of IEnumerator being IDisposable and will call it.
                // https://blogs.msdn.microsoft.com/dancre/2008/03/15/yield-and-usings-your-dispose-may-not-be-called/
                return YieldCollectionContents().GetEnumerator();
            }

            throw new InvalidOperationException("The enumerator has already been fetched. Cannot Enumerate this object twice.");
        }

        /// <inheritdoc />
        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        /// <summary>
        /// Dispose of the underlying <see cref="BlockingCollection{T}"/>
        /// If <see cref="GetEnumerator"/> has been called, this will do nothing. We rely on the consumer instead.
        /// </summary>
        // ReSharper disable once InheritdocConsiderUsage
        public void Dispose()
        {
            if ((Interlocked.Exchange(ref CurrentState, EnumerableRequested) & EnumerableRequested) == 0)
            {
                Collection?.Dispose();
            }
        }
    }
}

Here is a test case that works

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Threading.Tasks;
using WhereIWork.Utilities.Async;

namespace ConcurrencyTests
{

    public class Program
    {
        private static IEnumerable<Task<IEnumerable<Task<int>>>> TestScenario()
        {
            return Enumerable.Range(0, 5).Select(async y =>
            {
                await Task.Delay(y * 100);
                return Enumerable.Range(1,5).Select(async z =>
                {
                    await Task.Delay(z * 300);
                    var value = y*5 + z;
                    if (value % 7 == 3)
                    {
                        throw new Exception("What! " + value);
                    }

                    return value;
                });
            });
        }



        public static async Task Main()
        {
            var watch = new Stopwatch();
            watch.Start();
            foreach (var value in TestScenario().FoldAsync())
            {
                try
                {
                    var result = await value;
                    Console.WriteLine(watch.ElapsedMilliseconds + " | " + result);
                }
                catch (Exception e)
                {
                    Console.WriteLine(e);
                }
            }

            Console.WriteLine(watch.ElapsedMilliseconds);
            Console.ReadLine();
        }
    }
}
\$\endgroup\$
  • \$\begingroup\$ The words Cancellation, cancelled, and cancelling have lost all meaning and are probably spelled a variety of ways with one or two l's. <see cref"en.wikipedia.org/wiki/Semantic_satiation" /> \$\endgroup\$ – Jean-Bernard Pellerin Feb 15 at 3:29
  • 2
    \$\begingroup\$ Are you sure this works? I get The collection has been disposed. exception when trying to use it. Could you provide some test code? \$\endgroup\$ – t3chb0t Feb 15 at 10:33
  • \$\begingroup\$ I would change your flags enum to not hardcode values, for example, TaskFetched = 5 should be TaskFetched = Task | Fetched \$\endgroup\$ – Brad M Feb 15 at 15:42
  • \$\begingroup\$ @t3chb0t Can you please share your test case? \$\endgroup\$ – Jean-Bernard Pellerin Feb 15 at 16:10
  • \$\begingroup\$ This is how I tested it in LINQPad - I renamed your method :P. \$\endgroup\$ – t3chb0t Feb 15 at 16:18

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.