Reduce decal allocs (#7226)

Content.Server/Decals/DecalSystem.cs

@@ -3,6 +3,7 @@ using Content.Server.Administration.Managers;
 using Content.Shared.Administration;
 using Content.Shared.Decals;
 using Content.Shared.Maps;
+using Microsoft.Extensions.ObjectPool;
 using Robust.Server.Player;
 using Robust.Shared.Enums;
 using Robust.Shared.Map;
@@ -14,10 +15,23 @@ namespace Content.Server.Decals
     {
         [Dependency] private readonly IPlayerManager _playerManager = default!;
         [Dependency] private readonly IAdminManager _adminManager = default!;
+        [Dependency] private readonly SharedTransformSystem _transform = default!;
 
         private readonly Dictionary<GridId, HashSet<Vector2i>> _dirtyChunks = new();
         private readonly Dictionary<IPlayerSession, Dictionary<GridId, HashSet<Vector2i>>> _previousSentChunks = new();
 
+        // If this ever gets parallelised then you'll want to increase the pooled count.
+        private ObjectPool<HashSet<Vector2i>> _chunkIndexPool =
+            new DefaultObjectPool<HashSet<Vector2i>>(
+                new DefaultPooledObjectPolicy<HashSet<Vector2i>>(), 64);
+
+        private ObjectPool<Dictionary<GridId, HashSet<Vector2i>>> _chunkViewerPool =
+            new DefaultObjectPool<Dictionary<GridId, HashSet<Vector2i>>>(
+                new DefaultPooledObjectPolicy<Dictionary<GridId, HashSet<Vector2i>>>(), 64);
+
+        // Pool if we ever parallelise.
+        private HashSet<EntityUid> _viewers = new(64);
+
         public override void Initialize()
         {
             base.Initialize();
@@ -299,14 +313,15 @@ namespace Content.Server.Decals
 
             foreach (var session in Filter.Broadcast().Recipients)
             {
-                if(session is not IPlayerSession playerSession || playerSession.Status != SessionStatus.InGame)
+                if (session is not IPlayerSession { Status: SessionStatus.InGame } playerSession)
                     continue;
 
                 var chunks = GetChunksForSession(playerSession);
-                var updatedChunks = new Dictionary<GridId, HashSet<Vector2i>>();
+                var updatedChunks = _chunkViewerPool.Get();
                 foreach (var (gridId, gridChunks) in chunks)
                 {
-                    var newChunks = new HashSet<Vector2i>(gridChunks);
+                    var newChunks = _chunkIndexPool.Get();
+                    newChunks.UnionWith(gridChunks);
                     if (_previousSentChunks[playerSession].TryGetValue(gridId, out var previousChunks))
                     {
                         newChunks.ExceptWith(previousChunks);
@@ -325,8 +340,14 @@ namespace Content.Server.Decals
                 }
 
                 if (updatedChunks.Count == 0)
+                {
+                    ReturnToPool(chunks);
+                    // Even if updatedChunks is empty we'll still return it to the pool as it may have been allocated higher.
+                    ReturnToPool(updatedChunks);
                     continue;
+                }
 
+                ReturnToPool(_previousSentChunks[playerSession]);
                 _previousSentChunks[playerSession] = chunks;
 
                 //send all gridChunks to client
@@ -336,6 +357,16 @@ namespace Content.Server.Decals
             _dirtyChunks.Clear();
         }
 
+        private void ReturnToPool(Dictionary<GridId, HashSet<Vector2i>> chunks)
+        {
+            foreach (var (_, previous) in chunks)
+            {
+                _chunkIndexPool.Return(previous);
+            }
+
+            _chunkViewerPool.Return(chunks);
+        }
+
         private void SendChunkUpdates(IPlayerSession session, Dictionary<GridId, HashSet<Vector2i>> updatedChunks)
         {
             var updatedDecals = new Dictionary<GridId, Dictionary<Vector2i, Dictionary<uint, Decal>>>();
@@ -357,7 +388,7 @@ namespace Content.Server.Decals
 
         private HashSet<EntityUid> GetSessionViewers(IPlayerSession session)
         {
-            var viewers = new HashSet<EntityUid>();
+            var viewers = _viewers;
             if (session.Status != SessionStatus.InGame || session.AttachedEntity is null)
                 return viewers;
 
@@ -373,7 +404,35 @@ namespace Content.Server.Decals
 
         private Dictionary<GridId, HashSet<Vector2i>> GetChunksForSession(IPlayerSession session)
         {
-            return GetChunksForViewers(GetSessionViewers(session));
+            var viewers = GetSessionViewers(session);
+            var chunks = GetChunksForViewers(viewers);
+            viewers.Clear();
+            return chunks;
+        }
+
+        private Dictionary<GridId, HashSet<Vector2i>> GetChunksForViewers(HashSet<EntityUid> viewers)
+        {
+            var chunks = _chunkViewerPool.Get();
+            var xformQuery = GetEntityQuery<TransformComponent>();
+
+            foreach (var viewerUid in viewers)
+            {
+                var (bounds, mapId) = CalcViewBounds(viewerUid, xformQuery.GetComponent(viewerUid));
+
+                foreach (var grid in MapManager.FindGridsIntersecting(mapId, bounds))
+                {
+                    if (!chunks.ContainsKey(grid.Index))
+                        chunks[grid.Index] = _chunkIndexPool.Get();
+
+                    var enumerator = new ChunkIndicesEnumerator(_transform.GetInvWorldMatrix(grid.GridEntityId, xformQuery).TransformBox(bounds), ChunkSize);
+
+                    while (enumerator.MoveNext(out var indices))
+                    {
+                        chunks[grid.Index].Add(indices.Value);
+                    }
+                }
+            }
+            return chunks;
         }
     }
 }

Content.Shared/Decals/SharedDecalSystem.cs

@@ -12,7 +12,6 @@ namespace Content.Shared.Decals
         [Dependency] protected readonly IPrototypeManager PrototypeManager = default!;
         [Dependency] private readonly IConfigurationManager _configurationManager = default!;
         [Dependency] protected readonly IMapManager MapManager = default!;
-        [Dependency] protected readonly SharedTransformSystem Transforms = default!;
 
         protected readonly Dictionary<GridId, Dictionary<uint, Vector2i>> ChunkIndex = new();
 
@@ -84,40 +83,13 @@ namespace Content.Shared.Decals
 
         protected virtual bool RemoveDecalHook(GridId gridId, uint uid) => true;
 
-        private (Box2 view, MapId mapId) CalcViewBounds(in EntityUid euid)
+        protected (Box2 view, MapId mapId) CalcViewBounds(in EntityUid euid, TransformComponent xform)
         {
-            var xform = EntityManager.GetComponent<TransformComponent>(euid);
-
             var view = Box2.UnitCentered.Scale(_viewSize).Translated(xform.WorldPosition);
             var map = xform.MapID;
 
             return (view, map);
         }
-
-        protected Dictionary<GridId, HashSet<Vector2i>> GetChunksForViewers(HashSet<EntityUid> viewers)
-        {
-            var chunks = new Dictionary<GridId, HashSet<Vector2i>>();
-            var xformQuery = GetEntityQuery<TransformComponent>();
-
-            foreach (var viewerUid in viewers)
-            {
-                var (bounds, mapId) = CalcViewBounds(viewerUid);
-
-                foreach (var grid in MapManager.FindGridsIntersecting(mapId, bounds))
-                {
-                    if (!chunks.ContainsKey(grid.Index))
-                        chunks[grid.Index] = new HashSet<Vector2i>();
-
-                    var enumerator = new ChunkIndicesEnumerator(Transforms.GetInvWorldMatrix(grid.GridEntityId, xformQuery).TransformBox(bounds), ChunkSize);
-
-                    while (enumerator.MoveNext(out var indices))
-                    {
-                        chunks[grid.Index].Add(indices.Value);
-                    }
-                }
-            }
-            return chunks;
-        }
     }
 
     // TODO: Pretty sure paul was moving this somewhere but just so people know