How to await the results of an IAsyncEnumerable>, with a specific level of concurrency

Question

I have an asynchronous stream of tasks, that is generated by applying an async lambda to a stream of items:

IAsyncEnumerable<int> streamOfItems = AsyncEnumerable.Range(1, 10);
IAsyncEnumerable<Task<string>> streamOfTasks = streamOfItems.Select(async x =>
{
    await Task.Delay(100);
    return x.ToString();
})

The methods AsyncEnumerable.Range and Select above are provided from the System.Linq.Async package.

The result I want is a stream of results, expressed as an IAsyncEnumerable<string>. The results must be streamed in the same order as the originated tasks. Also the enumeration of the stream must be throttled, so than no more than a specified number of tasks are active at any given time.

I would like a solution in the form of an extension method on the IAsyncEnumerable<Task<T>> type, so that I could chain it multiple times and form a processing pipeline, similar in functionality with a TPL Dataflow pipeline, but expressed fluently. Below is the signature of the desirable extension method:

public async static IAsyncEnumerable<TResult> AwaitResults<TResult>(
    this IAsyncEnumerable<Task<TResult>> source,
    int concurrencyLevel);

Accepting also a CancellationToken as argument would be a nice feature.

---

Update: For completeness I am including an example of a fluent processing pipeline formed by chaining twice the AwaitResults method. This pipeline starts with a PLINQ block, just to demonstrate that mixing PLINQ and Linq.Async is possible.

int[] results = await Partitioner
    .Create(Enumerable.Range(1, 20), EnumerablePartitionerOptions.NoBuffering)
    .AsParallel()
    .AsOrdered()
    .WithDegreeOfParallelism(2)
    .WithMergeOptions(ParallelMergeOptions.NotBuffered)
    .Select(x =>
    {
        Thread.Sleep(100); // Simulate some CPU-bound operation
        return x;
    })
    .ToAsyncEnumerable()
    .Select(async x =>
    {
        await Task.Delay(300); // Simulate some I/O operation
        return x;
    })
    .AwaitResults(concurrencyLevel: 5)
    .Select(x => Task.Run(() =>
    {
        Thread.Sleep(100); // Simulate another CPU-bound operation
        return x;
    }))
    .AwaitResults(concurrencyLevel: 2)
    .ToArrayAsync();

Console.WriteLine($"Results: {String.Join(", ", results)}");

Expected output:

Results: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20

---

Note: In retrospect the AwaitResults method should probably be named Merge, and the concurrencyLevel argument should be named maxConcurrent, because its functionality resembles the Merge operator that exists in the Rx library. The System.Interactive.Async package does include an operator named Merge that produces IAsyncEnumerable<T>s, but none of its overloads operate on IAsyncEnumerable<Task<T>> sources. It operates on IEnumerable<IAsyncEnumerable<TSource>> and IAsyncEnumerable<IAsyncEnumerable<TSource>> sources. A parameter bufferCapacity could also be added, in order to control explicitly the size of the buffer needed for the awaiting/merging operation.

Answer 1

Here is an implementation of the AwaitResults method. It is based on a SemaphoreSlim for controlling the concurrency level, and on a Channel<Task<TResult>> for storing the running tasks. The source IAsyncEnumerable<Task<TResult>> sequence is enumerated on the ThreadPool, and the hot tasks are pushed into the channel. Each task is wrapped into a higher-level task, that releases the semaphore when it completes.

The main part of the method is the yielding loop, where the tasks are dequeued from the channel one by one, and then awaited sequentially. This way the results are yielded in the same order as the tasks in the source sequence.

/// <summary>
/// Enumerates a sequence of tasks and returns a sequence of results.
/// The source sequence is enumerated with a specific maximum concurrency.
/// The enumeration of the source is suspended when the internal buffer
/// has reached its full capacity. For unlimited capacity configure with -1.
/// </summary>
public async static IAsyncEnumerable<TResult> AwaitResults<TResult>(
    this IAsyncEnumerable<Task<TResult>> source,
    int maxConcurrency = 1,
    int capacity = -1,
    [EnumeratorCancellation]CancellationToken cancellationToken = default)
{
    ArgumentNullException.ThrowIfNull(source);
    if (maxConcurrency < 1)
        throw new ArgumentOutOfRangeException(nameof(maxConcurrency));
    if (capacity == -1) capacity = Int32.MaxValue;
    if (capacity < 1)
        throw new ArgumentOutOfRangeException(nameof(capacity));

    Channel<Task<TResult>> channel = Channel
        .CreateBounded<Task<TResult>>(capacity);
    using SemaphoreSlim semaphore = new(maxConcurrency - 1, maxConcurrency);
    using CancellationTokenSource linkedCts = CancellationTokenSource
        .CreateLinkedTokenSource(cancellationToken);

    Task producer = Task.Run(async () =>
    {
        try
        {
            await foreach (Task<TResult> task in source
                .WithCancellation(linkedCts.Token).ConfigureAwait(false))
            {
                await channel.Writer.WriteAsync(OnCompletionRelease(task))
                    .ConfigureAwait(false); // Without cancellation
                await semaphore.WaitAsync(linkedCts.Token)
                    .ConfigureAwait(false); // Acquire before MoveNextAsync
            }
        }
        catch (ChannelClosedException) { } // Ignore
        finally { channel.Writer.TryComplete(); }
    });

    async Task<TResult> OnCompletionRelease(Task<TResult> task)
    {
        try { return await task.ConfigureAwait(false); }
        finally { semaphore.Release(); }
    }

    try
    {
        await foreach (Task<TResult> task in channel.Reader.ReadAllAsync()
            .ConfigureAwait(false))
                yield return await task.ConfigureAwait(false);
        await producer.ConfigureAwait(false); // Propagate possible source error
    }
    finally
    {
        // Await all pending operations before completing the enumeration.
        try { linkedCts.Cancel(); }
        finally
        {
            if (!producer.IsCompleted)
            {
                // Unblock pending writes.
                channel.Writer.TryComplete();
                await Task.WhenAny(producer).ConfigureAwait(false);
            }
            if (!channel.Reader.Completion.IsCompleted)
            {
                // Drain the channel.
                while (channel.Reader.TryRead(out Task<TResult> task))
                    await Task.WhenAny(task).ConfigureAwait(false);
                Debug.Assert(channel.Reader.Completion.IsCompleted);
            }
        }
    }
}

The finally block at the bottom plays the important role of preventing any of the tasks to become fire-and-forget, in case the consumer of the resulting sequence abandons the enumeration prematurely. In case more than one tasks complete with failure, only the exception of the first faulted task will be propagated (first in order, not chronologically), and all other exceptions will trigger the TaskScheduler.UnobservedTaskException event.

Online demo.

Note: Canceling the cancellationToken may not cancel the whole operation instantaneously. For maximum responsiveness the same cancellationToken should be used for cancelling the individual tasks.