Benchmarking various ways to await an array of `ValueTask`'s

Program.cs

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using System.Threading.Tasks.Sources;
using BenchmarkDotNet.Attributes;
using BenchmarkDotNet.Configs;
using BenchmarkDotNet.Order;
using BenchmarkDotNet.Running;
using BenchmarkDotNet.Reports;
using ValueTaskSupplement;
using BenchmarkDotNet.Columns;
using BenchmarkDotNet.Diagnosers;
using System.Runtime.CompilerServices;

BenchmarkRunner.Run<Benchmarks>();

[ConfigSource]
public class Benchmarks
{
    private sealed class ConfigSourceAttribute : Attribute, IConfigSource
    {
        public IConfig Config { get; }

        public ConfigSourceAttribute() => Config = new SimpleConfig();
    }

    [Params(2, 10)]
    public int N { get; set; }

    [Benchmark]
    public async ValueTask Loop() 
    {
        var tasks = new ValueTask[N];
        for (int i = 0; i < tasks.Length; i++)
        {
            tasks[i] = DoWork();
        }

        for (int i = 0; i < tasks.Length; i++)
        {
            await tasks[i];
        }
    }

    [Benchmark]
    public async ValueTask LoopPooled() 
    {
        var tasks = ArrayPool<ValueTask>.Shared.Rent(N);
        for (int i = 0; i < N; i++)
        {
            tasks[i] = DoWork();
        }

        for (int i = 0; i < N; i++)
        {
            await tasks[i];
        }

        ArrayPool<ValueTask>.Shared.Return(tasks);
    }

    [Benchmark(Baseline = true)]
    [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))]
    public async ValueTask LoopPooledMethodBuilder() 
    {
        var tasks = ArrayPool<ValueTask>.Shared.Rent(N);
        for (int i = 0; i < N; i++)
        {
            tasks[i] = DoWork();
        }

        for (int i = 0; i < N; i++)
        {
            await tasks[i];
        }

        ArrayPool<ValueTask>.Shared.Return(tasks);
    }

    [Benchmark]
    public async ValueTask TaskWhenAll()
    {
        var tasks = new Task[N];
        for (int i = 0; i < tasks.Length; i++)
        {
            tasks[i] = DoWork().AsTask();
        }

        await Task.WhenAll(tasks);
    }

    [Benchmark]
    public async ValueTask TaskWhenAllPooled()
    {
        var tasks = ArrayPool<Task>.Shared.Rent(N);
        for (int i = 0; i < N; i++)
        {
            tasks[i] = DoWork().AsTask();
        }

        await Task.WhenAll(tasks.Take(N));

        ArrayPool<Task>.Shared.Return(tasks);
    }

    [Benchmark]
    [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))]
    public async ValueTask TaskWhenAllPooledMethodBuilder()
    {
        var tasks = ArrayPool<Task>.Shared.Rent(N);
        for (int i = 0; i < N; i++)
        {
            tasks[i] = DoWork().AsTask();
        }

        await Task.WhenAll(tasks.Take(N));

        ArrayPool<Task>.Shared.Return(tasks);
    }

    [Benchmark]
    public ValueTask ValueTaskSupplement()
    {
        var tasks = new ValueTask[N];
        for (int i = 0; i < tasks.Length; i++)
        {
            tasks[i] = DoWork();
        }

        return ValueTaskEx.WhenAll(tasks);
    }

    [Benchmark]
    public ValueTask ValueTaskSource() 
    {
        var tasks = ArrayPool<ValueTask>.Shared.Rent(N);

        for (int i = 0; i < N; i++)
        {
            tasks[i] = DoWork();
        }

        return PooledValueTaskWhenAllPromise.Wait(tasks, N);
    }

    [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))]
    internal static async ValueTask DoWork()
    {
        _ = await Task.Run(() => 1 + 3);
    }
}

internal sealed class PooledValueTaskWhenAllPromise : IValueTaskSource
{
    private ManualResetValueTaskSourceCore<int> _source = new () { RunContinuationsAsynchronously = true };

    private static PooledValueTaskWhenAllPromise? _cache;

    private ValueTask[] _tasks = Array.Empty<ValueTask>();
    private int _tasksCount = 0;
    private List<Exception> exceptions = new ();
    private int _completed;

    internal static ValueTask Wait(ValueTask[] tasks, int count) 
    {
        var vts = _cache;
        if (vts is not null)
        {
            _cache = null;
            vts.Init(tasks, count);
            return new ValueTask(vts, vts._source.Version);
        }
        vts = new PooledValueTaskWhenAllPromise();
        vts.Init(tasks, count);
        return new ValueTask(vts, vts._source.Version);
    }

    private void Init(ValueTask[] tasks, int count)
    {
        _tasks = tasks;
        _tasksCount = count;
        _completed = 0;
        
        var span = tasks.AsSpan(0, count);
        for (int i = 0; i < span.Length; i++)
        {
            ref readonly var task = ref span[i];
            var awaiter = task.GetAwaiter();
            if (!awaiter.IsCompleted)
            {
                awaiter.OnCompleted(() => 
                {
                    try
                    {
                        awaiter.GetResult();
                    }
                    catch (Exception ex)
                    {
                        exceptions.Add(ex);
                    }
                    
                    if (Interlocked.Increment(ref _completed) == _tasksCount)
                    {
                        _source.SetResult(0);
                    }
                });
            }
            else 
            {
                try
                {
                    awaiter.GetResult();
                }
                catch (Exception ex)
                {
                    exceptions.Add(ex);
                }
                if (Interlocked.Increment(ref _completed) == _tasksCount)
                {
                    _source.SetResult(0);
                }
            }
        }
    }

    private void Return()
    {
        _source.Reset();
        exceptions.Clear();
        if (exceptions.Capacity > 4)
            exceptions.Capacity = 4;
        ArrayPool<ValueTask>.Shared.Return(_tasks);
        _tasks = Array.Empty<ValueTask>();
        _tasksCount = 0;

        if (_cache is null)
            _cache = this;
    }

    private PooledValueTaskWhenAllPromise()
    {
    }

    public void GetResult(short token)
    {
        try 
        {
            if (exceptions.Count > 0)
            {
                throw new AggregateException(exceptions);
            }
            _source.GetResult(token);
        } 
        finally 
        {
            Return();
        }
    }

    public ValueTaskSourceStatus GetStatus(short token) => 
        _source.GetStatus(token);

    public void OnCompleted(Action<object?> continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) => 
        _source.OnCompleted(continuation, state, token, flags);
}

internal sealed class SimpleConfig : ManualConfig
{
    public SimpleConfig()
    {
        SummaryStyle = SummaryStyle.Default.WithRatioStyle(RatioStyle.Trend);
        AddColumn(RankColumn.Arabic);
        Orderer = new DefaultOrderer(SummaryOrderPolicy.FastestToSlowest, MethodOrderPolicy.Declared);
        AddDiagnoser(MemoryDiagnoser.Default);
    }
}

scratch-dotnet.csproj

<Project Sdk="Microsoft.NET.Sdk">

  <PropertyGroup>
    <OutputType>Exe</OutputType>
    <TargetFramework>net8.0</TargetFramework>
    <RootNamespace>scratch_dotnet</RootNamespace>
    <ImplicitUsings>disable</ImplicitUsings>
    <Nullable>enable</Nullable>
    <RestoreAdditionalProjectSources>
      https://www.myget.org/F/benchmarkdotnet/api/v3/index.json
    </RestoreAdditionalProjectSources>
  </PropertyGroup>

  <ItemGroup>
    <PackageReference Include="BenchmarkDotNet" Version="0.13.13-nightly.20240401.151" />
    <PackageReference Include="ValueTaskSupplement" Version="1.1.0" />
  </ItemGroup>

</Project>

DoWork() = Task.Yield()

DoWork() = Averaging 100k random doubles