ForEachAsync extension method (a way to run an async operation on each item of a sequence in parallel)

Question

In a recent project I worked on we faced some issues due to an excess of parallelization (thousands of threads were created and the overall result was a degradation of performance and several spikes in CPU usage).

What we needed to solve this problem was a way to apply an asynchronous operation to each item of a sequence in a parallel fashion, with the possibility of specifying a maximum degree of parallelism.

By looking at this Stack Overflow question I jumped into Stephan Toub's ForEachAsync. Starting from there I implemented the following extension methods:

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

namespace Lib.Concurrency.Extensions
{
    /// <summary>
    /// Extension methods for enumerables
    /// </summary>
    public static class EnumerableExtensions
    {
        /// <summary>
        /// Executes an asynchronous operation for each item inside a source sequence. These operations are run concurrently in a parallel fashion. The invokation returns a task which completes when all of the asynchronous operations (one for each item inside the source sequence) complete. It is possible to constrain the maximum number of parallel operations.
        /// </summary>
        /// <typeparam name="T">The type of the items inside <paramref name="source"/></typeparam>
        /// <param name="source">The source sequence</param>
        /// <param name="maxDegreeOfParallelism">The maximum number of operations that are able to run in parallel</param>
        /// <param name="operation">The asynchronous operation to be executed for each item inside <paramref name="source"/></param>
        /// <returns>A task which completes when all of the asynchronous operations (one for each item inside <paramref name="source"/>) complete</returns>
        /// <exception cref="ArgumentNullException"><paramref name="source"/> is <c>null</c>.</exception>
        /// <exception cref="ArgumentNullException"><paramref name="operation"/> is <c>null</c>.</exception>
        /// <exception cref="ArgumentOutOfRangeException"><paramrefname="maxDegreeOfParallelism"/> is less than or equal to zero</exception>
        public static Task ForEachAsync<T>(
            this IEnumerable<T> source,
            int maxDegreeOfParallelism,
            Func<T, Task> operation)
        {
            if (source == null)
                throw new ArgumentNullException(nameof(source));

            if (operation == null)
                throw new ArgumentNullException(nameof(operation));

            EnsureValidMaxDegreeOfParallelism(maxDegreeOfParallelism);

            var tasks = from partition in Partitioner.Create(source).GetPartitions(maxDegreeOfParallelism)
                                    select Task.Run(async () =>
                                    {
                                        using (partition)
                                        {
                                            while (partition.MoveNext())
                                            {
                                                await operation(partition.Current).ConfigureAwait(false);
                                            }
                                        }
                                    });
            return Task.WhenAll(tasks);
        }

        /// <summary>
        /// Executes an asynchronous operation for each item inside a source sequence. These operations are run concurrently in a parallel fashion. The invokation returns a task whose result is a sequence containing the results of all the asynchronous operations (in source sequence order). It is possible to constrain the maximum number of parallel operations.
        /// </summary>
        /// <typeparam name="TSource">The type of the items inside the source sequence</typeparam>
        /// <typeparam name="TResult">The type of the object produced by invoking <paramref name="operation"/> on any item of <paramref name="source"/></typeparam>
        /// <param name="source">The source sequence</param>
        /// <param name="maxDegreeOfParallelism">The maximum number of operations that are able to run in parallel</param>
        /// <param name="operation">The asynchronous operation to be executed for each item inside <paramref name="source"/>. This operation will produce a result of type <typeparamref name="TResult"/></param>
        /// <returns>A task which completes when all of the asynchronous operations (one for each item inside <paramref name="source"/>) complete. This task will produce a sequence of objects of type <typeparamref name="TResult"/> which are the results (in source sequence order) of applying <paramref name="operation"/> to all items in <paramref name="source"/></returns> 
        /// <exception cref="ArgumentNullException"><paramref name="source"/> is <c>null</c>.</exception>
        /// <exception cref="ArgumentNullException"><paramref name="operation"/> is <c>null</c>.</exception>
        /// <exception cref="ArgumentOutOfRangeException"><paramref name="maxDegreeOfParallelism"/> is less than or equal to zero.</exception>
        public static async Task<IEnumerable<TResult>> ForEachAsync<TSource, TResult>(
            this IEnumerable<TSource> source,
            int maxDegreeOfParallelism,
            Func<TSource, Task<TResult>> operation)
        {
            if (source == null)
                throw new ArgumentNullException(nameof(source));

            if (operation == null)
                throw new ArgumentNullException(nameof(operation));

            EnsureValidMaxDegreeOfParallelism(maxDegreeOfParallelism);

            var resultsByPositionInSource = new ConcurrentDictionary<long, TResult>();

            var tasks = from partition in Partitioner.Create(source).GetOrderablePartitions(maxDegreeOfParallelism)
                                    select Task.Run(async () =>
                                    {
                                        using (partition)
                                        {
                                            while (partition.MoveNext())
                                            {
                                                var positionInSource = partition.Current.Key;
                                                var item = partition.Current.Value;

                                                var result = await operation(item).ConfigureAwait(false);

                                                resultsByPositionInSource.TryAdd(positionInSource, result);
                                            }
                                        }
                                    });
            await Task.WhenAll(tasks).ConfigureAwait(false);

            return Enumerable.Range(0, resultsByPositionInSource.Count)
                .Select(position => resultsByPositionInSource[position]);
        }

        private static void EnsureValidMaxDegreeOfParallelism(int maxDegreeOfParallelism)
        {
            if (maxDegreeOfParallelism <= 0)
            {
                throw new ArgumentOutOfRangeException(
                    nameof(maxDegreeOfParallelism),
                    $"Invalid value for the maximum degree of parallelism: {maxDegreeOfParallelism}. The maximum degree of parallelism must be a positive integer.");
            }
        }
    }
}

Can you spot any error or issue with this code? Any suggestion to improve this implementation is welcome (I have already planned new overloads to offer support for cancellation).

Update ( 10th September 2018 )

After some testing with the version of the code showed above we decided to opt for a different implementation based on SemaphoreSlim class.

The issue we found with the previously posted version is due to the fact that fixed a maximum degree of parallelism of n, then exactly n partitions will be created and then exactly n tasks will be created.

The desired behavior is different: if the maximum degree of parallelism is set to n, then the number of parallel tasks should be less then or equal to n. For instance given a sequence of m items, with m < n, then we expect m parallel operations. This was not possible with the implementation showed above.

Here is the final version of the code (support for cancellation is still missing):

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

namespace Deltatre.Utils.Concurrency.Extensions
{
    /// <summary>
    /// Extension methods for enumerables
    /// </summary>
    public static class EnumerableExtensions
    {
        /// <summary>
        /// Executes an asynchronous operation for each item inside a source sequence. These operations are run concurrently in a parallel fashion. The invokation returns a task which completes when all of the asynchronous operations (one for each item inside the source sequence) complete. It is possible to constrain the maximum number of parallel operations.
        /// </summary>
        /// <typeparam name="T">The type of the items inside <paramref name="source"/></typeparam>
        /// <param name="source">The source sequence</param>
        /// <param name="operation">The asynchronous operation to be executed for each item inside <paramref name="source"/></param>
        /// <param name="maxDegreeOfParallelism">The maximum number of operations that are able to run in parallel. If null, no limits will be set for the maximum number of parallel operations (same behaviour as Task.WhenAll)</param>
        /// <returns>A task which completes when all of the asynchronous operations (one for each item inside <paramref name="source"/>) complete</returns>
        /// <exception cref="ArgumentNullException"><paramref name="source"/> is <c>null</c>.</exception>
        /// <exception cref="ArgumentNullException"><paramref name="operation"/> is <c>null</c>.</exception>
        /// <exception cref="ArgumentOutOfRangeException"><paramref name="maxDegreeOfParallelism"/> is less than or equal to zero.</exception>
        public static Task ForEachAsync<T>(
            this IEnumerable<T> source,
            Func<T, Task> operation,
            int? maxDegreeOfParallelism = null)
        {
            if (source == null)
                throw new ArgumentNullException(nameof(source));

            if (operation == null)
                throw new ArgumentNullException(nameof(operation));

            EnsureValidMaxDegreeOfParallelism(maxDegreeOfParallelism);

            return (maxDegreeOfParallelism == null)
                ? ApplyOperationToAllItems(source, operation)
                : ApplyOperationToAllItemsWithConstrainedParallelism(source, operation, maxDegreeOfParallelism.Value);
        }

        private static Task ApplyOperationToAllItems<T>(
            IEnumerable<T> items,
            Func<T, Task> operation)
        {
            var tasks = items.Select(operation);
            return Task.WhenAll(tasks);
        }

        private static async Task ApplyOperationToAllItemsWithConstrainedParallelism<T>(
            IEnumerable<T> items,
            Func<T, Task> operation, 
            int maxDegreeOfParallelism)
        {
            using (var throttler = new SemaphoreSlim(maxDegreeOfParallelism))
            {
                var tasks = new List<Task>();

                foreach (var item in items)
                {
                    await throttler.WaitAsync().ConfigureAwait(false);

#pragma warning disable IDE0039 // Use local function
                    Func<Task> bodyOfNewTask = async () =>
#pragma warning restore IDE0039 // Use local function
                    {
                        try
                        {
                            await operation(item).ConfigureAwait(false);
                        }
                        finally
                        {
                            throttler.Release();
                        }
                    };

                    tasks.Add(Task.Run(bodyOfNewTask));
                }

                await Task.WhenAll(tasks).ConfigureAwait(false);
            }
        }

        /// <summary>
        /// Executes an asynchronous operation for each item inside a source sequence. These operations are run concurrently in a parallel fashion. The invokation returns a task whose result is a sequence containing the results of all the asynchronous operations (in source sequence order). It is possible to constrain the maximum number of parallel operations.
        /// </summary>
        /// <typeparam name="TSource">The type of the items inside the source sequence</typeparam>
        /// <typeparam name="TResult">The type of the object produced by invoking <paramref name="operation"/> on any item of <paramref name="source"/></typeparam>
        /// <param name="source">The source sequence</param>
        /// <param name="operation">The asynchronous operation to be executed for each item inside <paramref name="source"/>. This operation will produce a result of type <typeparamref name="TResult"/></param>
        /// <param name="maxDegreeOfParallelism">The maximum number of operations that are able to run in parallel. If null, no limits will be set for the maximum number of parallel operations (same behaviour as Task.WhenAll)</param>
        /// <returns>A task which completes when all of the asynchronous operations (one for each item inside <paramref name="source"/>) complete. This task will produce a sequence of objects of type <typeparamref name="TResult"/> which are the results (in source sequence order) of applying <paramref name="operation"/> to all items in <paramref name="source"/></returns> 
        /// <exception cref="ArgumentNullException"><paramref name="source"/> is <c>null</c>.</exception>
        /// <exception cref="ArgumentNullException"><paramref name="operation"/> is <c>null</c>.</exception>
        /// <exception cref="ArgumentOutOfRangeException"><paramref name="maxDegreeOfParallelism"/> is less than or equal to zero.</exception>
        public static Task<TResult[]> ForEachAsync<TSource, TResult>(
            this IEnumerable<TSource> source,
            Func<TSource, Task<TResult>> operation,
            int? maxDegreeOfParallelism = null)
        {
            if (source == null)
                throw new ArgumentNullException(nameof(source));

            if (operation == null)
                throw new ArgumentNullException(nameof(operation));

            EnsureValidMaxDegreeOfParallelism(maxDegreeOfParallelism);

            return (maxDegreeOfParallelism == null)
                ? ApplyOperationToAllItems(source, operation)
                : ApplyOperationToAllItemsWithConstrainedParallelism(source, operation, maxDegreeOfParallelism.Value);
        }

        private static Task<TResult[]> ApplyOperationToAllItems<TItem, TResult>(
            IEnumerable<TItem> items,
            Func<TItem, Task<TResult>> operation)
        {
            var tasks = items.Select(operation);
            return Task.WhenAll(tasks);
        }

        private static async Task<TResult[]> ApplyOperationToAllItemsWithConstrainedParallelism<TItem, TResult>(
            IEnumerable<TItem> items,
            Func<TItem, Task<TResult>> operation,
            int maxDegreeOfParallelism)
        {
            var resultsByPositionInSource = new ConcurrentDictionary<long, TResult>();

            using (var throttler = new SemaphoreSlim(maxDegreeOfParallelism))
            {
                var tasks = new List<Task>();

                foreach (var itemWithIndex in items.Select((item, index) => new { item, index }))
                {
                    await throttler.WaitAsync().ConfigureAwait(false);

#pragma warning disable IDE0039 // Use local function
                    Func<Task> bodyOfNewTask = async () =>
#pragma warning restore IDE0039 // Use local function
                    {
                        try
                        {
                            var item = itemWithIndex.item;
                            var positionInSource = itemWithIndex.index;

                            var result = await operation(item).ConfigureAwait(false);

                            resultsByPositionInSource.TryAdd(positionInSource, result);
                        }
                        finally
                        {
                            throttler.Release();
                        }
                    };

                    tasks.Add(Task.Run(bodyOfNewTask));
                }

                await Task.WhenAll(tasks).ConfigureAwait(false);
            }

            return Enumerable
                .Range(0, resultsByPositionInSource.Count)
                .Select(position => resultsByPositionInSource[position])
                .ToArray();
        }

        private static void EnsureValidMaxDegreeOfParallelism(int? maxDegreeOfParallelism)
        {
            if (maxDegreeOfParallelism <= 0)
            {
                throw new ArgumentOutOfRangeException(
                    nameof(maxDegreeOfParallelism),
                    $"Invalid value for the maximum degree of parallelism: {maxDegreeOfParallelism}. The maximum degree of parallelism must be a positive integer.");
            }
        }
    }
} 

Answer 1

Rather than writing something custom, you could use the TLP Dataflow library.

public static Task ForEachAsync<TSource>(
    this IEnumerable<TSource> items,
    Func<TSource, Task> action,
    int maxDegreesOfParallelism)
{
    var actionBlock = new ActionBlock<TSource>(action, new ExecutionDataflowBlockOptions
    {
        MaxDegreeOfParallelism = maxDegreesOfParallelism
    });

    foreach (var item in items)
    {
        actionBlock.Post(item);
    }

    actionBlock.Complete();

    return actionBlock.Completion;
}

Check out this Fiddle to see it in action.

EDIT

If you need the results:

public static async Task<IEnumerable<TResult>> ForEachAsync<TSource, TResult>(
    this IEnumerable<TSource> items,
    Func<TSource, Task<TResult>> action,
    int maxDegreesOfParallelism)
{
    var transformBlock = new TransformBlock<TSource, TResult>(action, new ExecutionDataflowBlockOptions
    {
        MaxDegreeOfParallelism = maxDegreesOfParallelism
    });

    var bufferBlock = new BufferBlock<TResult>();

    using (transformBlock.LinkTo(bufferBlock, new DataflowLinkOptions {PropagateCompletion = true}))
    {
        foreach (var item in items)
        {
            transformBlock.Post(item);
        }

        transformBlock.Complete();
        await transformBlock.Completion;
    }

    bufferBlock.TryReceiveAll(out var result);
    return result;
}

Answer 2

I see one big problem and one smaller problem with this code.

The big problem is the complete lack of handling of exceptions. If just one of the tasks fails, there's no way to get the results of any of the rest of the tasks, because an AggregateException is going to be thrown out. The best solution to this problem that I've found is to make the return type be IEnumerable<Task<TResult>> rather than IEnumerable<TResult>/TResult[]. That way you encapsulate neatly for each task its result if it succeeded, or its exception if it didn't.

The smaller problem is that it's overly complicated. There's no need for partitioners or semaphores. The same result can be obtained using Task.WhenAny to wait for the first of the active tasks to finish, whereupon you replace it with a new task and loop:

    public static async Task<IEnumerable<Task<TResult>>> ForEachAsync<TSource, TResult>(
        this IEnumerable<TSource> source,
        Func<TSource, Task<TResult>> operation,
        int maxParallel = int.MaxValue)
    {
        var result = new List<Task<TResult>>();
        var active = new HashSet<Task<TResult>>();
        var enumerator = source.GetEnumerator();
        while (active.Count < maxParallel && enumerator.MoveNext())
        {
            var task = operation(enumerator.Current);
            result.Add(task);
            active.Add(task);
        }

        while (active.Count > 0)
        {
            Task finished = await Task.WhenAny(active).ConfigureAwait(false);
            active.Remove(finished as Task<TResult>);

            if (enumerator.MoveNext())
            {
                var task = operation(enumerator.Current);
                result.Add(task);
                active.Add(task);
            }
        }

        return result;
    }