From fd0bb9fd8c92d9b256a23e94ee0a42f2daa5af63 Mon Sep 17 00:00:00 2001
From: mirrorcult <lunarautomaton6@gmail.com>
Date: Mon, 4 Apr 2022 23:00:52 -0700
Subject: [PATCH] Storage system opts (#7427)

---
 .../Components/ServerStorageComponent.cs      | 11 +++++++
 .../Storage/EntitySystems/StorageSystem.cs    | 32 +++++++++----------
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/Content.Server/Storage/Components/ServerStorageComponent.cs b/Content.Server/Storage/Components/ServerStorageComponent.cs
index c715e85876..3a557b4e51 100644
--- a/Content.Server/Storage/Components/ServerStorageComponent.cs
+++ b/Content.Server/Storage/Components/ServerStorageComponent.cs
@@ -366,6 +366,7 @@ namespace Content.Server.Storage.Components
                 SubscribedSessions.Add(session);
             }
 
+            _entityManager.EnsureComponent<ActiveStorageComponent>(Owner);
             if (SubscribedSessions.Count == 1)
                 UpdateStorageVisualization();
         }
@@ -388,7 +389,14 @@ namespace Content.Server.Storage.Components
             CloseNestedInterfaces(session);
 
             if (SubscribedSessions.Count == 0)
+            {
                 UpdateStorageVisualization();
+
+                if (_entityManager.HasComponent<ActiveStorageComponent>(Owner))
+                {
+                    _entityManager.RemoveComponent<ActiveStorageComponent>(Owner);
+                }
+            }
         }
 
         /// <summary>
@@ -635,4 +643,7 @@ namespace Content.Server.Storage.Components
             SoundSystem.Play(Filter.Pvs(Owner), StorageSoundCollection.GetSound(), Owner, AudioParams.Default);
         }
     }
+
+    [RegisterComponent]
+    public sealed class ActiveStorageComponent : Component {}
 }
diff --git a/Content.Server/Storage/EntitySystems/StorageSystem.cs b/Content.Server/Storage/EntitySystems/StorageSystem.cs
index 44602bbae0..5034c1a8df 100644
--- a/Content.Server/Storage/EntitySystems/StorageSystem.cs
+++ b/Content.Server/Storage/EntitySystems/StorageSystem.cs
@@ -13,6 +13,7 @@ using Robust.Server.Player;
 using Robust.Shared.Containers;
 using Robust.Shared.Random;
 using Robust.Shared.Timing;
+using Robust.Shared.Utility;
 
 namespace Content.Server.Storage.EntitySystems
 {
@@ -23,8 +24,6 @@ namespace Content.Server.Storage.EntitySystems
         [Dependency] private readonly IRobustRandom _random = default!;
         [Dependency] private readonly DisposalUnitSystem _disposalSystem = default!;
 
-        private readonly List<IPlayerSession> _sessionCache = new();
-
         /// <inheritdoc />
         public override void Initialize()
         {
@@ -88,7 +87,7 @@ namespace Content.Server.Storage.EntitySystems
         /// <inheritdoc />
         public override void Update(float frameTime)
         {
-            foreach (var component in EntityManager.EntityQuery<ServerStorageComponent>())
+            foreach (var (_, component) in EntityManager.EntityQuery<ActiveStorageComponent, ServerStorageComponent>())
             {
                 CheckSubscribedEntities(component);
             }
@@ -261,32 +260,33 @@ namespace Content.Server.Storage.EntitySystems
 
         private void CheckSubscribedEntities(ServerStorageComponent storageComp)
         {
+            var xform = Transform(storageComp.Owner);
+            var storagePos = xform.WorldPosition;
+            var storageMap = xform.MapID;
 
-            // We have to cache the set of sessions because Unsubscribe modifies the original.
-            _sessionCache.Clear();
-            _sessionCache.AddRange(storageComp.SubscribedSessions);
+            var remove = new RemQueue<IPlayerSession>();
 
-            if (_sessionCache.Count == 0)
-                return;
-
-            var storagePos = EntityManager.GetComponent<TransformComponent>(storageComp.Owner).WorldPosition;
-            var storageMap = EntityManager.GetComponent<TransformComponent>(storageComp.Owner).MapID;
-
-            foreach (var session in _sessionCache)
+            foreach (var session in storageComp.SubscribedSessions)
             {
                 // The component manages the set of sessions, so this invalid session should be removed soon.
                 if (session.AttachedEntity is not {} attachedEntity || !EntityManager.EntityExists(attachedEntity))
                     continue;
 
-                if (storageMap != EntityManager.GetComponent<TransformComponent>(attachedEntity).MapID)
+                var attachedXform = Transform(attachedEntity);
+                if (storageMap != attachedXform.MapID)
                     continue;
 
-                var distanceSquared = (storagePos - EntityManager.GetComponent<TransformComponent>(attachedEntity).WorldPosition).LengthSquared;
+                var distanceSquared = (storagePos - attachedXform.WorldPosition).LengthSquared;
                 if (distanceSquared > SharedInteractionSystem.InteractionRangeSquared)
                 {
-                    storageComp.UnsubscribeSession(session);
+                    remove.Add(session);
                 }
             }
+
+            foreach (var session in remove)
+            {
+                storageComp.UnsubscribeSession(session);
+            }
         }
     }
 }