diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Merge.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Merge.cs index cfbcff5d..5bc7649c 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Merge.cs +++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/Merge.cs @@ -69,8 +69,6 @@ sealed class _Merge : MoveNextSource, IUniTaskAsyncEnumerator readonly Queue<(T, Exception)> resultQueue = new Queue<(T, Exception)>(); readonly CancellationToken cancellationToken; - bool allCompleted; - public T Current { get; private set; } public _Merge(IUniTaskAsyncEnumerable[] sources, CancellationToken cancellationToken) @@ -100,7 +98,7 @@ public UniTask MoveNextAsync() else { Current = queuedValue; - completionSource.TrySetResult(Volatile.Read(ref allCompleted)); + completionSource.TrySetResult(!IsCompletedAll()); } return new UniTask(this, completionSource.Version); } @@ -155,39 +153,22 @@ void GetResultAt(int index, UniTask.Awaiter awaiter) lock (states) { states[index] = hasNext ? MergeSourceState.Pending : MergeSourceState.Completed; - - // Check to complete - if (!hasNext) - { - allCompleted = true; - for (var i = 0; i < length; i++) - { - if (states[i] != MergeSourceState.Completed) - { - allCompleted = false; - break; - } - } - } } } catch (Exception ex) { - if (completionSource.GetStatus(completionSource.Version).IsCompleted()) + if (!completionSource.TrySetException(ex)) { lock (resultQueue) { resultQueue.Enqueue((default, ex)); } } - else - { - completionSource.TrySetException(ex); - } return; } - if (hasNext || Volatile.Read(ref allCompleted)) + var completed = IsCompletedAll(); + if (hasNext || completed) { if (completionSource.GetStatus(completionSource.Version).IsCompleted()) { @@ -199,7 +180,7 @@ void GetResultAt(int index, UniTask.Awaiter awaiter) else { Current = enumerators[index].Current; - completionSource.TrySetResult(!Volatile.Read(ref allCompleted)); + completionSource.TrySetResult(!completed); } } } @@ -220,6 +201,21 @@ bool TryDequeue(out T value, out Exception ex) ex = default; return false; } - } + + bool IsCompletedAll() + { + lock (states) + { + for (var i = 0; i < length; i++) + { + if (states[i] != MergeSourceState.Completed) + { + return false; + } + } + return true; + } + } + } } } \ No newline at end of file