using System;

using System.Collections.Generic;

using System.Linq;

using System.Threading;

using System.Threading.Tasks;

using System.Threading.Channels;

using System.Runtime.CompilerServices;

using System.Diagnostics;

​

public static class Program

{

    public static async Task Main()

    {

        IAsyncEnumerable<int> query = Enumerable.Range(1, 10)

            .ToAsyncEnumerable()

            .Select(async x =>

            {

                int delay = Random.Shared.Next(1, 11) * 50;

                Print($"Task #{x} started, Delay: {delay}");

                await Task.Delay(delay);

                Print($"Task #{x} finished");

                return x;

            })

            .AwaitResults(maxConcurrency: 2);

​

        List<int> results = new();

        Print($"Before await foreach");

        await foreach (int item in query)

        {

            results.Add(item);

            Print($"--Received: {item}");

        }

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

    }

​

    /// <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);

                }

            }

        }

    }

​

    private static void Print(object value)

    {

        Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} [{Thread.CurrentThread.ManagedThreadId}] > {value}");

    }

}

​