fix: Custom messaging manager exception (#3417)

<!-- Replace this block with what this PR does and why. Describe what
you'd like reviewers to know, how you applied the engineering
principles, and any interesting tradeoffs made. Delete bullet points
below that don't apply, and update the changelog section as appropriate.
-->

<!-- Add short version of the JIRA ticket to the PR title (e.g. "feat:
new shiny feature [MTT-123]") -->
[MTTB-1230](https://jira.unity3d.com/browse/MTTB-1230)
fixes: #3414 

## Changelog

- Fixed: Unregistering a message handler inside the message handler no
longer throws an exception.

## Testing and Documentation

- Includes a unit test.

<!-- Uncomment and mark items off with a * if this PR deprecates any
API:
### Deprecated API
- [ ] An `[Obsolete]` attribute was added along with a `(RemovedAfter
yyyy-mm-dd)` entry.
- [ ] An [api updater] was added.
- [ ] Deprecation of the API is explained in the CHANGELOG.
- [ ] The users can understand why this API was removed and what they
should use instead.
-->

## Backport

<!-- If this is a backport:
 - Add the following to the PR title: "\[Backport\] ..." .
 - Link to the original PR.
If this needs a backport - state this here
If a backport is not needed please provide the reason why.
If the "Backports" section is not present it will lead to a CI test
failure.
-->

Backported by #3420

Author: EmandM

com.unity.netcode.gameobjects/CHANGELOG.md

@@ -13,6 +13,7 @@ Additional documentation and release notes are available at [Multiplayer Documen
 
 ### Fixed
 
+- Fixed an exception being thrown when unregistering a custom message handler from within the registered callback. (#3417)
 
 ### Changed
 

com.unity.netcode.gameobjects/Runtime/Messaging/CustomMessageManager.cs

@@ -180,14 +180,18 @@ internal void InvokeNamedMessage(ulong hash, ulong sender, FastBufferReader read
                 // We dont know what size to use. Try every (more collision prone)
                 if (m_NamedMessageHandlers32.TryGetValue(hash, out HandleNamedMessageDelegate messageHandler32))
                 {
+                    // handler can remove itself, cache the name for metrics
+                    var messageName = m_MessageHandlerNameLookup32[hash];
                     messageHandler32(sender, reader);
-                    m_NetworkManager.NetworkMetrics.TrackNamedMessageReceived(sender, m_MessageHandlerNameLookup32[hash], bytesCount);
+                    m_NetworkManager.NetworkMetrics.TrackNamedMessageReceived(sender, messageName, bytesCount);
                 }
 
                 if (m_NamedMessageHandlers64.TryGetValue(hash, out HandleNamedMessageDelegate messageHandler64))
                 {
+                    // handler can remove itself, cache the name for metrics
+                    var messageName = m_MessageHandlerNameLookup64[hash];
                     messageHandler64(sender, reader);
-                    m_NetworkManager.NetworkMetrics.TrackNamedMessageReceived(sender, m_MessageHandlerNameLookup64[hash], bytesCount);
+                    m_NetworkManager.NetworkMetrics.TrackNamedMessageReceived(sender, messageName, bytesCount);
                 }
             }
             else
@@ -198,15 +202,19 @@ internal void InvokeNamedMessage(ulong hash, ulong sender, FastBufferReader read
                     case HashSize.VarIntFourBytes:
                         if (m_NamedMessageHandlers32.TryGetValue(hash, out HandleNamedMessageDelegate messageHandler32))
                         {
+                            // handler can remove itself, cache the name for metrics
+                            var messageName = m_MessageHandlerNameLookup32[hash];
                             messageHandler32(sender, reader);
-                            m_NetworkManager.NetworkMetrics.TrackNamedMessageReceived(sender, m_MessageHandlerNameLookup32[hash], bytesCount);
+                            m_NetworkManager.NetworkMetrics.TrackNamedMessageReceived(sender, messageName, bytesCount);
                         }
                         break;
                     case HashSize.VarIntEightBytes:
                         if (m_NamedMessageHandlers64.TryGetValue(hash, out HandleNamedMessageDelegate messageHandler64))
                         {
+                            // handler can remove itself, cache the name for metrics
+                            var messageName = m_MessageHandlerNameLookup64[hash];
                             messageHandler64(sender, reader);
-                            m_NetworkManager.NetworkMetrics.TrackNamedMessageReceived(sender, m_MessageHandlerNameLookup64[hash], bytesCount);
+                            m_NetworkManager.NetworkMetrics.TrackNamedMessageReceived(sender, messageName, bytesCount);
                         }
                         break;
                 }

com.unity.netcode.gameobjects/Tests/Runtime/Messaging/NamedMessageTests.cs

@@ -279,5 +279,38 @@ public unsafe void ErrorMessageIsPrintedWhenAttemptingToSendNamedMessageWithTooB
                 Assert.IsTrue(message.Contains($"Given message size ({msgSize} bytes) is greater than the maximum"), $"Unexpected exception: {message}");
             }
         }
+
+        [Test]
+        public void NamedMessageHandlerIsUnregisteredWithoutException()
+        {
+            var messageName = Guid.NewGuid().ToString();
+            const int numMessagesToSend = 3;
+            const int expectedMessageHandlerCallCount = 1;
+
+            var messageHandlerCalled = 0;
+            m_ServerNetworkManager.CustomMessagingManager.RegisterNamedMessageHandler(
+                messageName,
+                (_, _) =>
+                {
+                    messageHandlerCalled++;
+                    m_ServerNetworkManager.CustomMessagingManager.UnregisterNamedMessageHandler(messageName);
+                });
+
+            var messageContent = new ForceNetworkSerializeByMemcpy<Guid>(Guid.NewGuid());
+            var writer = new FastBufferWriter(1300, Allocator.Temp);
+            using (writer)
+            {
+                writer.WriteValueSafe(messageContent);
+                for (var i = 0; i < numMessagesToSend; i++)
+                {
+                    m_ServerNetworkManager.CustomMessagingManager.SendNamedMessage(
+                        messageName,
+                        m_ServerNetworkManager.LocalClientId,
+                        writer);
+                }
+            }
+
+            Assert.AreEqual(expectedMessageHandlerCallCount, messageHandlerCalled);
+        }
     }
 }