PHXH
/
UniTask
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add UniTaskAsyncEnumerable.Merge

master
hadashiA 2023-09-08 13:04:10 +09:00
parent c170af5642
commit f303d9d7e8
4 changed files with 362 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

137
src/UniTask.NetCoreTests/Linq/Merge.cs Normal file
View File

@ -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());
}
}
}

2
src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/Error.cs
View File

@ -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);
} }

221
src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Merge.cs Normal file
View File

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

3
src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Merge.cs.meta Normal file
View File

@ -0,0 +1,3 @@
fileFormatVersion: 2
guid: ca56812f160c45d0bacb4339819edf1a
timeCreated: 1694133666