frozone

Author: metalgearsloth

Robust.Server.Testing/RobustServerSimulation.cs

@@ -251,6 +251,7 @@ public ISimulation InitializeInstance()
             container.Register<ISerializationManager, SerializationManager>();
             container.Register<IRobustRandom, RobustRandom>();
             container.Register<IPrototypeManager, ServerPrototypeManager>();
+            container.Register<IPrototypeManagerInternal, ServerPrototypeManager>();
             container.Register<IComponentFactory, ComponentFactory>();
             container.Register<IEntitySystemManager, EntitySystemManager>();
             container.Register<IManifoldManager, CollisionManager>();

Robust.Shared/Prototypes/PrototypeManager.EntityLoad.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Frozen;
 using System.Collections.Generic;
 using Robust.Shared.EntitySerialization;
 using Robust.Shared.GameObjects;
@@ -8,7 +9,8 @@ namespace Robust.Shared.Prototypes;
 
 public abstract partial class PrototypeManager
 {
-    private readonly Dictionary<string, PrototypeComponentRegistration[]> _entityComponentRegistrations = new();
+    private FrozenDictionary<string, PrototypeComponentRegistration[]> _entityComponentRegistrations
+        = FrozenDictionary<string, PrototypeComponentRegistration[]>.Empty;
 
     void IPrototypeManagerInternal.LoadEntity(Entity<MetaDataComponent> ent, IEntityLoadContext? context)
     {
@@ -34,6 +36,34 @@ private PrototypeComponentRegistration[] GetComponentRegistrations(EntityPrototy
         if (_entityComponentRegistrations.TryGetValue(prototype.ID, out var cached))
             return cached;
 
+        // This should only happen if an entity is being loaded before prototype resolution has rebuilt the cache.
+        return BuildComponentRegistrations(prototype);
+    }
+
+    private void RebuildEntityComponentRegistrationCache()
+    {
+        if (!TryGetInstances<EntityPrototype>(out var prototypes))
+        {
+            ClearEntityComponentRegistrationCache();
+            return;
+        }
+
+        var registrations = new Dictionary<string, PrototypeComponentRegistration[]>(prototypes.Count);
+        foreach (var (id, prototype) in prototypes)
+        {
+            registrations[id] = BuildComponentRegistrations(prototype);
+        }
+
+        _entityComponentRegistrations = registrations.ToFrozenDictionary();
+    }
+
+    private void ClearEntityComponentRegistrationCache()
+    {
+        _entityComponentRegistrations = FrozenDictionary<string, PrototypeComponentRegistration[]>.Empty;
+    }
+
+    private PrototypeComponentRegistration[] BuildComponentRegistrations(EntityPrototype prototype)
+    {
         var registrations = new PrototypeComponentRegistration[prototype.Components.Count];
         var i = 0;
 
@@ -42,7 +72,6 @@ private PrototypeComponentRegistration[] GetComponentRegistrations(EntityPrototy
             registrations[i++] = new PrototypeComponentRegistration(name, entry, _factory.GetRegistration(name));
         }
 
-        _entityComponentRegistrations[prototype.ID] = registrations;
         return registrations;
     }
 

Robust.Shared/Prototypes/PrototypeManager.cs

@@ -295,6 +295,7 @@ public void Clear()
         {
             _kindNames.Clear();
             _kinds = FrozenDictionary<Type, KindData>.Empty;
+            ClearEntityComponentRegistrationCache();
         }
 
         /// <inheritdoc />
@@ -478,11 +479,16 @@ void AddToQueue(string id)
         private void Freeze(IEnumerable<KindData> kinds)
         {
             var st = RStopwatch.StartNew();
+            var frozeEntityPrototypes = false;
             foreach (var kind in kinds)
             {
                 kind.Freeze();
+                frozeEntityPrototypes |= kind.Type == typeof(EntityPrototype);
             }
 
+            if (frozeEntityPrototypes)
+                RebuildEntityComponentRegistrationCache();
+
             // fun fact: Sawmill can be null in tests????
             Sawmill?.Verbose($"Freezing prototype instances took {st.Elapsed.TotalMilliseconds:f2}ms");
         }
@@ -1170,7 +1176,6 @@ private void OnReload(PrototypesReloadedEventArgs args)
                 foreach (var id in modified.Modified.Keys)
                 {
                     _prototypeDataCache.Remove(id);
-                    _entityComponentRegistrations.Remove(id);
                 }
             }
 
@@ -1180,7 +1185,6 @@ private void OnReload(PrototypesReloadedEventArgs args)
             foreach (var id in removed)
             {
                 _prototypeDataCache.Remove(id);
-                _entityComponentRegistrations.Remove(id);
             }
         }