Add UniTaskAsyncEnumerable.Merge
parent
c170af5642
commit
f303d9d7e8
|
@ -0,0 +1,137 @@
|
||||||
|
using System;
|
||||||
|
using System.Threading;
|
||||||
|
using System.Threading.Tasks;
|
||||||
|
using Cysharp.Threading.Tasks;
|
||||||
|
using Cysharp.Threading.Tasks.Linq;
|
||||||
|
using FluentAssertions;
|
||||||
|
using Xunit;
|
||||||
|
|
||||||
|
namespace NetCoreTests.Linq
|
||||||
|
{
|
||||||
|
public class MergeTest
|
||||||
|
{
|
||||||
|
[Fact]
|
||||||
|
public async Task TwoSource()
|
||||||
|
{
|
||||||
|
var semaphore = new SemaphoreSlim(1, 1);
|
||||||
|
|
||||||
|
var a = UniTaskAsyncEnumerable.Create<string>(async (writer, _) =>
|
||||||
|
{
|
||||||
|
await UniTask.SwitchToThreadPool();
|
||||||
|
|
||||||
|
await semaphore.WaitAsync();
|
||||||
|
await writer.YieldAsync("A1");
|
||||||
|
semaphore.Release();
|
||||||
|
|
||||||
|
await semaphore.WaitAsync();
|
||||||
|
await writer.YieldAsync("A2");
|
||||||
|
semaphore.Release();
|
||||||
|
});
|
||||||
|
|
||||||
|
var b = UniTaskAsyncEnumerable.Create<string>(async (writer, _) =>
|
||||||
|
{
|
||||||
|
await UniTask.SwitchToThreadPool();
|
||||||
|
|
||||||
|
await semaphore.WaitAsync();
|
||||||
|
await writer.YieldAsync("B1");
|
||||||
|
await writer.YieldAsync("B2");
|
||||||
|
semaphore.Release();
|
||||||
|
|
||||||
|
await semaphore.WaitAsync();
|
||||||
|
await writer.YieldAsync("B3");
|
||||||
|
semaphore.Release();
|
||||||
|
});
|
||||||
|
|
||||||
|
var result = await a.Merge(b).ToArrayAsync();
|
||||||
|
result.Should().Equal("A1", "B1", "B2", "A2", "B3");
|
||||||
|
}
|
||||||
|
|
||||||
|
[Fact]
|
||||||
|
public async Task ThreeSource()
|
||||||
|
{
|
||||||
|
var semaphore = new SemaphoreSlim(0, 1);
|
||||||
|
|
||||||
|
var a = UniTaskAsyncEnumerable.Create<string>(async (writer, _) =>
|
||||||
|
{
|
||||||
|
await UniTask.SwitchToThreadPool();
|
||||||
|
|
||||||
|
await semaphore.WaitAsync();
|
||||||
|
await writer.YieldAsync("A1");
|
||||||
|
semaphore.Release();
|
||||||
|
|
||||||
|
await semaphore.WaitAsync();
|
||||||
|
await writer.YieldAsync("A2");
|
||||||
|
semaphore.Release();
|
||||||
|
});
|
||||||
|
|
||||||
|
var b = UniTaskAsyncEnumerable.Create<string>(async (writer, _) =>
|
||||||
|
{
|
||||||
|
await UniTask.SwitchToThreadPool();
|
||||||
|
|
||||||
|
await semaphore.WaitAsync();
|
||||||
|
await writer.YieldAsync("B1");
|
||||||
|
await writer.YieldAsync("B2");
|
||||||
|
semaphore.Release();
|
||||||
|
|
||||||
|
await semaphore.WaitAsync();
|
||||||
|
await writer.YieldAsync("B3");
|
||||||
|
semaphore.Release();
|
||||||
|
});
|
||||||
|
|
||||||
|
var c = UniTaskAsyncEnumerable.Create<string>(async (writer, _) =>
|
||||||
|
{
|
||||||
|
await UniTask.SwitchToThreadPool();
|
||||||
|
|
||||||
|
await writer.YieldAsync("C1");
|
||||||
|
semaphore.Release();
|
||||||
|
});
|
||||||
|
|
||||||
|
var result = await a.Merge(b, c).ToArrayAsync();
|
||||||
|
result.Should().Equal("C1", "A1", "B1", "B2", "A2", "B3");
|
||||||
|
}
|
||||||
|
|
||||||
|
[Fact]
|
||||||
|
public async Task Throw()
|
||||||
|
{
|
||||||
|
var a = UniTaskAsyncEnumerable.Create<string>(async (writer, _) =>
|
||||||
|
{
|
||||||
|
await writer.YieldAsync("A1");
|
||||||
|
|
||||||
|
});
|
||||||
|
|
||||||
|
var b = UniTaskAsyncEnumerable.Create<string>(async (writer, _) =>
|
||||||
|
{
|
||||||
|
throw new UniTaskTestException();
|
||||||
|
});
|
||||||
|
|
||||||
|
var enumerator = a.Merge(b).GetAsyncEnumerator();
|
||||||
|
(await enumerator.MoveNextAsync()).Should().Be(true);
|
||||||
|
enumerator.Current.Should().Be("A1");
|
||||||
|
|
||||||
|
await Assert.ThrowsAsync<UniTaskTestException>(async () => await enumerator.MoveNextAsync());
|
||||||
|
}
|
||||||
|
|
||||||
|
[Fact]
|
||||||
|
public async Task Cancel()
|
||||||
|
{
|
||||||
|
var cts = new CancellationTokenSource();
|
||||||
|
|
||||||
|
var a = UniTaskAsyncEnumerable.Create<string>(async (writer, _) =>
|
||||||
|
{
|
||||||
|
await writer.YieldAsync("A1");
|
||||||
|
});
|
||||||
|
|
||||||
|
var b = UniTaskAsyncEnumerable.Create<string>(async (writer, _) =>
|
||||||
|
{
|
||||||
|
await writer.YieldAsync("B1");
|
||||||
|
});
|
||||||
|
|
||||||
|
var enumerator = a.Merge(b).GetAsyncEnumerator(cts.Token);
|
||||||
|
(await enumerator.MoveNextAsync()).Should().Be(true);
|
||||||
|
enumerator.Current.Should().Be("A1");
|
||||||
|
|
||||||
|
cts.Cancel();
|
||||||
|
await Assert.ThrowsAsync<OperationCanceledException>(async () => await enumerator.MoveNextAsync());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
|
@ -39,7 +39,7 @@ namespace Cysharp.Threading.Tasks.Internal
|
||||||
}
|
}
|
||||||
|
|
||||||
[MethodImpl(MethodImplOptions.NoInlining)]
|
[MethodImpl(MethodImplOptions.NoInlining)]
|
||||||
public static void ThrowArgumentException<T>(string message)
|
public static void ThrowArgumentException(string message)
|
||||||
{
|
{
|
||||||
throw new ArgumentException(message);
|
throw new ArgumentException(message);
|
||||||
}
|
}
|
||||||
|
|
|
@ -0,0 +1,221 @@
|
||||||
|
using System;
|
||||||
|
using System.Collections.Generic;
|
||||||
|
using System.Linq;
|
||||||
|
using System.Threading;
|
||||||
|
using Cysharp.Threading.Tasks.Internal;
|
||||||
|
|
||||||
|
namespace Cysharp.Threading.Tasks.Linq
|
||||||
|
{
|
||||||
|
public static partial class UniTaskAsyncEnumerable
|
||||||
|
{
|
||||||
|
public static IUniTaskAsyncEnumerable<T> Merge<T>(this IUniTaskAsyncEnumerable<T> first, IUniTaskAsyncEnumerable<T> second)
|
||||||
|
{
|
||||||
|
Error.ThrowArgumentNullException(first, nameof(first));
|
||||||
|
Error.ThrowArgumentNullException(second, nameof(second));
|
||||||
|
|
||||||
|
return new Merge<T>(new [] { first, second });
|
||||||
|
}
|
||||||
|
|
||||||
|
public static IUniTaskAsyncEnumerable<T> Merge<T>(this IUniTaskAsyncEnumerable<T> first, IUniTaskAsyncEnumerable<T> second, IUniTaskAsyncEnumerable<T> third)
|
||||||
|
{
|
||||||
|
Error.ThrowArgumentNullException(first, nameof(first));
|
||||||
|
Error.ThrowArgumentNullException(second, nameof(second));
|
||||||
|
Error.ThrowArgumentNullException(third, nameof(third));
|
||||||
|
|
||||||
|
return new Merge<T>(new[] { first, second, third });
|
||||||
|
}
|
||||||
|
|
||||||
|
public static IUniTaskAsyncEnumerable<T> Merge<T>(this IEnumerable<IUniTaskAsyncEnumerable<T>> sources)
|
||||||
|
{
|
||||||
|
return new Merge<T>(sources.ToArray());
|
||||||
|
}
|
||||||
|
|
||||||
|
public static IUniTaskAsyncEnumerable<T> Merge<T>(params IUniTaskAsyncEnumerable<T>[] sources)
|
||||||
|
{
|
||||||
|
return new Merge<T>(sources);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
internal sealed class Merge<T> : IUniTaskAsyncEnumerable<T>
|
||||||
|
{
|
||||||
|
readonly IUniTaskAsyncEnumerable<T>[] sources;
|
||||||
|
|
||||||
|
public Merge(IUniTaskAsyncEnumerable<T>[] sources)
|
||||||
|
{
|
||||||
|
if (sources.Length <= 0)
|
||||||
|
{
|
||||||
|
Error.ThrowArgumentException("No source async enumerable to merge");
|
||||||
|
}
|
||||||
|
this.sources = sources;
|
||||||
|
}
|
||||||
|
|
||||||
|
public IUniTaskAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
|
||||||
|
=> new _Merge(sources, cancellationToken);
|
||||||
|
|
||||||
|
enum MergeSourceState
|
||||||
|
{
|
||||||
|
Pending,
|
||||||
|
Running,
|
||||||
|
Completed,
|
||||||
|
}
|
||||||
|
|
||||||
|
sealed class _Merge : MoveNextSource, IUniTaskAsyncEnumerator<T>
|
||||||
|
{
|
||||||
|
static readonly Action<object> GetResultAtAction = GetResultAt;
|
||||||
|
|
||||||
|
readonly int length;
|
||||||
|
readonly IUniTaskAsyncEnumerator<T>[] enumerators;
|
||||||
|
readonly MergeSourceState[] states;
|
||||||
|
readonly Queue<(T, Exception)> resultQueue = new Queue<(T, Exception)>();
|
||||||
|
readonly CancellationToken cancellationToken;
|
||||||
|
|
||||||
|
public T Current { get; private set; }
|
||||||
|
|
||||||
|
public _Merge(IUniTaskAsyncEnumerable<T>[] sources, CancellationToken cancellationToken)
|
||||||
|
{
|
||||||
|
this.cancellationToken = cancellationToken;
|
||||||
|
length = sources.Length;
|
||||||
|
states = ArrayPool<MergeSourceState>.Shared.Rent(length);
|
||||||
|
enumerators = ArrayPool<IUniTaskAsyncEnumerator<T>>.Shared.Rent(length);
|
||||||
|
for (var i = 0; i < length; i++)
|
||||||
|
{
|
||||||
|
enumerators[i] = sources[i].GetAsyncEnumerator(cancellationToken);
|
||||||
|
states[i] = MergeSourceState.Pending;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
public UniTask<bool> MoveNextAsync()
|
||||||
|
{
|
||||||
|
cancellationToken.ThrowIfCancellationRequested();
|
||||||
|
completionSource.Reset();
|
||||||
|
|
||||||
|
if (TryDequeue(out var queuedValue, out var queuedException))
|
||||||
|
{
|
||||||
|
if (queuedException != null)
|
||||||
|
{
|
||||||
|
completionSource.TrySetException(queuedException);
|
||||||
|
}
|
||||||
|
else
|
||||||
|
{
|
||||||
|
Current = queuedValue;
|
||||||
|
completionSource.TrySetResult(!IsCompletedAll());
|
||||||
|
}
|
||||||
|
return new UniTask<bool>(this, completionSource.Version);
|
||||||
|
}
|
||||||
|
|
||||||
|
for (var i = 0; i < length; i++)
|
||||||
|
{
|
||||||
|
lock (states)
|
||||||
|
{
|
||||||
|
if (states[i] != MergeSourceState.Pending)
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
states[i] = MergeSourceState.Running;
|
||||||
|
}
|
||||||
|
|
||||||
|
var awaiter = enumerators[i].MoveNextAsync().GetAwaiter();
|
||||||
|
if (awaiter.IsCompleted)
|
||||||
|
{
|
||||||
|
GetResultAt(i, awaiter);
|
||||||
|
}
|
||||||
|
else
|
||||||
|
{
|
||||||
|
awaiter.SourceOnCompleted(GetResultAtAction, StateTuple.Create(this, i, awaiter));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return new UniTask<bool>(this, completionSource.Version);
|
||||||
|
}
|
||||||
|
|
||||||
|
public async UniTask DisposeAsync()
|
||||||
|
{
|
||||||
|
for (var i = 0; i < length; i++)
|
||||||
|
{
|
||||||
|
await enumerators[i].DisposeAsync();
|
||||||
|
}
|
||||||
|
|
||||||
|
ArrayPool<MergeSourceState>.Shared.Return(states, true);
|
||||||
|
ArrayPool<IUniTaskAsyncEnumerator<T>>.Shared.Return(enumerators, true);
|
||||||
|
}
|
||||||
|
|
||||||
|
static void GetResultAt(object state)
|
||||||
|
{
|
||||||
|
var tuple = (StateTuple<_Merge, int, UniTask<bool>.Awaiter>)state;
|
||||||
|
tuple.Item1.GetResultAt(tuple.Item2, tuple.Item3);
|
||||||
|
}
|
||||||
|
|
||||||
|
void GetResultAt(int index, UniTask<bool>.Awaiter awaiter)
|
||||||
|
{
|
||||||
|
bool hasNext;
|
||||||
|
try
|
||||||
|
{
|
||||||
|
hasNext = awaiter.GetResult();
|
||||||
|
lock (states)
|
||||||
|
{
|
||||||
|
states[index] = hasNext ? MergeSourceState.Pending : MergeSourceState.Completed;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
catch (Exception ex)
|
||||||
|
{
|
||||||
|
if (!completionSource.TrySetException(ex))
|
||||||
|
{
|
||||||
|
lock (resultQueue)
|
||||||
|
{
|
||||||
|
resultQueue.Enqueue((default, ex));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
var completed = IsCompletedAll();
|
||||||
|
if (hasNext || completed)
|
||||||
|
{
|
||||||
|
if (completionSource.GetStatus(completionSource.Version).IsCompleted())
|
||||||
|
{
|
||||||
|
lock (resultQueue)
|
||||||
|
{
|
||||||
|
resultQueue.Enqueue((enumerators[index].Current, null));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
else
|
||||||
|
{
|
||||||
|
Current = enumerators[index].Current;
|
||||||
|
completionSource.TrySetResult(!completed);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
bool TryDequeue(out T value, out Exception ex)
|
||||||
|
{
|
||||||
|
lock (resultQueue)
|
||||||
|
{
|
||||||
|
if (resultQueue.Count > 0)
|
||||||
|
{
|
||||||
|
var result = resultQueue.Dequeue();
|
||||||
|
value = result.Item1;
|
||||||
|
ex = result.Item2;
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
value = default;
|
||||||
|
ex = default;
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
bool IsCompletedAll()
|
||||||
|
{
|
||||||
|
lock (states)
|
||||||
|
{
|
||||||
|
for (var i = 0; i < length; i++)
|
||||||
|
{
|
||||||
|
if (states[i] != MergeSourceState.Completed)
|
||||||
|
{
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
|
@ -0,0 +1,3 @@
|
||||||
|
fileFormatVersion: 2
|
||||||
|
guid: ca56812f160c45d0bacb4339819edf1a
|
||||||
|
timeCreated: 1694133666
|
Loading…
Reference in New Issue