Several callbacks can be added to a same message type and disposed (c…

…olyseus#197)
sticmac · Jul 8, 2022 · abaf471 · abaf471
1 parent 7020109
commit abaf471
Showing 1 changed file with 51 additions and 25 deletions.
diff --git a/Assets/Colyseus/Runtime/Scripts/Room/ColyseusRoom.cs b/Assets/Colyseus/Runtime/Scripts/Room/ColyseusRoom.cs
@@ -282,18 +282,48 @@ public async Task Send(string type, object message)
             await colyseusConnection.Send(bytes);
         }
 
+        /// <summary>
+        /// Disposable class to return with OnMessage function, so that a message callback can be removed.
+        /// </summary>
+        /// <typeparam name="MessageType">The type of object the message responds with</typeparam>
+        private class MessageRemover<MessageType> : IDisposable
+        {
+            private Action<MessageType> _targetHandler;
+            private ColyseusMessageHandler<MessageType> _registeredHandler;
+
+            public MessageRemover(Action<MessageType> targetHandler, ColyseusMessageHandler<MessageType> registeredHandler)
+            {
+                _targetHandler = targetHandler;
+                _registeredHandler = registeredHandler;
+            }
+
+            public void Dispose()
+            {
+                _registeredHandler.Action -= _targetHandler;
+            }
+        }
+
         /// <summary>
         ///     Method to add new message handlers to the room
         /// </summary>
         /// <param name="type">The type of message received</param>
         /// <param name="handler"></param>
         /// <typeparam name="MessageType">The type of object this message should respond with</typeparam>
-        public void OnMessage<MessageType>(string type, Action<MessageType> handler)
+        /// <returns>IDisposable object to remove the message handler when called</returns>
+        public IDisposable OnMessage<MessageType>(string type, Action<MessageType> handler)
         {
-            OnMessageHandlers.Add(type, new ColyseusMessageHandler<MessageType>
-            {
-                Action = handler
-            });
+            ColyseusMessageHandler<MessageType> messageHandler;
+            if (OnMessageHandlers.ContainsKey(type)) {
+                messageHandler = (OnMessageHandlers[type] as ColyseusMessageHandler<MessageType>);
+                messageHandler.Action += handler;
+            } else {
+                messageHandler = new ColyseusMessageHandler<MessageType>
+                {
+                    Action = handler
+                };
+                OnMessageHandlers.Add(type, messageHandler);
+            }
+            return new MessageRemover<MessageType>(handler, messageHandler);
         }
 
         /// <summary>
@@ -302,25 +332,21 @@ public void OnMessage<MessageType>(string type, Action<MessageType> handler)
         /// <param name="type">The type of message received</param>
         /// <param name="handler"></param>
         /// <typeparam name="MessageType">The type of object this message should respond with</typeparam>
-        public void OnMessage<MessageType>(byte type, Action<MessageType> handler)
+        /// <returns>IDisposable object to remove the message handler when called</returns>
+        public IDisposable OnMessage<MessageType>(byte type, Action<MessageType> handler)
         {
-            OnMessageHandlers.Add("i" + type, new ColyseusMessageHandler<MessageType>
-            {
-                Action = handler
-            });
+            return this.OnMessage<MessageType>("i" + type, handler);
         }
 
         /// <summary>
         ///     Method to add new message handlers to the room
         /// </summary>
         /// <param name="handler"></param>
         /// <typeparam name="MessageType">The type of object this message should respond with</typeparam>
-        public void OnMessage<MessageType>(Action<MessageType> handler) where MessageType : Schema.Schema, new()
+        /// <returns>IDisposable object to remove the message handler when called</returns>
+        public IDisposable OnMessage<MessageType>(Action<MessageType> handler) where MessageType : Schema.Schema, new()
         {
-            OnMessageHandlers.Add("s" + typeof(MessageType), new ColyseusMessageHandler<MessageType>
-            {
-                Action = handler
-            });
+            return this.OnMessage<MessageType>("s" + typeof(MessageType), handler);
         }
 
         /// <summary>
@@ -372,15 +398,15 @@ protected async void ParseMessage(byte[] bytes)
 
                 if (bytes.Length > offset)
                 {
-	                try {
-		                serializer.Handshake(bytes, offset);
-	                }
-	                catch (Exception e)
-	                {
-		                await Leave(false);
-		                OnError?.Invoke(ColyseusErrorCode.SCHEMA_MISMATCH, e.Message);
-		                return;
-	                }
+                    try {
+                        serializer.Handshake(bytes, offset);
+                    }
+                    catch (Exception e)
+                    {
+                        await Leave(false);
+                        OnError?.Invoke(ColyseusErrorCode.SCHEMA_MISMATCH, e.Message);
+                        return;
+                    }
                 }
 
                 OnJoin?.Invoke();
@@ -423,7 +449,7 @@ protected async void ParseMessage(byte[] bytes)
             }
             else if (code == ColyseusProtocol.ROOM_STATE)
             {
-	            SetState(bytes, 1);
+                SetState(bytes, 1);
             }
             else if (code == ColyseusProtocol.ROOM_STATE_PATCH)
             {