Update component query benchmarks (#27967)

* Add more component query benchmarks.

* Rename benchmark

Content.Benchmarks/ComponentQueryBenchmark.cs

@@ -0,0 +1,273 @@
+#nullable enable
+using System;
+using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
+using BenchmarkDotNet.Attributes;
+using BenchmarkDotNet.Configs;
+using Content.IntegrationTests;
+using Content.IntegrationTests.Pair;
+using Content.Shared.Clothing.Components;
+using Content.Shared.Doors.Components;
+using Content.Shared.Item;
+using Robust.Server.GameObjects;
+using Robust.Shared;
+using Robust.Shared.Analyzers;
+using Robust.Shared.GameObjects;
+using Robust.Shared.Map;
+using Robust.Shared.Map.Components;
+using Robust.Shared.Random;
+
+namespace Content.Benchmarks;
+
+/// <summary>
+/// Benchmarks for comparing the speed of various component fetching/lookup related methods, including directed event
+/// subscriptions
+/// </summary>
+[Virtual]
+[GroupBenchmarksBy(BenchmarkLogicalGroupRule.ByCategory)]
+[CategoriesColumn]
+public class ComponentQueryBenchmark
+{
+    public const string Map = "Maps/atlas.yml";
+
+    private TestPair _pair = default!;
+    private IEntityManager _entMan = default!;
+    private MapId _mapId = new(10);
+    private EntityQuery<ItemComponent> _itemQuery;
+    private EntityQuery<ClothingComponent> _clothingQuery;
+    private EntityQuery<MapComponent> _mapQuery;
+    private EntityUid[] _items = default!;
+
+    [GlobalSetup]
+    public void Setup()
+    {
+        ProgramShared.PathOffset = "../../../../";
+        PoolManager.Startup(typeof(QueryBenchSystem).Assembly);
+
+        _pair = PoolManager.GetServerClient().GetAwaiter().GetResult();
+        _entMan = _pair.Server.ResolveDependency<IEntityManager>();
+
+        _itemQuery = _entMan.GetEntityQuery<ItemComponent>();
+        _clothingQuery = _entMan.GetEntityQuery<ClothingComponent>();
+        _mapQuery = _entMan.GetEntityQuery<MapComponent>();
+
+        _pair.Server.ResolveDependency<IRobustRandom>().SetSeed(42);
+        _pair.Server.WaitPost(() =>
+        {
+            var success = _entMan.System<MapLoaderSystem>().TryLoad(_mapId, Map, out _);
+            if (!success)
+                throw new Exception("Map load failed");
+            _pair.Server.MapMan.DoMapInitialize(_mapId);
+        }).GetAwaiter().GetResult();
+
+        _items = new EntityUid[_entMan.Count<ItemComponent>()];
+        var i = 0;
+        var enumerator = _entMan.AllEntityQueryEnumerator<ItemComponent>();
+        while (enumerator.MoveNext(out var uid, out _))
+        {
+            _items[i++] = uid;
+        }
+    }
+
+    [GlobalCleanup]
+    public async Task Cleanup()
+    {
+        await _pair.DisposeAsync();
+        PoolManager.Shutdown();
+    }
+
+    #region TryComp
+
+    /// <summary>
+    /// Baseline TryComp benchmark. When the benchmark was created, around 40% of the items were clothing.
+    /// </summary>
+    [Benchmark(Baseline = true)]
+    [BenchmarkCategory("TryComp")]
+    public int TryComp()
+    {
+        var hashCode = 0;
+        foreach (var uid in _items)
+        {
+            if (_clothingQuery.TryGetComponent(uid, out var clothing))
+                hashCode = HashCode.Combine(hashCode, clothing.GetHashCode());
+        }
+        return hashCode;
+    }
+
+    /// <summary>
+    /// Variant of <see cref="TryComp"/> that is meant to always fail to get a component.
+    /// </summary>
+    [Benchmark]
+    [BenchmarkCategory("TryComp")]
+    public int TryCompFail()
+    {
+        var hashCode = 0;
+        foreach (var uid in _items)
+        {
+            if (_mapQuery.TryGetComponent(uid, out var map))
+                hashCode = HashCode.Combine(hashCode, map.GetHashCode());
+        }
+        return hashCode;
+    }
+
+    /// <summary>
+    /// Variant of <see cref="TryComp"/> that is meant to always succeed getting a component.
+    /// </summary>
+    [Benchmark]
+    [BenchmarkCategory("TryComp")]
+    public int TryCompSucceed()
+    {
+        var hashCode = 0;
+        foreach (var uid in _items)
+        {
+            if (_itemQuery.TryGetComponent(uid, out var item))
+                hashCode = HashCode.Combine(hashCode, item.GetHashCode());
+        }
+        return hashCode;
+    }
+
+    /// <summary>
+    /// Variant of <see cref="TryComp"/> that uses `Resolve()` to try get the component.
+    /// </summary>
+    [Benchmark]
+    [BenchmarkCategory("TryComp")]
+    public int Resolve()
+    {
+        var hashCode = 0;
+        foreach (var uid in _items)
+        {
+            DoResolve(uid, ref hashCode);
+        }
+        return hashCode;
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void DoResolve(EntityUid uid, ref int hash, ClothingComponent? clothing = null)
+    {
+        if (_clothingQuery.Resolve(uid, ref clothing, false))
+            hash = HashCode.Combine(hash, clothing.GetHashCode());
+    }
+
+    #endregion
+
+    #region Enumeration
+
+    [Benchmark]
+    [BenchmarkCategory("Item Enumerator")]
+    public int SingleItemEnumerator()
+    {
+        var hashCode = 0;
+        var enumerator = _entMan.AllEntityQueryEnumerator<ItemComponent>();
+        while (enumerator.MoveNext(out var item))
+        {
+            hashCode = HashCode.Combine(hashCode, item.GetHashCode());
+        }
+
+        return hashCode;
+    }
+
+    [Benchmark]
+    [BenchmarkCategory("Item Enumerator")]
+    public int DoubleItemEnumerator()
+    {
+        var hashCode = 0;
+        var enumerator = _entMan.AllEntityQueryEnumerator<ClothingComponent, ItemComponent>();
+        while (enumerator.MoveNext(out _, out var item))
+        {
+            hashCode = HashCode.Combine(hashCode, item.GetHashCode());
+        }
+
+        return hashCode;
+    }
+
+    [Benchmark]
+    [BenchmarkCategory("Item Enumerator")]
+    public int TripleItemEnumerator()
+    {
+        var hashCode = 0;
+        var enumerator = _entMan.AllEntityQueryEnumerator<ClothingComponent, ItemComponent, TransformComponent>();
+        while (enumerator.MoveNext(out _, out _, out var xform))
+        {
+            hashCode = HashCode.Combine(hashCode, xform.GetHashCode());
+        }
+
+        return hashCode;
+    }
+
+    [Benchmark]
+    [BenchmarkCategory("Airlock Enumerator")]
+    public int SingleAirlockEnumerator()
+    {
+        var hashCode = 0;
+        var enumerator = _entMan.AllEntityQueryEnumerator<AirlockComponent>();
+        while (enumerator.MoveNext(out var airlock))
+        {
+            hashCode = HashCode.Combine(hashCode, airlock.GetHashCode());
+        }
+
+        return hashCode;
+    }
+
+    [Benchmark]
+    [BenchmarkCategory("Airlock Enumerator")]
+    public int DoubleAirlockEnumerator()
+    {
+        var hashCode = 0;
+        var enumerator = _entMan.AllEntityQueryEnumerator<AirlockComponent, DoorComponent>();
+        while (enumerator.MoveNext(out _, out var door))
+        {
+            hashCode = HashCode.Combine(hashCode, door.GetHashCode());
+        }
+
+        return hashCode;
+    }
+
+    [Benchmark]
+    [BenchmarkCategory("Airlock Enumerator")]
+    public int TripleAirlockEnumerator()
+    {
+        var hashCode = 0;
+        var enumerator = _entMan.AllEntityQueryEnumerator<AirlockComponent, DoorComponent, TransformComponent>();
+        while (enumerator.MoveNext(out _, out _, out var xform))
+        {
+            hashCode = HashCode.Combine(hashCode, xform.GetHashCode());
+        }
+
+        return hashCode;
+    }
+
+    #endregion
+
+    [Benchmark(Baseline = true)]
+    [BenchmarkCategory("Events")]
+    public int StructEvents()
+    {
+        var ev = new QueryBenchEvent();
+        foreach (var uid in _items)
+        {
+            _entMan.EventBus.RaiseLocalEvent(uid, ref ev);
+        }
+
+        return ev.HashCode;
+    }
+}
+
+[ByRefEvent]
+public struct QueryBenchEvent
+{
+    public int HashCode;
+}
+
+public sealed class QueryBenchSystem : EntitySystem
+{
+    public override void Initialize()
+    {
+        base.Initialize();
+        SubscribeLocalEvent<ClothingComponent, QueryBenchEvent>(OnEvent);
+    }
+
+    private void OnEvent(EntityUid uid, ClothingComponent component, ref QueryBenchEvent args)
+    {
+        args.HashCode = HashCode.Combine(args.HashCode, component.GetHashCode());
+    }
+}

Content.Benchmarks/EntityQueryBenchmark.cs

@@ -1,137 +0,0 @@
-#nullable enable
-using System;
-using System.Threading.Tasks;
-using BenchmarkDotNet.Attributes;
-using Content.IntegrationTests;
-using Content.IntegrationTests.Pair;
-using Content.Shared.Clothing.Components;
-using Content.Shared.Item;
-using Robust.Server.GameObjects;
-using Robust.Shared;
-using Robust.Shared.Analyzers;
-using Robust.Shared.GameObjects;
-using Robust.Shared.Map;
-using Robust.Shared.Random;
-
-namespace Content.Benchmarks;
-
-[Virtual]
-public class EntityQueryBenchmark
-{
-    public const string Map = "Maps/atlas.yml";
-
-    private TestPair _pair = default!;
-    private IEntityManager _entMan = default!;
-    private MapId _mapId = new MapId(10);
-    private EntityQuery<ClothingComponent> _clothingQuery;
-
-    [GlobalSetup]
-    public void Setup()
-    {
-        ProgramShared.PathOffset = "../../../../";
-        PoolManager.Startup(null);
-
-        _pair = PoolManager.GetServerClient().GetAwaiter().GetResult();
-        _entMan = _pair.Server.ResolveDependency<IEntityManager>();
-
-        _pair.Server.ResolveDependency<IRobustRandom>().SetSeed(42);
-        _pair.Server.WaitPost(() =>
-        {
-            var success = _entMan.System<MapLoaderSystem>().TryLoad(_mapId, Map, out _);
-            if (!success)
-                throw new Exception("Map load failed");
-            _pair.Server.MapMan.DoMapInitialize(_mapId);
-        }).GetAwaiter().GetResult();
-
-        _clothingQuery = _entMan.GetEntityQuery<ClothingComponent>();
-
-        // Apparently ~40% of entities are items, and 1 in 6 of those are clothing.
-        /*
-        var entCount = _entMan.EntityCount;
-        var itemCount = _entMan.Count<ItemComponent>();
-        var clothingCount = _entMan.Count<ClothingComponent>();
-        var itemRatio = (float) itemCount / entCount;
-        var clothingRatio = (float) clothingCount / entCount;
-        Console.WriteLine($"Entities: {entCount}. Items: {itemRatio:P2}. Clothing: {clothingRatio:P2}.");
-        */
-    }
-
-    [GlobalCleanup]
-    public async Task Cleanup()
-    {
-        await _pair.DisposeAsync();
-        PoolManager.Shutdown();
-    }
-
-    [Benchmark]
-    public int HasComponent()
-    {
-        var hashCode = 0;
-        var enumerator = _entMan.AllEntityQueryEnumerator<ItemComponent>();
-        while (enumerator.MoveNext(out var uid, out var _))
-        {
-            if (_entMan.HasComponent<ClothingComponent>(uid))
-                hashCode = HashCode.Combine(hashCode, uid.Id);
-        }
-
-        return hashCode;
-    }
-
-    [Benchmark]
-    public int HasComponentQuery()
-    {
-        var hashCode = 0;
-        var enumerator = _entMan.AllEntityQueryEnumerator<ItemComponent>();
-        while (enumerator.MoveNext(out var uid, out var _))
-        {
-            if (_clothingQuery.HasComponent(uid))
-                hashCode = HashCode.Combine(hashCode, uid.Id);
-        }
-
-        return hashCode;
-    }
-
-    [Benchmark]
-    public int TryGetComponent()
-    {
-        var hashCode = 0;
-        var enumerator = _entMan.AllEntityQueryEnumerator<ItemComponent>();
-        while (enumerator.MoveNext(out var uid, out var _))
-        {
-            if (_entMan.TryGetComponent(uid, out ClothingComponent? clothing))
-                hashCode = HashCode.Combine(hashCode, clothing.GetHashCode());
-        }
-
-        return hashCode;
-    }
-
-    [Benchmark]
-    public int TryGetComponentQuery()
-    {
-        var hashCode = 0;
-        var enumerator = _entMan.AllEntityQueryEnumerator<ItemComponent>();
-        while (enumerator.MoveNext(out var uid, out var _))
-        {
-            if (_clothingQuery.TryGetComponent(uid, out var clothing))
-                hashCode = HashCode.Combine(hashCode, clothing.GetHashCode());
-        }
-
-        return hashCode;
-    }
-
-    /// <summary>
-    /// Enumerate all entities with both an item and clothing component.
-    /// </summary>
-    [Benchmark]
-    public int Enumerator()
-    {
-        var hashCode = 0;
-        var enumerator = _entMan.AllEntityQueryEnumerator<ItemComponent, ClothingComponent>();
-        while (enumerator.MoveNext(out var _, out var clothing))
-        {
-            hashCode = HashCode.Combine(hashCode, clothing.GetHashCode());
-        }
-
-        return hashCode;
-    }
-}

Content.Benchmarks/MapLoadBenchmark.cs

@@ -26,7 +26,7 @@ public class MapLoadBenchmark
     public void Setup()
     {
         ProgramShared.PathOffset = "../../../../";
-        PoolManager.Startup(null);
+        PoolManager.Startup();
 
         _pair = PoolManager.GetServerClient().GetAwaiter().GetResult();
         var server = _pair.Server;

Content.Benchmarks/PvsBenchmark.cs

@@ -49,7 +49,7 @@ public class PvsBenchmark
 #if !DEBUG
         ProgramShared.PathOffset = "../../../../";
 #endif
-        PoolManager.Startup(null);
+        PoolManager.Startup();
 
         _pair = PoolManager.GetServerClient().GetAwaiter().GetResult();
         _entMan = _pair.Server.ResolveDependency<IEntityManager>();

Content.Benchmarks/SpawnEquipDeleteBenchmark.cs

@@ -32,7 +32,7 @@ public class SpawnEquipDeleteBenchmark
     public async Task SetupAsync()
     {
         ProgramShared.PathOffset = "../../../../";
-        PoolManager.Startup(null);
+        PoolManager.Startup();
         _pair = await PoolManager.GetServerClient();
         var server = _pair.Server;
 

Content.IntegrationTests/PoolManager.Prototypes.cs

@@ -15,11 +15,8 @@ public static partial class PoolManager
                                        | BindingFlags.Public
                                        | BindingFlags.DeclaredOnly;
 
-    private static void DiscoverTestPrototypes(Assembly? assembly = null)
+    private static void DiscoverTestPrototypes(Assembly assembly)
     {
-        assembly ??= typeof(PoolManager).Assembly;
-        _testPrototypes.Clear();
-
         foreach (var type in assembly.GetTypes())
         {
             foreach (var field in type.GetFields(Flags))

Content.IntegrationTests/PoolManager.cs

@@ -42,6 +42,8 @@ public static partial class PoolManager
     private static bool _dead;
     private static Exception? _poolFailureReason;
 
+    private static HashSet<Assembly> _contentAssemblies = default!;
+
     public static async Task<(RobustIntegrationTest.ServerIntegrationInstance, PoolTestLogHandler)> GenerateServer(
         PoolSettings poolSettings,
         TextWriter testOut)
@@ -54,12 +56,7 @@ public static partial class PoolManager
                 LoadConfigAndUserData = false,
                 LoadContentResources = !poolSettings.NoLoadContent,
             },
-            ContentAssemblies = new[]
+            ContentAssemblies = _contentAssemblies.ToArray()
-            {
-                typeof(Shared.Entry.EntryPoint).Assembly,
-                typeof(Server.Entry.EntryPoint).Assembly,
-                typeof(PoolManager).Assembly
-            }
         };
 
         var logHandler = new PoolTestLogHandler("SERVER");
@@ -140,7 +137,7 @@ public static partial class PoolManager
             {
                 typeof(Shared.Entry.EntryPoint).Assembly,
                 typeof(Client.Entry.EntryPoint).Assembly,
-                typeof(PoolManager).Assembly
+                typeof(PoolManager).Assembly,
             }
         };
 
@@ -422,13 +419,26 @@ we are just going to end this here to save a lot of time. This is the exception
     /// <summary>
     /// Initialize the pool manager.
     /// </summary>
-    /// <param name="assembly">Assembly to search for to discover extra test prototypes.</param>
+    /// <param name="extraAssemblies">Assemblies to search for to discover extra prototypes and systems.</param>
-    public static void Startup(Assembly? assembly)
+    public static void Startup(params Assembly[] extraAssemblies)
     {
         if (_initialized)
             throw new InvalidOperationException("Already initialized");
 
         _initialized = true;
+        _contentAssemblies =
+        [
+            typeof(Shared.Entry.EntryPoint).Assembly,
+            typeof(Server.Entry.EntryPoint).Assembly,
+            typeof(PoolManager).Assembly
+        ];
+        _contentAssemblies.UnionWith(extraAssemblies);
+
+        _testPrototypes.Clear();
+        DiscoverTestPrototypes(typeof(PoolManager).Assembly);
+        foreach (var assembly in extraAssemblies)
+        {
             DiscoverTestPrototypes(assembly);
         }
     }
+}

Content.IntegrationTests/PoolManagerTestEventHandler.cs

@@ -13,7 +13,7 @@ public sealed class PoolManagerTestEventHandler
     [OneTimeSetUp]
     public void Setup()
     {
-        PoolManager.Startup(typeof(PoolManagerTestEventHandler).Assembly);
+        PoolManager.Startup();
         // If the tests seem to be stuck, we try to end it semi-nicely
         _ = Task.Delay(MaximumTotalTestingTimeLimit).ContinueWith(_ =>
         {

Content.MapRenderer/Program.cs

@@ -29,7 +29,7 @@ namespace Content.MapRenderer
             if (!CommandLineArguments.TryParse(args, out var arguments))
                 return;
 
-            PoolManager.Startup(null);
+            PoolManager.Startup();
             if (arguments.Maps.Count == 0)
             {
                 Console.WriteLine("Didn't specify any maps to paint! Loading the map list...");

Content.YAMLLinter/Program.cs

@@ -17,7 +17,7 @@ namespace Content.YAMLLinter
     {
         private static async Task<int> Main(string[] _)
         {
-            PoolManager.Startup(null);
+            PoolManager.Startup();
             var stopwatch = new Stopwatch();
             stopwatch.Start();