Merging IAsyncEnumberables

[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;
    }
  }
}