Core/AI: Add some checks to catch if the AI changes mid-update for players. Tag #22731 #22629.

2 files changed, 11 insertions, 10 deletions

diff --git a/src/server/game/Entities/Unit/Unit.cpp b/src/server/game/Entities/Unit/Unit.cpp
index 58b5af0b9ef..4262ced1fbb 100644
--- a/src/server/game/Entities/Unit/Unit.cpp
+++ b/src/server/game/Entities/Unit/Unit.cpp
@@ -292,7 +292,7 @@ Unit::Unit(bool isWorldObject) :
     m_procDeep(0), m_transformSpell(0), m_removedAurasCount(0), m_charmer(nullptr), m_charmed(nullptr),
     i_motionMaster(new MotionMaster(this)), m_regenTimer(0), m_vehicle(nullptr), m_vehicleKit(nullptr),
     m_unitTypeMask(UNIT_MASK_NONE), m_Diminishing(), m_combatManager(this), m_threatManager(this),
-    m_comboTarget(nullptr), m_comboPoints(0), m_spellHistory(new SpellHistory(this))
+    m_aiLocked(false), m_comboTarget(nullptr), m_comboPoints(0), m_spellHistory(new SpellHistory(this))
 {
     m_objectType |= TYPEMASK_UNIT;
     m_objectTypeId = TYPEID_UNIT;
@@ -9401,24 +9401,25 @@ uint32 Unit::GetCreatePowers(Powers power) const
     return 0;
 }
 
-void Unit::AIUpdateTick(uint32 diff, bool /*force*/)
+void Unit::AIUpdateTick(uint32 diff)
 {
-    if (!diff) // some places call with diff = 0, which does nothing (for now), see PR #22296
-        return;
     if (UnitAI* ai = GetAI())
+    {
+        m_aiLocked = true;
         ai->UpdateAI(diff);
+        m_aiLocked = false;
+    }
 }
 
 void Unit::SetAI(UnitAI* newAI)
 {
-    if (i_AI)
-        AIUpdateTick(0, true); // old AI gets a final tick if enabled
+    ASSERT(!m_aiLocked, "Attempt to replace AI during AI update tick");
     i_AI.reset(newAI);
-    AIUpdateTick(0, true); // new AI gets its initial tick
 }
 
 void Unit::ScheduleAIChange()
 {
+    ASSERT(!m_aiLocked, "Attempt to schedule AI change during AI update tick");
     bool const charmed = IsCharmed();
     // if charm is applied, we can't have disabled AI already, and vice versa
     if (charmed)
@@ -9434,9 +9435,9 @@ void Unit::ScheduleAIChange()
 
 void Unit::RestoreDisabledAI()
 {
+    ASSERT(!m_aiLocked, "Attempt to restore AI during UpdateAI tick");
     ASSERT((GetTypeId() == TYPEID_PLAYER) || i_disabledAI, "Attempt to restore disabled AI on creature without disabled AI");
     i_AI = std::move(i_disabledAI);
-    AIUpdateTick(0, true);
 }
 
 void Unit::AddToWorld()
@@ -9558,7 +9559,6 @@ void Unit::UpdateCharmAI()
         ASSERT(newAI);
         i_AI.reset(newAI);
         newAI->OnCharmed(true);
-        AIUpdateTick(0, true);
     }
     else
     {
diff --git a/src/server/game/Entities/Unit/Unit.h b/src/server/game/Entities/Unit/Unit.h
index 7d3288e092e..67b9b886b49 100644
--- a/src/server/game/Entities/Unit/Unit.h
+++ b/src/server/game/Entities/Unit/Unit.h
@@ -765,7 +765,7 @@ class TC_GAME_API Unit : public WorldObject
         virtual ~Unit();
 
         bool IsAIEnabled() const { return (i_AI != nullptr); }
-        void AIUpdateTick(uint32 diff, bool force = false);
+        void AIUpdateTick(uint32 diff);
         UnitAI* GetAI() const { return i_AI.get(); }
         void SetAI(UnitAI* newAI);
         void ScheduleAIChange();
@@ -1790,6 +1790,7 @@ class TC_GAME_API Unit : public WorldObject
         void UpdateCharmAI();
         void RestoreDisabledAI();
         std::unique_ptr<UnitAI> i_AI, i_disabledAI;
+        bool m_aiLocked;
 
         std::unordered_set<AbstractFollower*> m_followingMe;