// @nuget: System.Linq.Async

// @nuget: System.Interactive.Async

using System;

using System.Collections.Generic;

using System.Linq;

using System.Threading;

using System.Threading.Tasks;

using System.Runtime.CompilerServices;

using System.Diagnostics;

​

public static class Program

{

    public static async Task Main()

    {

        const int bufferTimeSpan = 1000;

        const int bufferCount = 4;

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

            .ToAsyncEnumerable()

            .Do(async x =>

            {

                int delay = Random.Shared.Next(0, 6) * 100;

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

                await Task.Delay(delay);

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

            })

            .Buffer(TimeSpan.FromMilliseconds(bufferTimeSpan), bufferCount);

​

        Print($"Buffer(TimeSpan.FromMilliseconds({bufferTimeSpan}), count: {bufferCount})");

        Print($"Before await foreach");

        List<int[]> results = new();

        await foreach (IList<int> batch in query)

        {

            results.Add(batch.ToArray());

            Print($"--Received: [{String.Join(", ", batch)}]");

        }

        Print($"Results: {String.Join(", ", results.Select(b => $"[{String.Join(", ", b)}]"))}");

    }

​

    /// <summary>

    /// Splits the elements of a sequence into chunks that are emitted when either

    /// they are full, or a given amount of time has elapsed after requesting the

    /// previous chunk.

    /// </summary>

    public static async IAsyncEnumerable<IList<TSource>> Buffer<TSource>(

        this IAsyncEnumerable<TSource> source, TimeSpan timeSpan, int count,

        [EnumeratorCancellation] CancellationToken cancellationToken = default)

    {

        ArgumentNullException.ThrowIfNull(source);

        if (timeSpan < TimeSpan.FromMilliseconds(1.0))

            throw new ArgumentOutOfRangeException(nameof(timeSpan));

        if (count < 1) throw new ArgumentOutOfRangeException(nameof(count));

​

        using CancellationTokenSource linkedCts = CancellationTokenSource

            .CreateLinkedTokenSource(cancellationToken);

        PeriodicTimer timer = null;

        Task<bool> StartTimer()

        {

            timer = new(timeSpan);

            return timer.WaitForNextTickAsync().AsTask();

        }

        IAsyncEnumerator<TSource> enumerator = source

            .GetAsyncEnumerator(linkedCts.Token);

        Task<bool> moveNext = null;

        try

        {

            List<TSource> buffer = new();

            TSource[] ConsumeBuffer()

            {

                timer?.Dispose();

                TSource[] array = buffer.ToArray();

                buffer.Clear();

                if (buffer.Capacity > count) buffer.Capacity = count;

                return array;

            }

            Task<bool> timerTickTask = StartTimer();

            while (true)

            {

                if (moveNext is null)

                {

                    if (timerTickTask.IsCompleted)

                    {

                        Debug.Assert(timerTickTask.Result);

                        yield return ConsumeBuffer();

                        timerTickTask = StartTimer();

                    }

                    moveNext = enumerator.MoveNextAsync().AsTask();

                }

                if (!moveNext.IsCompleted)

                {

                    Task completedTask = await Task.WhenAny(moveNext, timerTickTask)

                        .ConfigureAwait(false);

                    if (ReferenceEquals(completedTask, timerTickTask))

                    {

                        Debug.Assert(timerTickTask.IsCompleted);

                        Debug.Assert(timerTickTask.Result);

                        yield return ConsumeBuffer();

                        timerTickTask = StartTimer();

                        continue;

                    }

                }

                Debug.Assert(moveNext.IsCompleted);

                bool moved = await moveNext.ConfigureAwait(false);

                moveNext = null;

                if (!moved) break;

                TSource item = enumerator.Current;

                buffer.Add(item);

                if (buffer.Count == count)

                {

                    yield return ConsumeBuffer();

                    timerTickTask = StartTimer();

                }

            }

            if (buffer.Count > 0) yield return ConsumeBuffer();

        }

        finally

        {

            // Cancel the enumerator, for more responsive completion.

            try { linkedCts.Cancel(); }

            finally

            {

                // The last moveNext must be completed before disposing.

                if (moveNext is not null && !moveNext.IsCompleted)

                    await Task.WhenAny(moveNext).ConfigureAwait(false);

                await enumerator.DisposeAsync().ConfigureAwait(false);

                timer?.Dispose();

            }

        }

    }

​

    private static void Print(object value)

    {

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

    }

}

​