Buffer IAsyncEnumerable

[TonyValenti]

I have an AsyncEnumerable that pulls data from a (very slow) web service.
Once I get a result, I have to do some processing on it which also takes some time.
Sometimes I need only the first few results and sometimes I need every item.

Ideally I'd like a Buffer(n) extension method on IAsyncEnumerable that would pre-fetch up to n items.

As an alternative, is it safe to do the following?

1. Create a fixed-length Channel that blocks when full.
2. Task.Run() a method to fill up the channel and complete it.
3. Return the Channel.ReadAllAsync();

If I consume all items, I think this works fine, but my concern is that maybe the "fill it up" task would be blocked and live forever if I abandoned the enumerable after just pulling out a few items. Is my approach safe or is that a bad one?

[gfoidl]

do some processing on it which also takes some time ... sometimes I need every item

Just use Channels. await foreach the IAsyncEnumerable and write to the channel, the consumer (processing) is the channel-reader.

Sometimes I need only the first few results

Why not just create a T[count] (or a list) and fill this array in await foreach over the IAsyncEnumerable. You can track the item-count, there's no race, so it's safe. When the expected number of items are "buffered", hand this buffer over to the consumer.

It's possible with channels, too, but it sounds more complicated (KISS).
Also note that instead of "Channel that blocks when full" you could drop the newest item, so that there's less back-pressure on the IAsyncEnumberable.

[TonyValenti]

In the example you gave, if I understand correctly, it would essentially function as Chunk.

What I'm wanting to do is essentially still work over items in individually in the enumerable but I want the items cached locally if they are available so that I'm eliminating the network delay.

[gfoidl]

it would essentially function as Chunk

Not really, it is

Sometimes I need only the first few results

No you write about a different requirement: "work over items in individually in the enumerable but I want the items cached locally if they are available".
So when I got you right, see this trivial example:

using System.Threading.Channels;

Channel<IEnumerable<int>> channel = Channel.CreateUnbounded<IEnumerable<int>>();
IAsyncEnumerable<int> data = FetchData();
Task bufferTask = Buffer(data);
Task consumeTask = Consume(channel.Reader.ReadAllAsync());

await Task.WhenAll(bufferTask, consumeTask);

static async IAsyncEnumerable<int> FetchData(int count = 10)
{
    for (int i = 0; i < count; ++i)
    {
        await Task.Delay(Random.Shared.Next(300, 500));

        yield return i + 1;
    }
}

async Task Buffer(IAsyncEnumerable<int> source, int bufferSize = 3)
{
    List<int> buffer = new(capacity: bufferSize);

    await foreach (int item in source)
    {
        buffer.Add(item);

        if (buffer.Count == bufferSize)
        {
            await channel.Writer.WriteAsync(buffer.ToArray());
            buffer.Clear();
        }
    }

    if (buffer.Count > 0)
    {
        await channel.Writer.WriteAsync(buffer.ToArray());
    }

    channel.Writer.Complete();
}

static async Task Consume(IAsyncEnumerable<IEnumerable<int>> bufferedData)
{
    await foreach (IEnumerable<int> data in bufferedData)
    {
        foreach (int item in data)
        {
            Console.Write($"{item} ");
        }
        Console.WriteLine();
    }
}

This could be labeled as "chunking", and is what you want?!

This is quite imperative, so of course you could make this more OOP-like.
As you are under control of all the parts here, buffer.ToArray-allocation could be avoided too by re-using that array -- but I'd do this only if profiling says so.

---

Or do you want to cache (buffer, chunk, cache, what are we talking about now?) the whole IAsyncEnumerable? Then just use a list.

[TonyValenti]

@gfoidl That wasn't exactly what I wanted but you inspired me toward the answer I was looking for. Here is what I wanted:

public static class AsyncExtensions {
    public static async IAsyncEnumerable<T> Buffer<T>(this IAsyncEnumerable<T> This, int Count) {

        var options = new BoundedChannelOptions(Count)
        {
            SingleReader = true,
            SingleWriter = true,
            AllowSynchronousContinuations = false,
        };

        var chan = System.Threading.Channels.Channel.CreateBounded<T>(options);
        var reader = chan.Reader;
        var writer = chan.Writer;

        var IE = This.GetAsyncEnumerator();

        async Task<bool> FillAsync() {
            var More = true;
            while (More && reader.Count < Count) {
                try {
                    More = await IE.MoveNextAsync()
                        ;

                    if (More) {
                        await writer.WriteAsync(IE.Current);
                    }
                    else {
                        writer.Complete();
                    }
                } catch(Exception ex) {
                    More = false;
                    writer.Complete(ex);
                }
            }

            return More;
        }

        var FillTask = Task.FromResult(true);
        var KeepFilling = true;

        var Continue = true;
        while (Continue) {

            if (KeepFilling && FillTask.IsCompleted && reader.Count < Count) {
                KeepFilling = await FillTask;

                if (KeepFilling) {
                    FillTask = Task.Run(() => FillAsync());
                }
            }

            Continue =  await reader.WaitToReadAsync();

            if (Continue) {
                var ret = await reader.ReadAsync();
                yield return ret;
            }

        }



    }
}

[angularsen]

Maybe not what OP asked for, but I stumbled in here looking for chunking/batching similar to System.Linq.Enumerable.Chunk and I found this gist helpful:
https://gist.github.com/rquackenbush/96157ec6940ea1dd33774670cdc4786c

    public static async IAsyncEnumerable<T[]> BatchAsync<T>(
        this IAsyncEnumerable<T> source,
        int batchSize,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        var batch = new List<T>(batchSize);

        await foreach (var item in source.WithCancellation(cancellationToken))
        {
            if (cancellationToken.IsCancellationRequested)
                yield break;

            batch.Add(item);

            if (batch.Count >= batchSize)
            {
                yield return batch.ToArray();

                batch.Clear();
            }
        }

        if (batch.Count > 0)
            yield return batch.ToArray();
    }

Some other implementations as well:

• MoreAsyncLINQ Batch() (src: Batch.cs)
• System.Interactive.Async Buffer() (src: Buffer.cs)

Found in similar discussion of MoreLINQ package:
morelinq/MoreLINQ#773 (comment)