fix mergeEnumerable in case of reentrancy on SubscribeAsync

Author: fedeAlterio

src/R3Async.Tests/Operators/MergeEnumerableTest.cs

@@ -51,6 +51,46 @@ public async Task MergeEnumerable_Basic()
         results.OrderBy(x => x).ShouldBe(new[] { 1, 2, 3 });
     }
 
+    [Fact]
+    public async Task MergeEnumerableSync_Basic()
+    {
+        var tcs1 = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var tcs2 = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        var obs1 = AsyncObservable.Create<int>(async (observer, token) =>
+        {
+            await observer.OnNextAsync(1, token);
+            await observer.OnNextAsync(2, token);
+            await observer.OnCompletedAsync(Result.Success);
+            tcs1.TrySetResult();
+            return AsyncDisposable.Empty;
+        });
+
+        var obs2 = AsyncObservable.Create<int>(async (observer, token) =>
+        {
+            await observer.OnNextAsync(3, token);
+            await observer.OnCompletedAsync(Result.Success);
+            tcs2.TrySetResult();
+            return AsyncDisposable.Empty;
+        });
+
+        var merge = new[] { obs1, obs2 }.Merge();
+        var results = new List<int>();
+        var completedTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        await using var subscription = await merge.SubscribeAsync(
+            async (x, token) => results.Add(x),
+            async (ex, token) => { },
+            async result => completedTcs.TrySetResult(result.IsSuccess),
+            CancellationToken.None);
+
+        await tcs1.Task;
+        await tcs2.Task;
+        await completedTcs.Task;
+
+        results.OrderBy(x => x).ShouldBe(new[] { 1, 2, 3 });
+    }
+
     [Fact]
     public async Task MergeEnumerable_Empty()
     {

src/R3Async/Operators/Merge.cs

@@ -242,7 +242,8 @@ sealed class MergeEnumerableSubscription : IAsyncDisposable
             readonly TaskCompletionSource<bool> _subscriptionFinished = new(TaskCreationOptions.RunContinuationsAsynchronously);
             readonly AsyncLocal<bool> _reentrant = new();
             int _active;
-            int _disposed;
+            bool _enumerationCompleted;
+            bool _disposed;
             readonly AsyncObserver<T> _observer;
 
             public MergeEnumerableSubscription(AsyncObserver<T> observer, IEnumerable<AsyncObservable<T>> sources)
@@ -261,7 +262,11 @@ public async void StartAsync()
                     {
                         foreach (var src in _sources)
                         {
-                            Interlocked.Increment(ref _active);
+                            _disposedCancellationToken.ThrowIfCancellationRequested();
+                            lock (_onSomethingGate)
+                            {
+                                _active++;
+                            }
 
                             var innerObserver = new InnerAsyncObserver(this);
                             await _innerDisposables.AddAsync(innerObserver);
@@ -280,7 +285,14 @@ public async void StartAsync()
                             }
                         }
 
-                        if (Volatile.Read(ref _active) == 0)
+                        bool shouldComplete;
+                        lock (_onSomethingGate)
+                        {
+                            _enumerationCompleted = true;
+                            shouldComplete = _active == 0;
+                        }
+
+                        if (shouldComplete)
                         {
                             await CompleteAsync(Result.Success);
                         }
@@ -305,7 +317,7 @@ async ValueTask OnNextAsync(T value, CancellationToken token)
                 using var linked = CancellationTokenSource.CreateLinkedTokenSource(_disposedCancellationToken, token);
                 using (await _onSomethingGate.LockAsync())
                 {
-                    if (_disposed == 1) return;
+                    if (_disposed) return;
                     await _observer.OnNextAsync(value, linked.Token);
                 }
             }
@@ -315,7 +327,7 @@ async ValueTask OnErrorResumeAsync(Exception ex, CancellationToken token)
                 using var linked = CancellationTokenSource.CreateLinkedTokenSource(_disposedCancellationToken, token);
                 using (await _onSomethingGate.LockAsync())
                 {
-                    if (_disposed == 1) return;
+                    if (_disposed) return;
                     await _observer.OnErrorResumeAsync(ex, linked.Token);
                 }
             }
@@ -327,21 +339,28 @@ ValueTask OnCompletedAsync(Result result)
                     return CompleteAsync(result);
                 }
 
-                if (Interlocked.Decrement(ref _active) == 0)
+                bool shouldComplete;
+                lock (_onSomethingGate)
                 {
-                    return CompleteAsync(Result.Success);
+                    _active--;
+                    shouldComplete = _active == 0 && _enumerationCompleted;
                 }
 
-                return default;
+                return shouldComplete ? CompleteAsync(Result.Success) : default;
             }
 
             async ValueTask CompleteAsync(Result? result)
             {
-                if (Interlocked.Exchange(ref _disposed, 1) == 1)
+                using (await _onSomethingGate.LockAsync())
                 {
-                    if (result?.Exception is not null and var ex)
-                        UnhandledExceptionHandler.OnUnhandledException(ex);
-                    return;
+                    if (_disposed)
+                    {
+                        if (result?.Exception is not null and var ex)
+                            UnhandledExceptionHandler.OnUnhandledException(ex);
+                        return;
+                    }
+
+                    _disposed = true;
                 }
 
                 _cts.Cancel();

src/R3Async/R3Async.csproj

@@ -11,7 +11,7 @@
     <Description>Async version of R3 (Variant of async rx.net)</Description>
     <PackageReadmeFile>README.md</PackageReadmeFile>
     <RepositoryUrl>https://github.com/fedeAlterio/R3Async</RepositoryUrl>
-    <Version>0.0.11</Version>
+    <Version>0.0.12</Version>
     <PackageLicenseFile>LICENSE</PackageLicenseFile>
     <PackageIcon>icon.png</PackageIcon>
     <PackageProjectUrl>https://github.com/fedeAlterio/R3Async</PackageProjectUrl>