Merge pull request #19 from fretje/unsubscribeweak

Fix bug where RemoveWeak would sometimes remove too much

Author: KuraiAndras

CHANGELOG.md

@@ -6,7 +6,8 @@
 - Pipelines now fully written in github actions
 - Updated dependencies
 - Added documentation comments to `ICourier`, `CourierInjector` and `CourierOptions`
-- Use Invoke iso DynamicInvoke for better performance 
+- Use Invoke iso DynamicInvoke for better performance
+- Fix bug where RemoveWeak would sometimes remove too many handlers
 
 # 6.0.0
 - Update MediatR to version 12

MediatR.Courier.Tests/WeakReferenceTests.cs

@@ -36,6 +36,31 @@ public async Task UnSubscribedAsyncActionNotInvoked()
         Assert.Equal(1, handler.ReceivedMessageCount);
     }
 
+    [Fact]
+    public async Task UnSubscribedAsyncActionNotInvokedButOthersAre()
+    {
+        var mediatRCourier = new MediatRCourier(new());
+
+        var handler = new Handler(new(), new());
+        var handler2 = new Handler(new(), new());
+
+        mediatRCourier.SubscribeWeak<TestNotification>(handler.NotificationAction);
+        mediatRCourier.SubscribeWeak<TestNotification>(handler.NotificationAction2);
+        mediatRCourier.SubscribeWeak<TestNotification>(handler2.NotificationAction);
+        mediatRCourier.SubscribeWeak<TestNotification>(handler2.NotificationAction2);
+
+        await mediatRCourier.Handle(new TestNotification(), CancellationToken.None);
+
+        mediatRCourier.UnSubscribe<TestNotification>(handler.NotificationAction);
+
+        await mediatRCourier.Handle(new TestNotification(), CancellationToken.None);
+
+        Assert.Equal(1, handler.ReceivedMessageCount);
+        Assert.Equal(2, handler.ReceivedMessageCount2);
+        Assert.Equal(2, handler2.ReceivedMessageCount);
+        Assert.Equal(2, handler2.ReceivedMessageCount2);
+    }
+
     [Fact]
     [System.Diagnostics.CodeAnalysis.SuppressMessage("Critical Code Smell", "S1215:\"GC.Collect\" should not be called", Justification = "We want to test behavior around GC collection")]
     [System.Diagnostics.CodeAnalysis.SuppressMessage("CodeQuality", "IDE0079:Remove unnecessary suppression", Justification = "False positive")]
@@ -71,8 +96,13 @@ private sealed class Counter
     private sealed class Handler
     {
         private readonly Counter _counter;
+        private readonly Counter? _counter2;
 
-        public Handler(Counter counter) => _counter = counter;
+        public Handler(Counter counter, Counter? counter2 = null)
+        {
+            _counter = counter;
+            _counter2 = counter2;
+        }
 
 #pragma warning disable S1172 // Unused method parameters should be removed
         public async Task NotificationAction(TestNotification _, CancellationToken cancellationToken)
@@ -83,6 +113,18 @@ public async Task NotificationAction(TestNotification _, CancellationToken cance
             await Task.Delay(TimeSpan.FromMilliseconds(1), cancellationToken);
         }
 
+#pragma warning disable S1172 // Unused method parameters should be removed
+        public async Task NotificationAction2(TestNotification _, CancellationToken cancellationToken)
+#pragma warning restore S1172 // Unused method parameters should be removed
+        {
+            if (_counter2 is null) return;
+
+            _counter2.ReceivedMessageCount++;
+
+            await Task.Delay(TimeSpan.FromMilliseconds(1), cancellationToken).ConfigureAwait(false);
+        }
+
         public int ReceivedMessageCount => _counter.ReceivedMessageCount;
+        public int? ReceivedMessageCount2 => _counter2?.ReceivedMessageCount;
     }
 }

MediatR.Courier/MediatRCourier.cs

@@ -158,7 +158,9 @@ private void RemoveWeak(Delegate handler, Type notificationType)
 
         var remainingSubscribers = new ConcurrentBag<(WeakReference<object> target, MethodInfo methodInfo, bool needsToken)>
         (
-            subscribers.Where(subscriber => !subscriber.target.TryGetTarget(out var weakHandler) && weakHandler != handler.Target)
+            subscribers.Where(subscriber =>
+                (subscriber.target.TryGetTarget(out var weakHandler) && weakHandler != handler.Target)
+                || subscriber.methodInfo != handler.Method)
         );
 
         _weakActions.TryRemove(notificationType, out _);