fix: remove ownership key not found exception (#1532)

Author: 0xFA11

com.unity.netcode.gameobjects/CHANGELOG.md

@@ -25,8 +25,8 @@ Additional documentation and release notes are available at [Multiplayer Documen
 - NetworkConfig will no longer throw an OverflowException in GetConfig() when ForceSamePrefabs is enabled and the number of prefabs causes the config blob size to exceed 1300 bytes. (#1385)
 - Fixed NetworkVariable not calling NetworkSerialize on INetworkSerializable types (#1383)
 - Fixed NullReferenceException on ImportReferences call in NetworkBehaviourILPP (#1434)
-
 - Fixed NetworkObjects not being despawned before they are destroyed during shutdown for client, host, and server instances. (#1390)
+- Fixed KeyNotFound exception when removing ownership of a newly spawned NetworkObject that is already owned by the server. (#1500)
 
 ### Changed
 

com.unity.netcode.gameobjects/Runtime/Spawning/NetworkSpawnManager.cs

@@ -187,28 +187,44 @@ internal void RemoveOwnership(NetworkObject networkObject)
                 throw new SpawnStateException("Object is not spawned");
             }
 
-            for (int i = NetworkManager.ConnectedClients[networkObject.OwnerClientId].OwnedObjects.Count - 1;
-                i > -1;
-                i--)
+            // If we made it here then we are the server and if the server is determined to already be the owner
+            // then ignore the RemoveOwnership invocation.
+            if (networkObject.OwnerClientId == NetworkManager.ServerClientId)
             {
-                if (NetworkManager.ConnectedClients[networkObject.OwnerClientId].OwnedObjects[i] == networkObject)
+                return;
+            }
+
+            // Make sure the connected client entry exists before trying to remove ownership.
+            if (TryGetNetworkClient(networkObject.OwnerClientId, out NetworkClient networkClient))
+            {
+                for (int i = networkClient.OwnedObjects.Count - 1; i > -1; i--)
                 {
-                    NetworkManager.ConnectedClients[networkObject.OwnerClientId].OwnedObjects.RemoveAt(i);
+                    if (networkClient.OwnedObjects[i] == networkObject)
+                    {
+                        networkClient.OwnedObjects.RemoveAt(i);
+                    }
                 }
-            }
 
-            networkObject.OwnerClientIdInternal = null;
+                networkObject.OwnerClientIdInternal = null;
 
-            var message = new ChangeOwnershipMessage
-            {
-                NetworkObjectId = networkObject.NetworkObjectId,
-                OwnerClientId = networkObject.OwnerClientId
-            };
-            var size = NetworkManager.SendMessage(message, NetworkDelivery.ReliableSequenced, NetworkManager.ConnectedClientsIds);
+                var message = new ChangeOwnershipMessage
+                {
+                    NetworkObjectId = networkObject.NetworkObjectId,
+                    OwnerClientId = networkObject.OwnerClientId
+                };
+                var size = NetworkManager.SendMessage(message, NetworkDelivery.ReliableSequenced, NetworkManager.ConnectedClientsIds);
 
-            foreach (var client in NetworkManager.ConnectedClients)
+                foreach (var client in NetworkManager.ConnectedClients)
+                {
+                    NetworkManager.NetworkMetrics.TrackOwnershipChangeSent(client.Key, networkObject, size);
+                }
+            }
+            else
             {
-                NetworkManager.NetworkMetrics.TrackOwnershipChangeSent(client.Key, networkObject, size);
+                if (NetworkLog.CurrentLogLevel <= LogLevel.Normal)
+                {
+                    NetworkLog.LogWarning($"No connected clients prior to removing ownership for {networkObject.name}.  Make sure you are not initializing or shutting down when removing ownership.");
+                }
             }
         }
 

com.unity.netcode.gameobjects/Tests/Runtime/NetworkObject/NetworkObjectOwnershipTests.cs

@@ -115,6 +115,9 @@ public IEnumerator TestOwnershipCallbacks()
                 Assert.That(clientNetworkManager.SpawnManager.SpawnedObjects.ContainsKey(dummyNetworkObjectId));
             }
 
+            // Verifies that removing the ownership when the default (server) is already set does not cause
+            // a Key Not Found Exception
+            m_ServerNetworkManager.SpawnManager.RemoveOwnership(dummyNetworkObject);
 
             var serverObject = m_ServerNetworkManager.SpawnManager.SpawnedObjects[dummyNetworkObjectId];
             var clientObject = m_ClientNetworkManagers[0].SpawnManager.SpawnedObjects[dummyNetworkObjectId];