pengtaihao
/
UniTask
mirror of https://github.com/Cysharp/UniTask
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Reduce the lock

pull/498/head
hadashiA 2023-09-13 19:12:11 +09:00
parent 937d3adf66
commit 3ba64412f8
1 changed files with 58 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Merge.cs
View File

@ -66,9 +66,11 @@ namespace Cysharp.Threading.Tasks.Linq
readonly int length; readonly int length;
readonly IUniTaskAsyncEnumerator<T>[] enumerators; readonly IUniTaskAsyncEnumerator<T>[] enumerators;
readonly MergeSourceState[] states; readonly MergeSourceState[] states;
readonly Queue<(T, Exception)> queuedResult = new Queue<(T, Exception)>(); readonly Queue<(T, Exception, bool)> queuedResult = new Queue<(T, Exception, bool)>();
readonly CancellationToken cancellationToken; readonly CancellationToken cancellationToken;
int moveNextCompleted;
public T Current { get; private set; } public T Current { get; private set; }
public _Merge(IUniTaskAsyncEnumerable<T>[] sources, CancellationToken cancellationToken) public _Merge(IUniTaskAsyncEnumerable<T>[] sources, CancellationToken cancellationToken)
@ -88,33 +90,35 @@ namespace Cysharp.Threading.Tasks.Linq
{ {
cancellationToken.ThrowIfCancellationRequested(); cancellationToken.ThrowIfCancellationRequested();
completionSource.Reset(); completionSource.Reset();
Interlocked.Exchange(ref moveNextCompleted, 0);
lock (queuedResult) if (HasQueuedResult() && Interlocked.CompareExchange(ref moveNextCompleted, 1, 0) == 0)
{ {
if (queuedResult.Count > 0) (T, Exception, bool) value;
lock (states)
{ {
var result = queuedResult.Dequeue(); value = queuedResult.Dequeue();
var queuedValue = result.Item1; }
var queuedException = result.Item2; var resultValue = value.Item1;
var exception = value.Item2;
if (queuedException != null) var hasNext = value.Item3;
if (exception != null)
{ {
completionSource.TrySetException(queuedException); completionSource.TrySetException(exception);
} }
else else
{ {
Current = queuedValue; Current = resultValue;
completionSource.TrySetResult(!IsCompletedAll()); completionSource.TrySetResult(hasNext);
} }
return new UniTask<bool>(this, completionSource.Version); return new UniTask<bool>(this, completionSource.Version);
} }
}
for (var i = 0; i < length; i++) for (var i = 0; i < length; i++)
{ {
lock (queuedResult) lock (states)
{ {
if (states[i] == (int)MergeSourceState.Pending) if (states[i] == MergeSourceState.Pending)
{ {
states[i] = MergeSourceState.Running; states[i] = MergeSourceState.Running;
} }
@ -158,42 +162,60 @@ namespace Cysharp.Threading.Tasks.Linq
void GetResultAt(int index, UniTask<bool>.Awaiter awaiter) void GetResultAt(int index, UniTask<bool>.Awaiter awaiter)
{ {
bool hasNext; bool hasNext;
bool completedAll;
try try
{ {
hasNext = awaiter.GetResult(); hasNext = awaiter.GetResult();
} }
catch (Exception ex) catch (Exception ex)
{ {
if (!completionSource.TrySetException(ex)) if (Interlocked.CompareExchange(ref moveNextCompleted, 1, 0) == 0)
{ {
lock (queuedResult) completionSource.TrySetException(ex);
}
else
{ {
queuedResult.Enqueue((default, ex)); lock (states)
{
queuedResult.Enqueue((default, ex, default));
} }
} }
return; return;
} }
lock (queuedResult) lock (states)
{ {
states[index] = hasNext ? MergeSourceState.Pending : MergeSourceState.Completed; states[index] = hasNext ? MergeSourceState.Pending : MergeSourceState.Completed;
var completedAll = !hasNext && IsCompletedAll(); completedAll = !hasNext && IsCompletedAll();
}
if (hasNext || completedAll) if (hasNext || completedAll)
{ {
if (completionSource.GetStatus(completionSource.Version).IsCompleted()) if (Interlocked.CompareExchange(ref moveNextCompleted, 1, 0) == 0)
{
queuedResult.Enqueue((enumerators[index].Current, null));
}
else
{ {
Current = enumerators[index].Current; Current = enumerators[index].Current;
completionSource.TrySetResult(!completedAll); completionSource.TrySetResult(!completedAll);
} }
else
{
lock (states)
{
queuedResult.Enqueue((enumerators[index].Current, null, !completedAll));
}
} }
} }
} }
bool HasQueuedResult()
{
lock (states)
{
return queuedResult.Count > 0;
}
}
bool IsCompletedAll() bool IsCompletedAll()
{
lock (states)
{ {
for (var i = 0; i < length; i++) for (var i = 0; i < length; i++)
{ {
@ -202,6 +224,7 @@ namespace Cysharp.Threading.Tasks.Linq
return false; return false;
} }
} }
}
return true; return true;
} }
} }