[development] fixed issue with dangling budget pointers if the budget tracker is torn down and rebuilt (#6801)

Signed-off-by: AMZN-ScottR 24445312+AMZN-ScottR@users.noreply.github.com

Code/Framework/AzCore/AzCore/Component/ComponentApplication.cpp

@@ -565,10 +565,6 @@ namespace AZ
         m_entityActivatedEvent.DisconnectAllHandlers();
         m_entityDeactivatedEvent.DisconnectAllHandlers();
 
-#if !defined(_RELEASE)
-        m_budgetTracker.Reset();
-#endif
-
         DestroyAllocator();
     }
 
@@ -758,6 +754,12 @@ namespace AZ
         static_cast<SettingsRegistryImpl*>(m_settingsRegistry.get())->ClearNotifiers();
         static_cast<SettingsRegistryImpl*>(m_settingsRegistry.get())->ClearMergeEvents();
 
+#if !defined(_RELEASE)
+        // the budget tracker must be cleaned up prior to module unloading to ensure
+        // budgets initialized cross boundary are freed properly
+        m_budgetTracker.Reset();
+#endif
+
         // Uninit and unload any dynamic modules.
         m_moduleManager->UnloadModules();
         ComponentApplicationLifecycle::SignalEvent(*m_settingsRegistry, "GemsUnloaded", R"({})");

Code/Framework/AzCore/AzCore/Debug/Budget.h

@@ -62,12 +62,16 @@ namespace AZ::Debug
 //
 // Anywhere the budget is used, the budget must be declared (either in a header or in the source file itself)
 // AZ_DECLARE_BUDGET(AzCore);
-#define AZ_DEFINE_BUDGET(name)                                                                                                             \
-    ::AZ::Debug::Budget* AZ_BUDGET_GETTER(name)()                                                                                          \
-    {                                                                                                                                      \
-        constexpr static uint32_t crc = AZ_CRC_CE(#name);                                                                                  \
-        static ::AZ::Debug::Budget* budget = ::AZ::Debug::BudgetTracker::GetBudgetFromEnvironment(#name, crc);                             \
-        return budget;                                                                                                                     \
+#define AZ_DEFINE_BUDGET(name)                                                          \
+    ::AZ::Debug::Budget* AZ_BUDGET_GETTER(name)()                                       \
+    {                                                                                   \
+        static ::AZ::Debug::Budget* budget = nullptr;                                   \
+        if (budget == nullptr)                                                          \
+        {                                                                               \
+            constexpr static uint32_t crc = AZ_CRC_CE(#name);                           \
+            ::AZ::Debug::BudgetTracker::GetBudgetFromEnvironment(budget, #name, crc);   \
+        }                                                                               \
+        return budget;                                                                  \
     }
 #endif
 

Code/Framework/AzCore/AzCore/Debug/BudgetTracker.cpp

@@ -13,23 +13,24 @@
 #include <AzCore/Interface/Interface.h>
 #include <AzCore/Memory/Memory.h>
 #include <AzCore/std/containers/unordered_map.h>
+#include <AzCore/std/containers/unordered_set.h>
 #include <AzCore/std/parallel/scoped_lock.h>
 
 namespace AZ::Debug
 {
     struct BudgetTracker::BudgetTrackerImpl
     {
-        AZStd::unordered_map<const char*, Budget> m_budgets;
+        AZStd::unordered_map<AZStd::string_view, Budget> m_budgets;
+        AZStd::unordered_set<Budget**> m_externalBudgetRefs;
     };
 
-    Budget* BudgetTracker::GetBudgetFromEnvironment(const char* budgetName, uint32_t crc)
+    void BudgetTracker::GetBudgetFromEnvironment(Budget*& extBudgetRef, const char* budgetName, uint32_t crc)
     {
         BudgetTracker* tracker = Interface<BudgetTracker>::Get();
         if (tracker)
         {
-            return &tracker->GetBudget(budgetName, crc);
+            tracker->GetBudget(extBudgetRef, budgetName, crc);
         }
-        return nullptr;
     }
 
     BudgetTracker::~BudgetTracker()
@@ -54,17 +55,24 @@ namespace AZ::Debug
         if (m_impl)
         {
             Interface<BudgetTracker>::Unregister(this);
+
+            for (auto budgetRef : m_impl->m_externalBudgetRefs)
+            {
+                *budgetRef = nullptr;
+            }
+
             delete m_impl;
             m_impl = nullptr;
         }
     }
 
-    Budget& BudgetTracker::GetBudget(const char* budgetName, uint32_t crc)
+    void BudgetTracker::GetBudget(Budget*& extBudgetRef, const char* budgetName, uint32_t crc)
     {
         AZStd::scoped_lock lock{ m_mutex };
 
-        auto it = m_impl->m_budgets.try_emplace(budgetName, budgetName, crc).first;
+        m_impl->m_externalBudgetRefs.insert(&extBudgetRef);
 
-        return it->second;
+        auto iter = m_impl->m_budgets.try_emplace(budgetName, budgetName, crc).first;
+        extBudgetRef = &iter->second;
     }
 } // namespace AZ::Debug

Code/Framework/AzCore/AzCore/Debug/BudgetTracker.h

@@ -20,7 +20,7 @@ namespace AZ::Debug
     {
     public:
         AZ_TYPE_INFO(BudgetTracker, "{E14A746D-BFFE-4C02-90FB-4699B79864A5}");
-        static Budget* GetBudgetFromEnvironment(const char* budgetName, uint32_t crc);
+        static void GetBudgetFromEnvironment(Budget*& extBudgetRef, const char* budgetName, uint32_t crc);
 
         ~BudgetTracker();
 
@@ -28,7 +28,7 @@ namespace AZ::Debug
         bool Init();
         void Reset();
 
-        Budget& GetBudget(const char* budgetName, uint32_t crc);
+        void GetBudget(Budget*& extBudgetRef, const char* budgetName, uint32_t crc);
 
     private:
         struct BudgetTrackerImpl;