AI sleeping (#1708)

* AI sleeping

AI no longer update when dead.

* It was easier to merge master and re-apply it.

* Update AiControllerComponent.cs

Co-authored-by: Metal Gear Sloth <metalgearsloth@gmail.com>
Co-authored-by: DrSmugleaf <DrSmugleaf@users.noreply.github.com>

Content.Server/AI/Utility/AiLogic/UtilityAI.cs

@@ -6,10 +6,13 @@ using Content.Server.AI.Utility.Actions;
 using Content.Server.AI.Utility.BehaviorSets;
 using Content.Server.AI.WorldState;
 using Content.Server.AI.WorldState.States.Utility;
+using Content.Server.GameObjects.EntitySystems.AI;
 using Content.Server.GameObjects.EntitySystems.AI.LoadBalancer;
 using Content.Server.GameObjects.EntitySystems.JobQueues;
 using Content.Shared.GameObjects.Components.Damage;
 using Robust.Server.AI;
+using Robust.Shared.GameObjects;
+using Robust.Shared.GameObjects.Systems;
 using Robust.Shared.Interfaces.GameObjects;
 using Robust.Shared.IoC;
 using Robust.Shared.Log;
@@ -62,6 +65,13 @@ namespace Content.Server.AI.Utility.AiLogic
             {
                 SortActions();
             }
+
+            if (BehaviorSets.Count == 1 && !EntitySystem.Get<AiSystem>().IsAwake(this))
+            {
+                IoCManager.Resolve<IEntityManager>()
+                    .EventBus
+                    .RaiseEvent(EventSource.Local, new SleepAiMessage(this, false));
+            }
         }
 
         public void RemoveBehaviorSet(Type behaviorSet)
@@ -73,6 +83,13 @@ namespace Content.Server.AI.Utility.AiLogic
                 BehaviorSets.Remove(behaviorSet);
                 SortActions();
             }
+
+            if (BehaviorSets.Count == 0)
+            {
+                IoCManager.Resolve<IEntityManager>()
+                    .EventBus
+                    .RaiseEvent(EventSource.Local, new SleepAiMessage(this, true));
+            }
         }
 
         /// <summary>
@@ -133,7 +150,23 @@ namespace Content.Server.AI.Utility.AiLogic
 
         private void DeathHandle(HealthChangedEventArgs eventArgs)
         {
+            var oldDeadState = _isDead;
             _isDead = eventArgs.Damageable.CurrentDamageState == DamageState.Dead || eventArgs.Damageable.CurrentDamageState == DamageState.Critical;
+
+            if (oldDeadState != _isDead)
+            {
+                var entityManager = IoCManager.Resolve<IEntityManager>();
+                
+                switch (_isDead)
+                {
+                    case true:
+                        entityManager.EventBus.RaiseEvent(EventSource.Local, new SleepAiMessage(this, true));
+                        break;
+                    case false:
+                        entityManager.EventBus.RaiseEvent(EventSource.Local, new SleepAiMessage(this, false));
+                        break;
+                }
+            }
         }
 
         private void ReceivedAction()
@@ -167,16 +200,6 @@ namespace Content.Server.AI.Utility.AiLogic
 
         public override void Update(float frameTime)
         {
-            // If we can't do anything then there's no point thinking
-            if (_isDead || BehaviorSets.Count == 0)
-            {
-                _actionCancellation?.Cancel();
-                _blackboard.GetState<LastUtilityScoreState>().SetValue(0.0f);
-                CurrentAction?.Shutdown();
-                CurrentAction = null;
-                return;
-            }
-
             // If we asked for a new action we don't want to dump the existing one.
             if (_actionRequest != null)
             {

Content.Server/GameObjects/Components/Movement/AiControllerComponent.cs

@@ -1,7 +1,9 @@
-using Content.Shared.GameObjects.Components.Movement;
+using Content.Server.GameObjects.EntitySystems.AI;
+using Content.Shared.GameObjects.Components.Movement;
 using Robust.Server.AI;
 using Robust.Shared.GameObjects;
 using Robust.Shared.GameObjects.Components;
+using Robust.Shared.GameObjects.Systems;
 using Robust.Shared.Map;
 using Robust.Shared.Maths;
 using Robust.Shared.Serialization;
@@ -45,6 +47,8 @@ namespace Content.Server.GameObjects.Components.Movement
             // This component requires a collidable component.
             if (!Owner.HasComponent<ICollidableComponent>())
                 Owner.AddComponent<CollidableComponent>();
+            
+            EntitySystem.Get<AiSystem>().ProcessorInitialize(this);
         }
 
         /// <inheritdoc />

Content.Server/GameObjects/EntitySystems/AI/AiSystem.cs

@@ -1,4 +1,5 @@
-using System;
+#nullable enable
+using System;
 using System.Collections.Generic;
 using Content.Server.GameObjects.Components.Movement;
 using Content.Shared.GameObjects.Components.Movement;
@@ -19,23 +20,31 @@ namespace Content.Server.GameObjects.EntitySystems.AI
     [UsedImplicitly]
     internal class AiSystem : EntitySystem
     {
-#pragma warning disable 649
+        [Dependency] private readonly IDynamicTypeFactory _typeFactory = default!;
-        [Dependency] private readonly IPauseManager _pauseManager;
+        [Dependency] private readonly IReflectionManager _reflectionManager = default!;
-        [Dependency] private readonly IDynamicTypeFactory _typeFactory;
-        [Dependency] private readonly IReflectionManager _reflectionManager;
-#pragma warning restore 649
 
         private readonly Dictionary<string, Type> _processorTypes = new Dictionary<string, Type>();
         
+        /// <summary>
+        ///     To avoid iterating over dead AI continuously they can wake and sleep themselves when necessary.
+        /// </summary>
+        private readonly HashSet<AiLogicProcessor> _awakeAi = new HashSet<AiLogicProcessor>();
+        
+        // To avoid modifying awakeAi while iterating over it.
+        private readonly List<SleepAiMessage> _queuedSleepMessages = new List<SleepAiMessage>();
+
+        public bool IsAwake(AiLogicProcessor processor) => _awakeAi.Contains(processor);
+        
         /// <inheritdoc />
         public override void Initialize()
         {
             base.Initialize();
+            SubscribeLocalEvent<SleepAiMessage>(HandleAiSleep);
 
             var processors = _reflectionManager.GetAllChildren<AiLogicProcessor>();
             foreach (var processor in processors)
             {
-                var att = (AiLogicProcessorAttribute)Attribute.GetCustomAttribute(processor, typeof(AiLogicProcessorAttribute));
+                var att = (AiLogicProcessorAttribute) Attribute.GetCustomAttribute(processor, typeof(AiLogicProcessorAttribute))!;
                 // Tests should pick this up
                 DebugTools.AssertNotNull(att);
                 _processorTypes.Add(att.SerializeName, processor);
@@ -45,23 +54,35 @@ namespace Content.Server.GameObjects.EntitySystems.AI
         /// <inheritdoc />
         public override void Update(float frameTime)
         {
-            foreach (var comp in ComponentManager.EntityQuery<AiControllerComponent>())
+            foreach (var message in _queuedSleepMessages)
             {
-                if (_pauseManager.IsEntityPaused(comp.Owner))
+                switch (message.Sleep)
                 {
-                    continue;
+                    case true:
+                        _awakeAi.Remove(message.Processor);
+                        break;
+                    case false:
+                        _awakeAi.Add(message.Processor);
+                        break;
                 }
+            }
             
-                ProcessorInitialize(comp);
+            _queuedSleepMessages.Clear();
-
-                var processor = comp.Processor;
             
+            foreach (var processor in _awakeAi)
+            {
                 processor.Update(frameTime);
             }
         }
 
+        private void HandleAiSleep(SleepAiMessage message)
+        {
+            _queuedSleepMessages.Add(message);
+        }
+
         /// <summary>
-        /// Will start up the controller's processor if not already done so
+        ///     Will start up the controller's processor if not already done so.
+        ///     Also add them to the awakeAi for updates.
         /// </summary>
         /// <param name="controller"></param>
         public void ProcessorInitialize(AiControllerComponent controller)
@@ -70,6 +91,7 @@ namespace Content.Server.GameObjects.EntitySystems.AI
             controller.Processor = CreateProcessor(controller.LogicName);
             controller.Processor.SelfEntity = controller.Owner;
             controller.Processor.Setup();
+            _awakeAi.Add(controller.Processor);
         }
 
         private AiLogicProcessor CreateProcessor(string name)
@@ -94,7 +116,7 @@ namespace Content.Server.GameObjects.EntitySystems.AI
                                 + "\n    processorId: Class that inherits AiLogicProcessor and has an AiLogicProcessor attribute."
                                 + "\n    entityID: Uid of entity to add the AiControllerComponent to. Open its VV menu to find this.";
 
-            public void Execute(IConsoleShell shell, IPlayerSession player, string[] args)
+            public void Execute(IConsoleShell shell, IPlayerSession? player, string[] args)
             {
                 if(args.Length != 2)
                 {

Content.Server/GameObjects/EntitySystems/AI/SleepAiMessage.cs

@@ -0,0 +1,24 @@
+using Robust.Server.AI;
+using Robust.Shared.GameObjects;
+
+namespace Content.Server.GameObjects.EntitySystems.AI
+{
+    /// <summary>
+    ///     Indicates whether an AI should be updated by the AiSystem or not.
+    ///     Useful to sleep AI when they die or otherwise should be inactive.
+    /// </summary>
+    internal sealed class SleepAiMessage : EntitySystemMessage
+    {
+        /// <summary>
+        ///     Sleep or awake.
+        /// </summary>
+        public bool Sleep { get; }
+        public AiLogicProcessor Processor { get; }
+        
+        public SleepAiMessage(AiLogicProcessor processor, bool sleep)
+        {
+            Processor = processor;
+            Sleep = sleep;
+        }
+    }
+}