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

93
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;
if (queuedException != null)
{
completionSource.TrySetException(queuedException);
}
else
{
Current = queuedValue;
completionSource.TrySetResult(!IsCompletedAll());
}
return new UniTask<bool>(this, completionSource.Version);
} }
var resultValue = value.Item1;
var exception = value.Item2;
var hasNext = value.Item3;
if (exception != null)
{
completionSource.TrySetException(exception);
}
else
{
Current = resultValue;
completionSource.TrySetResult(hasNext);
}
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,48 +162,67 @@ 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
{
lock (states)
{ {
queuedResult.Enqueue((default, ex)); 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 (Interlocked.CompareExchange(ref moveNextCompleted, 1, 0) == 0)
{ {
if (completionSource.GetStatus(completionSource.Version).IsCompleted()) Current = enumerators[index].Current;
completionSource.TrySetResult(!completedAll);
}
else
{
lock (states)
{ {
queuedResult.Enqueue((enumerators[index].Current, null)); queuedResult.Enqueue((enumerators[index].Current, null, !completedAll));
}
else
{
Current = enumerators[index].Current;
completionSource.TrySetResult(!completedAll);
} }
} }
} }
} }
bool HasQueuedResult()
{
lock (states)
{
return queuedResult.Count > 0;
}
}
bool IsCompletedAll() bool IsCompletedAll()
{ {
for (var i = 0; i < length; i++) lock (states)
{ {
if (states[i] != MergeSourceState.Completed) for (var i = 0; i < length; i++)
{ {
return false; if (states[i] != MergeSourceState.Completed)
{
return false;
}
} }
} }
return true; return true;