Remove message source when link detached

Author: Havret

src/ArtemisNetClient.Testing/SharedMessageSource.cs

@@ -23,13 +23,25 @@ public void AddMessageSource(MessageSource messageSource)
         _messageSources.Add(messageSource);
     }
 
+    public void RemoveMessageSource(MessageSource messageSource)
+    {
+        _messageSources.Remove(messageSource);
+    }
+
+    public int MessageSourcesCount => _messageSources.Count;
+    
     public void Enqueue(Message message)
     {
         if (Info.FilterExpression != null && FilterExpressionMatches(message) == false)
         {
             return;
         }
 
+        if (_messageSources.Count == 0)
+        {
+            return;
+        }
+
         if (message.GroupId is { Length: > 0 } groupId)
         {
             var cursor = ArtemisBucketHelper.GetBucket(groupId, _messageSources.Count);

src/ArtemisNetClient.Testing/TestKit.cs

@@ -15,7 +15,7 @@ public TestKit(Endpoint endpoint)
         _host = new ContainerHost(endpoint.Address);
         _host.AddressResolver = AddressResolver;
         _host.Listeners[0].HandlerFactory = _ => new Handler();
-        _host.RegisterLinkProcessor(new TestLinkProcessor(OnMessage, OnMessageSource));
+        _host.RegisterLinkProcessor(new TestLinkProcessor(OnMessage, OnMessageSource, OnMessageSourceClosed));
         _host.Open();
     }
 
@@ -91,6 +91,29 @@ private void OnMessageSource(MessageSourceInfo info, MessageSource messageSource
         }
     }
 
+    private void OnMessageSourceClosed(MessageSourceInfo info, MessageSource messageSource)
+    {
+        if (!_messageSources.TryGetValue(info.Address, out var messageSources))
+        {
+            return;
+        }
+
+        lock (messageSources)
+        {
+            var existingMessageSource = messageSources.FirstOrDefault(x => x.Info == info);
+            if (existingMessageSource == null)
+            {
+                return;
+            }
+            
+            existingMessageSource.RemoveMessageSource(messageSource);
+            if (messageSources.Count == 0)
+            {
+                messageSources.Remove(existingMessageSource);
+            }
+        }
+    }
+
     public ISubscription Subscribe(string address)
     {
         var subscription = new Subscription();

src/ArtemisNetClient.Testing/TestLinkProcessor.cs

@@ -12,11 +12,15 @@ internal class TestLinkProcessor : ILinkProcessor
 {
     private readonly Action<Message> _onMessage;
     private readonly Action<MessageSourceInfo, MessageSource> _onMessageSource;
+    private readonly Action<MessageSourceInfo, MessageSource> _onMessageSourceClosed;
 
-    public TestLinkProcessor(Action<Message> onMessage, Action<MessageSourceInfo, MessageSource> onMessageSource)
+    public TestLinkProcessor(Action<Message> onMessage,
+        Action<MessageSourceInfo, MessageSource> onMessageSource,
+        Action<MessageSourceInfo, MessageSource> onMessageSourceClosed)
     {
         _onMessage = onMessage;
         _onMessageSource = onMessageSource;
+        _onMessageSourceClosed = onMessageSourceClosed;
     }
 
     public void Process(AttachContext attachContext)
@@ -35,6 +39,10 @@ public void Process(AttachContext attachContext)
 
             // override OnDispose so it won't throw NRE when message is null
             attachContext.Link.SetOnDispose((_, _, _, _) => { });
+            attachContext.Link.AddSafeClosed((_, _) =>
+            {
+                _onMessageSourceClosed(messageSourceInfo, messageSource);
+            });
 
             attachContext.Link.CompleteAttach(attachContext.Attach, null);
         }