Reduce the lock

master
hadashiA 2023-09-13 19:12:11 +09:00
parent 937d3adf66
commit 3ba64412f8
1 changed files with 58 additions and 35 deletions

View File

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