From b6a9fe5299fcfb520eab4e53b56df2a9a1de8df5 Mon Sep 17 00:00:00 2001
From: metalgearsloth <comedian_vs_clown@hotmail.com>
Date: Mon, 20 Nov 2023 00:05:01 +1100
Subject: [PATCH] Re-use RobustTree pool

---
 Robust.Server/GameStates/PvsSystem.cs | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/Robust.Server/GameStates/PvsSystem.cs b/Robust.Server/GameStates/PvsSystem.cs
index e20912762f3..8fd5895b672 100644
--- a/Robust.Server/GameStates/PvsSystem.cs
+++ b/Robust.Server/GameStates/PvsSystem.cs
@@ -84,6 +84,9 @@ private readonly ObjectPool<Dictionary<NetEntity, PvsEntityVisibility>> _visSetP
     private readonly ObjectPool<HashSet<EntityUid>> _uidSetPool
         = new DefaultObjectPool<HashSet<EntityUid>>(new SetPolicy<EntityUid>(), MaxVisPoolSize);
 
+    private readonly ObjectPool<HashSet<NetEntity>> _nentPool =
+        new DefaultObjectPool<HashSet<NetEntity>>(new SetPolicy<NetEntity>(), MaxVisPoolSize * 8);
+
     private readonly ObjectPool<Stack<NetEntity>> _stackPool
         = new DefaultObjectPool<Stack<NetEntity>>(
             new StackPolicy<NetEntity>(), MaxVisPoolSize);
@@ -95,8 +98,7 @@ private readonly ObjectPool<Stack<NetEntity>> _stackPool
     private readonly ObjectPool<HashSet<int>> _playerChunkPool =
         new DefaultObjectPool<HashSet<int>>(new SetPolicy<int>(), MaxVisPoolSize);
 
-    private readonly ObjectPool<RobustTree<NetEntity>> _treePool =
-        new DefaultObjectPool<RobustTree<NetEntity>>(new TreePolicy<NetEntity>(), MaxVisPoolSize);
+    private ObjectPool<RobustTree<NetEntity>> _treePool = default!;
 
     private readonly ObjectPool<Dictionary<MapChunkLocation, int>> _mapChunkPool =
         new DefaultObjectPool<Dictionary<MapChunkLocation, int>>(
@@ -130,6 +132,8 @@ public override void Initialize()
     {
         base.Initialize();
 
+        _treePool = new DefaultObjectPool<RobustTree<NetEntity>>(new TreePolicy<NetEntity>(_nentPool), MaxVisPoolSize);
+
         _eyeQuery = GetEntityQuery<EyeComponent>();
         _metaQuery = GetEntityQuery<MetaDataComponent>();
         _xformQuery = GetEntityQuery<TransformComponent>();
@@ -1371,10 +1375,16 @@ private EntityUid[] GetSessionViewers(ICommonSession session)
 
     public sealed class TreePolicy<T> : PooledObjectPolicy<RobustTree<T>> where T : notnull
     {
+        private readonly ObjectPool<HashSet<T>> _pool;
+
+        public TreePolicy(ObjectPool<HashSet<T>> pool)
+        {
+            _pool = pool;
+        }
+
         public override RobustTree<T> Create()
         {
-            var pool = new DefaultObjectPool<HashSet<T>>(new SetPolicy<T>(), MaxVisPoolSize);
-            return new RobustTree<T>(pool);
+            return new RobustTree<T>(_pool);
         }
 
         public override bool Return(RobustTree<T> obj)