private static async Task AwaitMoveNextAsyncAndDispose <T>(Task <bool> moveNextAsync, IAsyncEnumerator <T> enumerator)
{
if (enumerator != null)
{
await using (AsyncEnumerableExtensions.ConfigureAwait(enumerator, false))
{
if (moveNextAsync != null)
{
await moveNextAsync.ConfigureAwait(false);
}
}
}
}
public static IAsyncEnumerable <TSource> Amb <TSource>(params IAsyncEnumerable <TSource>[] sources)
{
if (sources == null)
{
throw Error.ArgumentNull(nameof(sources));
}
#if USE_ASYNC_ITERATOR
return(AsyncEnumerable.Create(Core));
async IAsyncEnumerator <TSource> Core(CancellationToken cancellationToken)
{
//
// REVIEW: See remarks on binary overload for changes compared to the original.
//
var n = sources.Length;
var enumerators = new IAsyncEnumerator <TSource> [n];
var moveNexts = new Task <bool> [n];
try
{
for (var i = 0; i < n; i++)
{
var enumerator = sources[i].GetAsyncEnumerator(cancellationToken);
enumerators[i] = enumerator;
moveNexts[i] = enumerator.MoveNextAsync().AsTask();
}
}
catch
{
var cleanup = new Task[n];
for (var i = 0; i < n; i++)
{
cleanup[i] = AwaitMoveNextAsyncAndDispose(moveNexts[i], enumerators[i]);
}
await Task.WhenAll(cleanup).ConfigureAwait(false);
throw;
}
var moveNextWinner = await Task.WhenAny(moveNexts).ConfigureAwait(false);
//
// NB: The use of IndexOf is fine. If task N completed by returning a ValueTask<bool>
// which is equivalent to the task returned by task M (where M < N), AsTask may
// return the same reference (e.g. due to caching of completed Boolean tasks). In
// such a case, IndexOf will find task M rather than N in the array, but both have
// an equivalent completion state (because they're reference equal). This only leads
// to a left-bias in selection of sources, but given Amb's "ambiguous" nature, this
// is acceptable.
//
var winnerIndex = Array.IndexOf(moveNexts, moveNextWinner);
var winner = enumerators[winnerIndex];
var loserCleanupTasks = new List <Task>(n - 1);
for (var i = 0; i < n; i++)
{
if (i != winnerIndex)
{
var loserCleanupTask = AwaitMoveNextAsyncAndDispose(moveNexts[i], enumerators[i]);
loserCleanupTasks.Add(loserCleanupTask);
}
}
var cleanupLosers = Task.WhenAll(loserCleanupTasks);
try
{
await using (AsyncEnumerableExtensions.ConfigureAwait(winner, false))
{
if (!await moveNextWinner.ConfigureAwait(false))
{
yield break;
}
yield return(winner.Current);
while (await winner.MoveNextAsync().ConfigureAwait(false))
{
yield return(winner.Current);
}
}
}
finally
{
await cleanupLosers.ConfigureAwait(false);
}
}
#else
return(new AmbAsyncIteratorN <TSource>(sources));
#endif
}
public static IAsyncEnumerable <TSource> Amb <TSource>(this IAsyncEnumerable <TSource> first, IAsyncEnumerable <TSource> second)
{
if (first == null)
{
throw Error.ArgumentNull(nameof(first));
}
if (second == null)
{
throw Error.ArgumentNull(nameof(second));
}
#if USE_ASYNC_ITERATOR
return(AsyncEnumerable.Create(Core));
async IAsyncEnumerator <TSource> Core(CancellationToken cancellationToken)
{
IAsyncEnumerator <TSource> firstEnumerator = null;
IAsyncEnumerator <TSource> secondEnumerator = null;
Task <bool> firstMoveNext = null;
Task <bool> secondMoveNext = null;
try
{
//
// REVIEW: We start both sequences unconditionally. An alternative implementation could be to just stick
// to the first sequence if we notice it already has a value (or exception) available. This would
// be similar to Task.WhenAny behavior (see CommonCWAnyLogic in dotnet/coreclr). We could consider
// adding a WhenAny combinator that does exactly that. We can even avoid calling AsTask.
//
firstEnumerator = first.GetAsyncEnumerator(cancellationToken);
firstMoveNext = firstEnumerator.MoveNextAsync().AsTask();
//
// REVIEW: Order of operations has changed here compared to the original, but is now in sync with the N-ary
// overload which performs GetAsyncEnumerator/MoveNextAsync in pairs, rather than phased.
//
secondEnumerator = second.GetAsyncEnumerator(cancellationToken);
secondMoveNext = secondEnumerator.MoveNextAsync().AsTask();
}
catch
{
// NB: AwaitMoveNextAsyncAndDispose checks for null for both arguments, reducing the need for many null
// checks over here.
var cleanup = new[]
{
AwaitMoveNextAsyncAndDispose(secondMoveNext, secondEnumerator),
AwaitMoveNextAsyncAndDispose(firstMoveNext, firstEnumerator)
};
await Task.WhenAll(cleanup).ConfigureAwait(false);
throw;
}
//
// REVIEW: Consider using the WhenAny combinator defined for Merge in TaskExt, which would avoid the need
// to convert to Task<bool> prior to calling Task.WhenAny.
//
var moveNextWinner = await Task.WhenAny(firstMoveNext, secondMoveNext).ConfigureAwait(false);
//
// REVIEW: An alternative option is to call DisposeAsync on the other and await it, but this has two drawbacks:
//
// 1. Concurrent DisposeAsync while a MoveNextAsync is in flight.
// 2. The winner elected by Amb is blocked to yield results until the loser unblocks.
//
IAsyncEnumerator <TSource> winner;
Task disposeLoser;
if (moveNextWinner == firstMoveNext)
{
winner = firstEnumerator;
disposeLoser = AwaitMoveNextAsyncAndDispose(secondMoveNext, secondEnumerator);
}
else
{
winner = secondEnumerator;
disposeLoser = AwaitMoveNextAsyncAndDispose(firstMoveNext, firstEnumerator);
}
try
{
await using (AsyncEnumerableExtensions.ConfigureAwait(winner, false))
{
if (!await moveNextWinner.ConfigureAwait(false))
{
yield break;
}
yield return(winner.Current);
while (await winner.MoveNextAsync().ConfigureAwait(false))
{
yield return(winner.Current);
}
}
}
finally
{
//
// REVIEW: This behavior differs from the original implementation in that we never discard any in flight
// asynchronous operations. If an exception occurs while enumerating the winner, it can be
// subsumed by an exception thrown due to cleanup of the loser. Also, if Amb is used to deal with
// a potentially long-blocking sequence, this implementation would transfer this blocking behavior
// to the resulting sequence. However, it seems never discarding a non-completed task should be a
// general design tenet, and fire-and-forget dispose behavior could be added as another "unsafe"
// operator, so all such sins are made explicit by the user. Nonetheless, this change is a breaking
// change for Ix Async.
//
await disposeLoser.ConfigureAwait(false);
}
}
#else
return(new AmbAsyncIterator <TSource>(first, second));
#endif
}
