Re-use RobustTree pool

metalgearsloth · Nov 19, 2023 · b6a9fe5 · b6a9fe5
1 parent abad5a1
commit b6a9fe5
Showing 1 changed file with 14 additions and 4 deletions.
diff --git 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)