Make DoAfter visible to all (#2183)

Co-authored-by: Metal Gear Sloth <metalgearsloth@gmail.com>

Content.Client/GameObjects/Components/DoAfterComponent.cs

@@ -3,10 +3,7 @@ using System;
 using System.Collections.Generic;
 using Content.Client.GameObjects.EntitySystems.DoAfter;
 using Content.Shared.GameObjects.Components;
-using Robust.Client.GameObjects;
 using Robust.Shared.GameObjects;
-using Robust.Shared.GameObjects.Systems;
-using Robust.Shared.Interfaces.GameObjects;
 using Robust.Shared.Interfaces.Network;
 using Robust.Shared.Interfaces.Timing;
 using Robust.Shared.IoC;
@@ -19,62 +16,112 @@ namespace Content.Client.GameObjects.Components
     {
         public override string Name => "DoAfter";
 
-        public IReadOnlyDictionary<byte, DoAfterMessage> DoAfters => _doAfters;
+        public IReadOnlyDictionary<byte, ClientDoAfter> DoAfters => _doAfters;
-        private readonly Dictionary<byte, DoAfterMessage> _doAfters = new Dictionary<byte, DoAfterMessage>();
+        private readonly Dictionary<byte, ClientDoAfter> _doAfters = new Dictionary<byte, ClientDoAfter>();
         
-        public readonly List<(TimeSpan CancelTime, DoAfterMessage Message)> CancelledDoAfters = 
+        public readonly List<(TimeSpan CancelTime, ClientDoAfter Message)> CancelledDoAfters = 
-                     new List<(TimeSpan CancelTime, DoAfterMessage Message)>();
+                     new List<(TimeSpan CancelTime, ClientDoAfter Message)>();
+
+        public DoAfterGui? Gui { get; set; }
 
         public override void HandleNetworkMessage(ComponentMessage message, INetChannel netChannel, ICommonSession? session = null)
         {
             base.HandleNetworkMessage(message, netChannel, session);
             switch (message)
             {
-                case DoAfterMessage msg:
-                    _doAfters.Add(msg.ID, msg);
-                    EntitySystem.Get<DoAfterSystem>().Gui?.AddDoAfter(msg);
-                    break;
                 case CancelledDoAfterMessage msg:
                     Cancel(msg.ID);
                     break;
             }
         }
 
-        public override void HandleMessage(ComponentMessage message, IComponent? component)
+        public override void OnAdd()
         {
-            base.HandleMessage(message, component);
+            base.OnAdd();
-            switch (message)
+            Enable();
+        }
+
+        public override void OnRemove()
+        {
+            base.OnRemove();
+            Disable();
+        }
+
+        /// <summary>
+        ///     For handling PVS so we dispose of controls if they go out of range
+        /// </summary>
+        public void Enable()
+        {
+            if (Gui != null && !Gui.Disposed)
+                return;
+            
+            Gui = new DoAfterGui {AttachedEntity = Owner};
+            
+            foreach (var (_, doAfter) in _doAfters)
             {
-                case PlayerDetachedMsg _:
+                Gui.AddDoAfter(doAfter);
-                    _doAfters.Clear();
+            }
-                    CancelledDoAfters.Clear();
+
-                    break;
+            foreach (var (_, cancelled) in CancelledDoAfters)
+            {
+                Gui.CancelDoAfter(cancelled.ID);
+            }
+        }
+
+        public void Disable()
+        {
+            Gui?.Dispose();
+        }
+
+        public override void HandleComponentState(ComponentState? curState, ComponentState? nextState)
+        {
+            base.HandleComponentState(curState, nextState);
+            if (!(curState is DoAfterComponentState state))
+                return;
+            
+            _doAfters.Clear();
+
+            foreach (var doAfter in state.DoAfters)
+            {
+                _doAfters.Add(doAfter.ID, doAfter);
+            }
+            
+            if (Gui == null || Gui.Disposed)
+                return;
+
+            foreach (var (_, doAfter) in _doAfters)
+            {
+                Gui.AddDoAfter(doAfter);
             }
         }
 
         /// <summary>
         ///     Remove a DoAfter without showing a cancellation graphic.
         /// </summary>
-        /// <param name="doAfter"></param>
+        /// <param name="clientDoAfter"></param>
-        public void Remove(DoAfterMessage doAfter)
+        public void Remove(ClientDoAfter clientDoAfter)
         {
-            if (_doAfters.ContainsKey(doAfter.ID))
+            if (_doAfters.ContainsKey(clientDoAfter.ID))
-            {
+                _doAfters.Remove(clientDoAfter.ID);
-                _doAfters.Remove(doAfter.ID);
-            }
 
+            var found = false;
+            
             for (var i = CancelledDoAfters.Count - 1; i >= 0; i--)
             {
                 var cancelled = CancelledDoAfters[i];
 
-                if (cancelled.Message == doAfter)
+                if (cancelled.Message == clientDoAfter)
                 {
                     CancelledDoAfters.RemoveAt(i);
+                    found = true;
                     break;
                 }
             }
-            
+
-            EntitySystem.Get<DoAfterSystem>().Gui?.RemoveDoAfter(doAfter.ID);
+            if (!found)
+                _doAfters.Remove(clientDoAfter.ID);
+
+            Gui?.RemoveDoAfter(clientDoAfter.ID);
         }
 
         /// <summary>
@@ -88,15 +135,16 @@ namespace Content.Client.GameObjects.Components
             foreach (var (_, cancelled) in CancelledDoAfters)
             {
                 if (cancelled.ID == id)
-                {
                     return;
-                }
             }
 
+            if (!_doAfters.ContainsKey(id))
+                return;
+            
             var doAfterMessage = _doAfters[id];
             currentTime ??= IoCManager.Resolve<IGameTiming>().CurTime;
             CancelledDoAfters.Add((currentTime.Value, doAfterMessage));
-            EntitySystem.Get<DoAfterSystem>().Gui?.CancelDoAfter(id);
+            Gui?.CancelDoAfter(id);
         }
     }
 }

Content.Client/GameObjects/EntitySystems/DoAfter/DoAfterGui.cs

@@ -32,7 +32,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
 
         // This behavior probably shouldn't be happening; so for whatever reason the control position is set the frame after
         // I got NFI why because I don't know the UI internals
-        private bool _firstDraw = true;
+        public bool FirstDraw { get; set; }
 
         public DoAfterGui()
         {
@@ -43,21 +43,18 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             LayoutContainer.SetGrowVertical(this, LayoutContainer.GrowDirection.Begin);
         }
 
-        /// <summary>
+        protected override void Dispose(bool disposing)
-        ///     Called when the mind is detached from an entity
-        /// </summary>
-        ///     Rather than just dispose of the Gui we'll just remove its child controls and re-use the control.
-        public void Detached()
         {
+            base.Dispose(disposing);
+            if (Disposed)
+                return;
+            
             foreach (var (_, control) in _doAfterControls)
             {
                 control.Dispose();
             }
+            
             _doAfterControls.Clear();
-            foreach (var (_, control) in _doAfterBars)
-            {
-                control.Dispose();
-            }
             _doAfterBars.Clear();
             _cancelledDoAfters.Clear();
         }
@@ -66,12 +63,10 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
         ///     Add the necessary control for a DoAfter progress bar.
         /// </summary>
         /// <param name="message"></param>
-        public void AddDoAfter(DoAfterMessage message)
+        public void AddDoAfter(ClientDoAfter message)
         {
             if (_doAfterControls.ContainsKey(message.ID))
-            {
                 return;
-            }
 
             var doAfterBar = new DoAfterBar
             {
@@ -108,18 +103,16 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
         public void RemoveDoAfter(byte id)
         {
             if (!_doAfterControls.ContainsKey(id))
-            {
                 return;
-            }
 
             var control = _doAfterControls[id];
             RemoveChild(control);
             _doAfterControls.Remove(id);
             _doAfterBars.Remove(id);
+            
             if (_cancelledDoAfters.ContainsKey(id))
-            {
                 _cancelledDoAfters.Remove(id);
-            }
+            
         }
 
         /// <summary>
@@ -130,12 +123,15 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
         public void CancelDoAfter(byte id)
         {
             if (_cancelledDoAfters.ContainsKey(id))
-            {
                 return;
-            }
 
-            var control = _doAfterControls[id];
+            if (!_doAfterBars.TryGetValue(id, out var doAfterBar))
-            _doAfterBars[id].Cancelled = true;
+            {
+                doAfterBar = new DoAfterBar();
+                _doAfterBars[id] = doAfterBar;
+            }
+            
+            doAfterBar.Cancelled = true;
             _cancelledDoAfters.Add(id, _gameTiming.CurTime);
         }
 
@@ -143,28 +139,24 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
         {
             base.FrameUpdate(args);
 
-            if (AttachedEntity?.IsValid() != true || !AttachedEntity.TryGetComponent(out DoAfterComponent? doAfterComponent))
+            if (AttachedEntity?.IsValid() != true || 
+                !AttachedEntity.TryGetComponent(out DoAfterComponent? doAfterComponent))
             {
                 return;
             }
 
             var doAfters = doAfterComponent.DoAfters;
-
+            if (doAfters.Count == 0)
-            // Nothing to render so we'll hide.
-            if (doAfters.Count == 0 && _cancelledDoAfters.Count == 0)
-            {
-                _firstDraw = true;
-                Visible = false;
                 return;
-            }
 
             // Set position ready for 2nd+ frames.
             _playerPosition = _eyeManager.CoordinatesToScreen(AttachedEntity.Transform.Coordinates);
             LayoutContainer.SetPosition(this, new Vector2(_playerPosition.X - Width / 2, _playerPosition.Y - Height - 30.0f));
 
-            if (_firstDraw)
+            if (FirstDraw)
             {
-                _firstDraw = false;
+                Visible = false;
+                FirstDraw = false;
                 return;
             }
 
@@ -176,9 +168,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             foreach (var (id, cancelTime) in _cancelledDoAfters)
             {
                 if ((currentTime - cancelTime).TotalSeconds > DoAfterSystem.ExcessTime)
-                {
                     toCancel.Add(id);
-                }
             }
 
             foreach (var id in toCancel)
@@ -190,9 +180,7 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
             foreach (var (id, message) in doAfters)
             {
                 if (_cancelledDoAfters.ContainsKey(id) || !_doAfterControls.ContainsKey(id))
-                {
                     continue;
-                }
 
                 var doAfterBar = _doAfterBars[id];
                 doAfterBar.Ratio = MathF.Min(1.0f,

Content.Client/GameObjects/EntitySystems/DoAfter/DoAfterSystem.cs

@@ -1,6 +1,8 @@
 #nullable enable
+using System.Collections.Generic;
 using System.Linq;
 using Content.Client.GameObjects.Components;
+using Content.Shared.GameObjects.EntitySystems;
 using JetBrains.Annotations;
 using Robust.Client.GameObjects.EntitySystems;
 using Robust.Shared.GameObjects.Systems;
@@ -27,133 +29,132 @@ namespace Content.Client.GameObjects.EntitySystems.DoAfter
         [Dependency] private readonly IGameTiming _gameTiming = default!;
         [Dependency] private readonly IEntityManager _entityManager = default!;
 
-        /// <summary>
-        ///     Rather than checking attached player every tick we'll just store it from the message.
-        /// </summary>
-        private IEntity? _player;
-
         /// <summary>
         ///     We'll use an excess time so stuff like finishing effects can show.
         /// </summary>
         public const float ExcessTime = 0.5f;
 
-        public DoAfterGui? Gui { get; private set; }
+        // Each component in range will have its own vBox which we need to keep track of so if they go out of range or
+        // come into range it needs altering
+        private HashSet<DoAfterComponent> _knownComponents = new HashSet<DoAfterComponent>();
 
+        private IEntity? _attachedEntity;
+        
         public override void Initialize()
         {
             base.Initialize();
-            SubscribeLocalEvent<PlayerAttachSysMessage>(message => HandlePlayerAttached(message.AttachedEntity));
+            SubscribeLocalEvent<PlayerAttachSysMessage>(HandlePlayerAttached);
         }
 
-        public override void Shutdown()
+        private void HandlePlayerAttached(PlayerAttachSysMessage message)
         {
-            base.Shutdown();
+            _attachedEntity = message.AttachedEntity;
-            Gui?.Dispose();
-            Gui = null;
-            _player = null;
         }
-
+        
-        private void HandlePlayerAttached(IEntity? entity)
-        {
-            _player = entity;
-            // Setup the GUI and pass the new data to it if applicable.
-            Gui?.Detached();
-
-            if (entity == null)
-            {
-                return;
-            }
-
-            Gui ??= new DoAfterGui();
-            Gui.AttachedEntity = entity;
-
-            if (entity.TryGetComponent(out DoAfterComponent? doAfterComponent))
-            {
-                foreach (var (_, doAfter) in doAfterComponent.DoAfters)
-                {
-                    Gui.AddDoAfter(doAfter);
-                }
-            }
-        }
-
         public override void Update(float frameTime)
         {
             base.Update(frameTime);
 
             var currentTime = _gameTiming.CurTime;
-
+            var foundComps = new HashSet<DoAfterComponent>();
-            if (_player?.IsValid() != true)
+            
-            {
+            // Can't see any I guess?
+            if (_attachedEntity == null || _attachedEntity.Deleted)
                 return;
-            }
 
-            if (!_player.TryGetComponent(out DoAfterComponent? doAfterComponent))
+            foreach (var comp in ComponentManager.EntityQuery<DoAfterComponent>())
             {
-                return;
+                if (!_knownComponents.Contains(comp))
-            }
-
-            var doAfters = doAfterComponent.DoAfters.ToList();
-            if (doAfters.Count == 0)
-            {
-                return;
-            }
-
-            var userGrid = _player.Transform.Coordinates;
-
-            // Check cancellations / finishes
-            foreach (var (id, doAfter) in doAfters)
-            {
-                var elapsedTime = (currentTime - doAfter.StartTime).TotalSeconds;
-
-                // If we've passed the final time (after the excess to show completion graphic) then remove.
-                if (elapsedTime > doAfter.Delay + ExcessTime)
                 {
-                    Gui?.RemoveDoAfter(id);
+                    _knownComponents.Add(comp);
-                    doAfterComponent.Remove(doAfter);
+                }
-                    continue;
+                
+                var doAfters = comp.DoAfters.ToList();
+
+                if (doAfters.Count == 0)
+                {
+                    if (comp.Gui != null)
+                        comp.Gui.FirstDraw = true;
+                    
+                    return;
                 }
 
-                // Don't predict cancellation if it's already finished.
+                var range = (comp.Owner.Transform.WorldPosition - _attachedEntity.Transform.WorldPosition).Length + 0.01f;
-                if (elapsedTime > doAfter.Delay)
-                {
-                    continue;
-                }
 
-                // Predictions
+                if (comp.Owner != _attachedEntity && !ExamineSystemShared.InRangeUnOccluded(
-                if (doAfter.BreakOnUserMove)
+                    _attachedEntity.Transform.MapPosition,
+                    comp.Owner.Transform.MapPosition, range,
+                    entity => entity == comp.Owner || entity == _attachedEntity))
                 {
-                    if (userGrid != doAfter.UserGrid)
+                    if (comp.Gui != null)
+                        comp.Gui.FirstDraw = true;
+                    
+                    return;
+                }
+                
+                comp.Enable();
+
+                var userGrid = comp.Owner.Transform.Coordinates;
+
+                // Check cancellations / finishes
+                foreach (var (id, doAfter) in doAfters)
+                {
+                    var elapsedTime = (currentTime - doAfter.StartTime).TotalSeconds;
+
+                    // If we've passed the final time (after the excess to show completion graphic) then remove.
+                    if (elapsedTime > doAfter.Delay + ExcessTime)
                     {
-                        doAfterComponent.Cancel(id, currentTime);
+                        comp.Remove(doAfter);
-                        continue;
-                    }
-                }
-
-                if (doAfter.BreakOnTargetMove)
-                {
-                    if (!_entityManager.TryGetEntity(doAfter.TargetUid, out var targetEntity))
-                    {
-                        // Cancel if the target entity doesn't exist.
-                        doAfterComponent.Cancel(id, currentTime);
                         continue;
                     }
 
-                    if (targetEntity.Transform.Coordinates != doAfter.TargetGrid)
+                    // Don't predict cancellation if it's already finished.
-                    {
+                    if (elapsedTime > doAfter.Delay)
-                        doAfterComponent.Cancel(id, currentTime);
                         continue;
+
+                    // Predictions
+                    if (doAfter.BreakOnUserMove)
+                    {
+                        if (userGrid != doAfter.UserGrid)
+                        {
+                            comp.Cancel(id, currentTime);
+                            continue;
+                        }
+                    }
+
+                    if (doAfter.BreakOnTargetMove)
+                    {
+                        var targetEntity = _entityManager.GetEntity(doAfter.TargetUid);
+
+                        if (targetEntity.Transform.Coordinates != doAfter.TargetGrid)
+                        {
+                            comp.Cancel(id, currentTime);
+                            continue;
+                        }
                     }
                 }
+
+                var count = comp.CancelledDoAfters.Count;
+                // Remove cancelled DoAfters after ExcessTime has elapsed
+                for (var i = count - 1; i >= 0; i--)
+                {
+                    var cancelled = comp.CancelledDoAfters[i];
+                    if ((currentTime - cancelled.CancelTime).TotalSeconds > ExcessTime)
+                    {
+                        comp.Remove(cancelled.Message);
+                    }
+                }
+                
+                // Remove any components that we no longer need to track
+                foundComps.Add(comp);
             }
 
-            var count = doAfterComponent.CancelledDoAfters.Count;
+            foreach (var comp in foundComps)
-            // Remove cancelled DoAfters after ExcessTime has elapsed
-            for (var i = count - 1; i >= 0; i--)
             {
-                var cancelled = doAfterComponent.CancelledDoAfters[i];
+                if (!_knownComponents.Contains(comp))
-                if ((currentTime - cancelled.CancelTime).TotalSeconds > ExcessTime)
                 {
-                    doAfterComponent.Remove(cancelled.Message);
+                    _knownComponents.Remove(comp);
+                    comp.Disable();
                 }
             }
         }

Content.Server/GameObjects/Components/DoAfterComponent.cs

@@ -2,11 +2,7 @@
 using System.Collections.Generic;
 using Content.Server.GameObjects.EntitySystems.DoAfter;
 using Content.Shared.GameObjects.Components;
-using Robust.Server.GameObjects;
-using Robust.Server.Interfaces.GameObjects;
 using Robust.Shared.GameObjects;
-using Robust.Shared.Interfaces.GameObjects;
-using Robust.Shared.Interfaces.Network;
 
 namespace Content.Server.GameObjects.Components
 {
@@ -20,30 +16,15 @@ namespace Content.Server.GameObjects.Components
         // we'll just send them the index. Doesn't matter if it wraps around.
         private byte _runningIndex;
 
-        public override void HandleMessage(ComponentMessage message, IComponent? component)
+        public override ComponentState GetComponentState()
         {
-            base.HandleMessage(message, component);
+            var toAdd = new List<ClientDoAfter>();
-            switch (message)
+            
-            {
+            foreach (var doAfter in DoAfters)
-                case PlayerAttachedMsg _:
-                    UpdateClient();
-                    break;
-            }
-        }
-
-        // Only sending data to the relevant client (at least, other clients don't need to know about do_after for now).
-        private void UpdateClient()
-        {
-            if (!TryGetConnectedClient(out var connectedClient))
-            {
-                return;
-            }
-
-            foreach (var (doAfter, id) in _doAfters)
             {
                 // THE ALMIGHTY PYRAMID
-                var message = new DoAfterMessage(
+                var clientDoAfter = new ClientDoAfter(
-                    id,
+                    _doAfters[doAfter],
                     doAfter.UserGrid,
                     doAfter.TargetGrid,
                     doAfter.StartTime,
@@ -51,65 +32,27 @@ namespace Content.Server.GameObjects.Components
                     doAfter.EventArgs.BreakOnUserMove,
                     doAfter.EventArgs.BreakOnTargetMove,
                     doAfter.EventArgs.Target?.Uid ?? EntityUid.Invalid);
-
+                
-                SendNetworkMessage(message, connectedClient);
+                toAdd.Add(clientDoAfter);
-            }
-        }
-
-        private bool TryGetConnectedClient(out INetChannel? connectedClient)
-        {
-            connectedClient = null;
-
-            if (!Owner.TryGetComponent(out IActorComponent? actorComponent))
-            {
-                return false;
             }
 
-            connectedClient = actorComponent.playerSession.ConnectedClient;
+            return new DoAfterComponentState(toAdd);
-            if (!connectedClient.IsConnected)
-            {
-                return false;
-            }
-
-            return true;
         }
 
         public void Add(DoAfter doAfter)
         {
             _doAfters.Add(doAfter, _runningIndex);
-
-            if (TryGetConnectedClient(out var connectedClient))
-            {
-                var message = new DoAfterMessage(
-                    _runningIndex,
-                    doAfter.UserGrid,
-                    doAfter.TargetGrid,
-                    doAfter.StartTime,
-                    doAfter.EventArgs.Delay,
-                    doAfter.EventArgs.BreakOnUserMove,
-                    doAfter.EventArgs.BreakOnTargetMove,
-                    doAfter.EventArgs.Target?.Uid ?? EntityUid.Invalid);
-
-                SendNetworkMessage(message, connectedClient);
-            }
-
             _runningIndex++;
+            Dirty();
         }
 
         public void Cancelled(DoAfter doAfter)
         {
             if (!_doAfters.TryGetValue(doAfter, out var index))
-            {
                 return;
-            }
-
-            if (TryGetConnectedClient(out var connectedClient))
-            {
-                var message = new CancelledDoAfterMessage(index);
-                SendNetworkMessage(message, connectedClient);
-            }
 
             _doAfters.Remove(doAfter);
+            SendNetworkMessage(new CancelledDoAfterMessage(index));
         }
 
         /// <summary>
@@ -120,9 +63,7 @@ namespace Content.Server.GameObjects.Components
         public void Finished(DoAfter doAfter)
         {
             if (!_doAfters.ContainsKey(doAfter))
-            {
                 return;
-            }
 
             _doAfters.Remove(doAfter);
         }

Content.Shared/GameObjects/Components/SharedDoAfterComponent.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using Robust.Shared.GameObjects;
 using Robust.Shared.Map;
 using Robust.Shared.Serialization;
@@ -12,6 +13,17 @@ namespace Content.Shared.GameObjects.Components
         public override uint? NetID => ContentNetIDs.DO_AFTER;
     }
 
+    [Serializable, NetSerializable]
+    public sealed class DoAfterComponentState : ComponentState
+    {
+        public List<ClientDoAfter> DoAfters { get; }
+
+        public DoAfterComponentState(List<ClientDoAfter> doAfters) : base(ContentNetIDs.DO_AFTER)
+        {
+            DoAfters = doAfters;
+        }
+    }
+
     [Serializable, NetSerializable]
     public sealed class CancelledDoAfterMessage : ComponentMessage
     {
@@ -27,7 +39,7 @@ namespace Content.Shared.GameObjects.Components
     ///     We send a trimmed-down version of the DoAfter for the client for it to use.
     /// </summary>
     [Serializable, NetSerializable]
-    public sealed class DoAfterMessage : ComponentMessage
+    public sealed class ClientDoAfter
     {
         // To see what these do look at DoAfter and DoAfterEventArgs
         public byte ID { get; }
@@ -47,7 +59,7 @@ namespace Content.Shared.GameObjects.Components
 
         public bool BreakOnTargetMove { get; }
 
-        public DoAfterMessage(byte id, EntityCoordinates userGrid, EntityCoordinates targetGrid, TimeSpan startTime, float delay, bool breakOnUserMove, bool breakOnTargetMove, EntityUid targetUid = default)
+        public ClientDoAfter(byte id, EntityCoordinates userGrid, EntityCoordinates targetGrid, TimeSpan startTime, float delay, bool breakOnUserMove, bool breakOnTargetMove, EntityUid targetUid = default)
         {
             ID = id;
             UserGrid = userGrid;