How to query two IAsyncEnumerables asynchronously

Viewed 608

I have two methods connected to two different sources of Foos which return two IAsyncEnumerable<Foo>. I need to fetch all Foos from both sources before being able to process them .

Problem : I would like to query both sources simultaneously (asynchronously), ie. not waiting for Source1 to complete the enumeration before starting to enumerate Source2. From my understanding, this is what happens into the method SequentialSourcesQuery example below, am I right?

With regular tasks, I would just start the first Task, then the second one, and call a await Task.WhenAll. But I am a bit confused on how to handle IAsyncEnumerable.

public class FoosAsync
{
    public async IAsyncEnumerable<Foo> Source1() { }

    public async IAsyncEnumerable<Foo> Source2() { }

    public async Task<List<Foo>> SequentialSourcesQuery()
    {
        List<Foo> foos = new List<Foo>();

        await foreach (Foo foo1 in Source1())
        {
            foos.Add(foo1);
        }

        await foreach (Foo foo2 in Source2())
        { //doesn't start until Source1 completed the enumeration? 
            foos.Add(foo2);
        }

        return foos;
    }
}
3 Answers

You could take advantage of the libraries System.Linq.Async and System.Interactive.Async (owned by the RxTeam who are part of the .NET Foundation). They contain operators like Merge and ToListAsync that could solve your problem easily.

// Merges elements from all of the specified async-enumerable sequences
// into a single async-enumerable sequence.
public static IAsyncEnumerable<TSource> Merge<TSource>(
    params IAsyncEnumerable<TSource>[] sources);

// Creates a list from an async-enumerable sequence.
public static ValueTask<List<TSource>> ToListAsync<TSource>(
    this IAsyncEnumerable<TSource> source,
    CancellationToken cancellationToken = default);

Putting everything together:

public Task<List<Foo>> SequentialSourcesQuery()
{
    return AsyncEnumerableEx.Merge(Source1(), Source2()).ToListAsync().AsTask();
}

By aware that these libraries have a focus on providing a rich set of features, and not on performance or efficiency. So if top-notch performance is important for your use case, niki.kante's solution will most probably outperform the above operator-based approach.

You could write another async local method which returns Task.

Func<IAsyncEnumerable<Foo>, Task<List<Foo>>> readValues = async (values) => {
        List<Foo> foos = new List<Foo>();
        await foreach (Foo foo1 in values)
        {
            foos.Add(foo1);
        }        
        return foos;
};

and call it like this:

Task<List<Foo>> task1 = readValues(Source1());
Task<List<Foo>> task2 = readValues(Source2());

await Task.WhenAll(task1, task2);

Whole code would be:

public class FoosAsync
{
    public async IAsyncEnumerable<Foo> Source1() { }

    public async IAsyncEnumerable<Foo> Source2() { }

    public async Task<List<Foo>> SequentialSourcesQuery()
    {
        var asyncEnumerator = Source1().GetAsyncEnumerator();
        Func<IAsyncEnumerable<Foo>, Task<List<Foo>>> readValues = async (values) => {
            List<Foo> foos2 = new List<Foo>();
            await foreach (Foo foo in values)
            {
                foos2.Add(foo);
            }        
            return foos2;
        };
        
        Task<List<Foo>> task1 = readValues(Source1());
        Task<List<Foo>> task2 = readValues(Source2());
        
        await Task.WhenAll(task1, task2);
        
        List<Foo> foos = new List<Foo>(task1.Result.Count + task2.Result.Count);
        foos.AddRange(task1.Result);
        foos.AddRange(task2.Result);

        return foos;
    }
}

If you have two IAsyncEnumerable<T> as a source and don't care about the order of incoming data, you could use a method like the following to interleave your data.

public static class AsyncEnumerableExt
{
    public static async IAsyncEnumerable<T> Interleave<T>(this IAsyncEnumerable<T> first, IAsyncEnumerable<T> second)
    {
        var enum1 = first.GetAsyncEnumerator();
        var enum2 = second.GetAsyncEnumerator();

        var nextWait1 = enum1.MoveNextAsync().AsTask();
        var nextWait2 = enum2.MoveNextAsync().AsTask();

        do
        {
            var task = await Task.WhenAny(nextWait1, nextWait2).ConfigureAwait(false);

            if (task == nextWait1)
            {
                yield return enum1.Current;

                nextWait1 = !await task.ConfigureAwait(false) ? null : enum1.MoveNextAsync().AsTask();
            }
            else if (task == nextWait2)
            {
                yield return enum2.Current;

                nextWait2 = !await task.ConfigureAwait(false) ? null : enum2.MoveNextAsync().AsTask();
            }
        } while (nextWait1 != null && nextWait2 != null);

        while (nextWait1 != null)
        {
            if (!await nextWait1.ConfigureAwait(false))
            {
                nextWait1 = null;
            }
            else
            {
                yield return enum1.Current;
                nextWait1 = enum1.MoveNextAsync().AsTask();
            }
        }

        while (nextWait2 != null)
        {
            if (!await nextWait2.ConfigureAwait(false))
            {
                nextWait2 = null;
            }
            else
            {
                yield return enum2.Current;
                nextWait2 = enum2.MoveNextAsync().AsTask();
            }
        }
    }
}

then you can consume the data with one await foreach and store the data in a list.

Related