.NET Fiddle

CSharp · Console

/Program.cs

// @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}");
    }
}