fix: subscribing to a message channel while unsubscribing is pending [MTT-3756] (#675)

* adding removal of pending sub/unsub when unsubscribing/subscribing

* Adding tests

* adding specificity to auto-injected GameObjects in MainMenu scene

* Moving subscribing/unsubscribing to LifeStateChangedEventMessages in ServerBossRoomState to OnNetworkSpawn/Despawn

Author: LPLafontaineB

Assets/Scenes/MainMenu.unity

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f9162f4f410f9bf87f75cb0b933e0902b542ef09f9b1c1b7a2a93ce01ec11bf6
-size 67373
+oid sha256:3ca89d724f7bc5f0ef930a26462369191017995509ae4cdaefafe43010e62491
+size 68542

Assets/Scripts/Gameplay/GameState/ServerBossRoomState.cs

@@ -71,12 +71,16 @@ void OnNetworkSpawn()
                 return;
             }
 
+            m_Subscription = m_LifeStateChangedEventMessageSubscriber.Subscribe(OnLifeStateChangedEventMessage);
+
             NetworkManager.Singleton.SceneManager.OnLoadComplete += OnServerLoadComplete;
             NetworkManager.Singleton.SceneManager.OnUnloadComplete += OnServerUnloadComplete;
         }
 
         void OnNetworkDespawn()
         {
+            m_Subscription?.Dispose();
+
             NetworkManager.Singleton.SceneManager.OnLoadComplete -= OnServerLoadComplete;
             NetworkManager.Singleton.SceneManager.OnUnloadComplete -= OnServerUnloadComplete;
         }
@@ -96,7 +100,6 @@ void OnServerLoadComplete(ulong clientId, string sceneName, LoadSceneMode loadSc
             NetworkManager.Singleton.SceneManager.OnSynchronizeComplete += OnSynchronizeComplete;
 
             SessionManager<SessionPlayerData>.Instance.OnSessionStarted();
-            m_Subscription = m_LifeStateChangedEventMessageSubscriber.Subscribe(OnLifeStateChangedEventMessage);
         }
 
         void OnServerUnloadComplete(ulong clientId, string sceneName)
@@ -109,7 +112,6 @@ void OnServerUnloadComplete(ulong clientId, string sceneName)
             NetworkManager.Singleton.OnClientDisconnectCallback -= OnClientDisconnect;
             NetworkManager.Singleton.SceneManager.OnLoadEventCompleted -= OnLoadEventCompleted;
             NetworkManager.Singleton.SceneManager.OnSynchronizeComplete -= OnSynchronizeComplete;
-            m_Subscription?.Dispose();
         }
 
         protected override void OnDestroy()

Assets/Scripts/Infrastructure/PubSub/MessageChannel.cs

@@ -8,11 +8,10 @@ public class MessageChannel<T> : IMessageChannel<T>
     {
         readonly List<Action<T>> m_MessageHandlers = new List<Action<T>>();
 
-        /// <summary>
-        /// This queue of actions that would either add or remove subscriber is used to prevent problems from immediate modification
-        /// of the list of subscribers. It could happen if one decides to unsubscribe in a message handler etc.
-        /// </summary>
-        readonly Queue<Action> m_PendingHandlers = new Queue<Action>();
+        /// This dictionary of handlers to be either added or removed is used to prevent problems from immediate
+        /// modification of the list of subscribers. It could happen if one decides to unsubscribe in a message handler
+        /// etc.A true value means this handler should be added, and a false one means it should be removed
+        readonly Dictionary<Action<T>, bool> m_PendingHandlers = new Dictionary<Action<T>, bool>();
 
         public bool IsDisposed { get; private set; } = false;
 
@@ -28,10 +27,18 @@ public virtual void Dispose()
 
         public virtual void Publish(T message)
         {
-            while (m_PendingHandlers.Count > 0)
+            foreach (var handler in m_PendingHandlers.Keys)
             {
-                m_PendingHandlers.Dequeue()?.Invoke();
+                if (m_PendingHandlers[handler])
+                {
+                    m_MessageHandlers.Add(handler);
+                }
+                else
+                {
+                    m_MessageHandlers.Remove(handler);
+                }
             }
+            m_PendingHandlers.Clear();
 
             foreach (var messageHandler in m_MessageHandlers)
             {
@@ -41,30 +48,46 @@ public virtual void Publish(T message)
 
         public virtual IDisposable Subscribe(Action<T> handler)
         {
-            Assert.IsTrue(!m_MessageHandlers.Contains(handler), "Attempting to subscribe with the same handler more than once");
-            // so we don't modify the handler list while iterating on it (which could happen if the handler unsubscribes). With this, we'd unsubscribe on next publish.
-            m_PendingHandlers.Enqueue(() => { DoSubscribe(handler); });
-
-            var subscription = new DisposableSubscription<T>(this, handler);
-            return subscription;
+            Assert.IsTrue(!IsSubscribed(handler), "Attempting to subscribe with the same handler more than once");
 
-            void DoSubscribe(Action<T> h)
+            if (m_PendingHandlers.ContainsKey(handler))
             {
-                if (h != null && !m_MessageHandlers.Contains(h))
+                if (!m_PendingHandlers[handler])
                 {
-                    m_MessageHandlers.Add(h);
+                    m_PendingHandlers.Remove(handler);
                 }
             }
+            else
+            {
+                m_PendingHandlers[handler] = true;
+            }
+
+            var subscription = new DisposableSubscription<T>(this, handler);
+            return subscription;
         }
 
         public void Unsubscribe(Action<T> handler)
         {
-            m_PendingHandlers.Enqueue(() => { DoUnsubscribe(handler); });
+            Assert.IsTrue(IsSubscribed(handler), "Attempting to unsubscribe with a handler that is not subscribed");
 
-            void DoUnsubscribe(Action<T> h)
+            if (m_PendingHandlers.ContainsKey(handler))
             {
-                m_MessageHandlers.Remove(h);
+                if (m_PendingHandlers[handler])
+                {
+                    m_PendingHandlers.Remove(handler);
+                }
             }
+            else
+            {
+                m_PendingHandlers[handler] = false;
+            }
+        }
+
+        bool IsSubscribed(Action<T> handler)
+        {
+            var isPendingRemoval = m_PendingHandlers.ContainsKey(handler) && !m_PendingHandlers[handler];
+            var isPendingAdding = m_PendingHandlers.ContainsKey(handler) && m_PendingHandlers[handler];
+            return m_MessageHandlers.Contains(handler) && !isPendingRemoval || isPendingAdding;
         }
     }
 }

Assets/Tests/Runtime/MessageChannelTest.cs

@@ -0,0 +1,88 @@
+using System;
+using NUnit.Framework;
+using Unity.Multiplayer.Samples.BossRoom.Shared.Infrastructure;
+using UnityEngine;
+
+namespace Unity.Multiplayer.Samples.BossRoom.Tests.Runtime
+{
+    public class MessageChannelTests
+    {
+        struct EmptyMessage { }
+
+        int m_NbSubscribers = 2;
+        int m_NbMessagesReceived;
+
+        IDisposable SubscribeToChannel(MessageChannel<EmptyMessage> channel)
+        {
+            var subscriptions = new DisposableGroup();
+            subscriptions.Add(channel.Subscribe(Subscription1));
+            subscriptions.Add(channel.Subscribe(Subscription2));
+
+            return subscriptions;
+        }
+
+        void Subscription1(EmptyMessage message)
+        {
+            m_NbMessagesReceived++;
+        }
+
+        void Subscription2(EmptyMessage message)
+        {
+            m_NbMessagesReceived++;
+        }
+
+
+        [SetUp]
+        public void Setup()
+        {
+            m_NbMessagesReceived = 0;
+        }
+
+        [Test]
+        public void MessagePublishedIsReceivedByAllSubscribers()
+        {
+            var messageChannel = new MessageChannel<EmptyMessage>();
+            var subscriptions = SubscribeToChannel(messageChannel);
+
+            messageChannel.Publish(new EmptyMessage());
+            Assert.AreEqual(m_NbSubscribers, m_NbMessagesReceived);
+            subscriptions.Dispose();
+        }
+
+        [Test]
+        public void MessagePublishedIsNotReceivedByAllSubscribersAfterUnsubscribing()
+        {
+            var messageChannel = new MessageChannel<EmptyMessage>();
+            var subscriptions = SubscribeToChannel(messageChannel);
+
+            messageChannel.Publish(new EmptyMessage());
+            Assert.AreEqual(m_NbSubscribers, m_NbMessagesReceived);
+
+            m_NbMessagesReceived = 0;
+
+            subscriptions.Dispose();
+
+            messageChannel.Publish(new EmptyMessage());
+            Assert.AreEqual(0, m_NbMessagesReceived);
+        }
+
+        [Test]
+        public void MessagePublishedIsReceivedByAllSubscribersAfterResubscribing()
+        {
+            var messageChannel = new MessageChannel<EmptyMessage>();
+            var subscriptions = SubscribeToChannel(messageChannel);
+
+            messageChannel.Publish(new EmptyMessage());
+            Assert.AreEqual(m_NbSubscribers, m_NbMessagesReceived);
+
+            m_NbMessagesReceived = 0;
+
+            subscriptions.Dispose();
+            subscriptions = SubscribeToChannel(messageChannel);
+
+            messageChannel.Publish(new EmptyMessage());
+            Assert.AreEqual(m_NbSubscribers, m_NbMessagesReceived);
+            subscriptions.Dispose();
+        }
+    }
+}

Assets/Tests/Runtime/MessageChannelTest.cs.meta

@@ -18,7 +18,7 @@ Additional documentation and release notes are available at [Multiplayer Documen
 ### Removed
 *
 ### Fixed
-*
+* Subscribing to a message channel while unsubscribing is pending (#675)
 
 ## [v1.3.0-pre] - 2022-06-23