Merging Parallel IAsyncEnumberables to Eliminate the Weakest Link

[toupswork]

I am thrilled that the .NET team has decided to invest in adding IAsyncEnumberable support; i.e., linq-like extension methods, for now.

IAsyncEnumberable is very important to me because I write a lot of APIs, and I want to squeeze out every last drop of performance.

What I love about them is how much easier they make it to perform computation and I/O concurrently on the same request. Like while the hardware is buffering to RAM via DMA or MMIO activity while .NET waits for the IOCP, the CPU can perform computation like JSON serialization or other computation concurrently on the same request, while minimizing memory usage, sense data in ≈ data out.

One of the more important aspects is merging IAsyncEnumberables. For example, I just finished a project that involved producing UI endpoint that needed to aggregate and multiplex a dozen API calls along with a DB query; then I did custom JSON serialization with varying width arrays using a null presense mask to save payload size.

Anyway, here is a custom Merge method I wrote until the team can get around to writing a better one. This one demonstrated 2x performance over the one that is part of Ix.Net in my own testing.

using System.Runtime.ExceptionServices;
using System.Threading.Channels;
using System.Runtime.CompilerServices;

namespace Application.Extensions;

public static class AsyncEnumerableExtensions
{
  extension<T>(AsyncEnumerable) where T : class
  {
    public static IAsyncEnumerable<T> MergeConcurrent(params IAsyncEnumerable<T>[] sources) =>
      MergeConcurrentCore(sources, sources.Length, sources.Length * 2);

    public static IAsyncEnumerable<T> MergeConcurrent(IAsyncEnumerable<T>[] sources, int concurrency) =>
      MergeConcurrentCore(sources, concurrency, sources.Length * 2);

    public static IAsyncEnumerable<T> MergeConcurrent(IAsyncEnumerable<T>[] sources, int concurrency, int capacity) =>
      MergeConcurrentCore(sources, concurrency, capacity);

  }


  private static IAsyncEnumerable<T> MergeConcurrentCore<T>(IAsyncEnumerable<T>[] sources, int concurrency, int capacity) where T: class
  {
    
    ArgumentNullException.ThrowIfNull(sources);
    ArgumentNullException.ThrowIfContainsNull(sources);

    if (sources.Length is 0) return AsyncEnumerable.Empty<T>();

    ArgumentOutOfRangeException.ThrowIfLessThan(concurrency, 1);
    ArgumentOutOfRangeException.ThrowIfLessThan(capacity, 1);

    return new MergeConcurrentAsyncIterator<T>(sources, Math.Min(concurrency, sources.Length), capacity);
  }
}


file sealed class MergeConcurrentAsyncIterator<T>(IEnumerable<IAsyncEnumerable<T>> sources, int concurrency, int capacity) : IAsyncEnumerable<T>
{
  private BoundedChannelOptions ChannelOptions { get; } = new(capacity)
  {
    SingleReader = true,
    SingleWriter = concurrency is 1,
    FullMode = BoundedChannelFullMode.Wait,
    AllowSynchronousContinuations = true,
  };

  private (ChannelWriter<T>, ChannelReader<T>) CreateChannel()
  {
    var channel = Channel.CreateBounded<T>(ChannelOptions);
    return (channel.Writer, channel.Reader);
  }

  private async Task ProduceAsync(ChannelWriter<T> writer, CancellationToken ct)
  {
    Exception? error = null;
    try
    {
      var pOptions = new ParallelOptions { CancellationToken = ct, MaxDegreeOfParallelism = concurrency };

      await Parallel.ForEachAsync(sources, pOptions, async ValueTask (source, token) => {
          await foreach(var item in source.WithCancellation(token))
            await writer.WriteAsync(item, token);
      });
    }
    catch (Exception ex)
    {
      error = ex;
    }
    finally
    {
      writer.TryComplete(error);
    }
  }
 
  public async IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken)
  {
    using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
    var (writer, reader) = CreateChannel();
    
    var producerTask = ProduceAsync(writer, cts.Token);

    try
    {
      var allItems = reader.ReadAllAsync(cts.Token);
      await using var enumerator = allItems.GetAsyncEnumerator(cancellationToken);
      while (await enumerator.MoveNextAsync()) yield return enumerator.Current;
    }
    finally
    {
      if (!producerTask.IsCompleted) await cts.CancelAsync();

      await producerTask;
    }
  }
}

[Lightczx]

I believe 'Concat'/'Append' is more likely to be adopted by .NET teams since there are Enumerable.Concat/Append things already.

[toupswork]

Well then they're leaving a lot of performance on the table when your running parallel I/O to acquire data --e.g., multiple api calls where pagination is done concurrently and not serially, multiple DB queries in parallel, or a mixture of DB and API calls.

Because when you use concat or the like, you are subject to the slowest link problem. If the I/O for one has 50ms pause while it's loading the data in a buffer, my solution will simply grap the next item already available--no weak link issue that dominates the entire throughput.

[Lightczx]

Ah, I see, so your method here is FCFS manner, and what I mentioned is simplly queueing one AsyncEnumerable after another.

[toupswork]

Ah, I see, so your method here is FCFS manner, and what I mentioned is simplly queueing one AsyncEnumerable after another.

Yes, that's exactly correct. Take a look at the implementation (link)

What is happening is that they are performing a serial enumeration starting with the first IAsyncEnumerable, the second, and so on. And when they loop them, it involves awaiting which can be instant if the input I/O stream (e.g., network card, database, filesystem) is coming in at the same rate or higher than the output stream. But if that turns out to not be the case, the entire stream is halted until the buffer is full again. And that might be very well fine if you are pulling from one stream, or you're pulling from multiple and doing so serially and it's not a lot of data.

BUT: That won't fly for me when I need to send down 20K rows of data to the UI in fewer than 250ms and when the column data of each row has to be aggregated from over a dozen APIs as well as 2 SQL queries, and because it's thousands of rows, the APIs I have to call require pagination, but I can't do serial pagination if I want to meet that 250ms requirement. So instead, I fetch the first page from the API, then use the pagination meta data from that to immediately construct the remaining pagination calls and fire them off all at once up to a specific concurrency count.

So now I have multiple IAsyncEnumerables that I have to essentially multiplex and I can't have one slow API call hold up the entire output stream.

Make sense? Then the other part is that this allows you me to keep the memory on my API equal to the difference between the data coming in and the data going out. Further, I can perform computation (like view model serialization) at the same time that I/O is occurring, so I am not limited to gather->compute->gather->repeat. It's more like gather x6 AND compute x6 -> repeat.

So what all of that means is that my API Server is squeezing out every single drop of performance that it possibly can

[TonyValenti]

I would definitely love something like this. We rolled our own (v1 was similar to this) but it would be great to have an official solution. The goal is to pull out whichever items are available first.

[timcassell]

I think the BCL should start with a non-parallel Merge function. The parallel version would fit nicely into the proposed parallel async Linq extensions (#114478).

P.S. I also have my own (non-parallel) efficient Merge implementation in my async library.