PythonProxyNotificationHandler::OnEventGenericHook acquires the Python GIL before executing (#3904)

- Protected python execution of OnEventGenericHook by trying to lock the GIL and show a descriptive error when it was not possible to lock.
- Improved mechanism to lock python mutex and GIL in PhytonSystemComponent. Acquiring/releasing GIL once per thread.
- Added unit test to verify errors are fired when trying to execute OnEventGenericHook from another thread (as it should not able to acquire the GIL).
- Improved python threading tests to actually using python buses.

Signed-off-by: moraaar moraaar@amazon.com

Code/Framework/AzToolsFramework/AzToolsFramework/API/EditorPythonConsoleBus.h

@@ -70,8 +70,12 @@ namespace AzToolsFramework
         //! Determines if the caller needs to wait for the Python VM to initialize (non-main thread only)
         virtual void WaitForInitialization() {}
 
-        //! Acquires the Python global interpreter lock (GIL) and executed the callback
+        //! Acquires the Python global interpreter lock (GIL) and execute the callback
         virtual void ExecuteWithLock(AZStd::function<void()> executionCallback) = 0;
+
+        //! Tries to acquire the Python global interpreter lock (GIL) and execute the callback.
+        //! @return Whether it was able to lock the mutex or not.
+        virtual bool TryExecuteWithLock(AZStd::function<void()> executionCallback) = 0;
     };
 
     //! A bus to handle post notifications to the console views of Python output

Gems/EditorPythonBindings/Code/Source/PythonProxyBus.cpp

@@ -20,6 +20,7 @@
 #include <AzCore/RTTI/AttributeReader.h>
 #include <AzCore/std/optional.h>
 #include <AzFramework/StringFunc/StringFunc.h>
+#include <AzToolsFramework/API/EditorPythonConsoleBus.h>
 
 namespace EditorPythonBindings
 {
@@ -261,19 +262,39 @@ namespace EditorPythonBindings
 
             static void OnEventGenericHook(void* userData, const char* eventName, int eventIndex, AZ::BehaviorValueParameter* result, int numParameters, AZ::BehaviorValueParameter* parameters)
             {
-                reinterpret_cast<PythonProxyNotificationHandler*>(userData)->OnEventGenericHook(eventName, eventIndex, result, numParameters, parameters);
+                auto editorPythonEventsInterface = AZ::Interface<AzToolsFramework::EditorPythonEventsInterface>::Get();
+                if (!editorPythonEventsInterface)
+                {
+                    return;
                 }
 
-            void OnEventGenericHook(const char* eventName, [[maybe_unused]] int eventIndex, AZ::BehaviorValueParameter* result, int numParameters, AZ::BehaviorValueParameter* parameters)
-            {
                 // find the callback for the event
-                const auto& callbackEntry = m_callbackMap.find(eventName);
+                auto* handler = reinterpret_cast<PythonProxyNotificationHandler*>(userData);
-                if (callbackEntry == m_callbackMap.end())
+                const auto& callbackEntry = handler->m_callbackMap.find(eventName);
+                if (callbackEntry == handler->m_callbackMap.end())
                 {
                     return;
                 }
-                pybind11::function callback = callbackEntry->second;
 
+                // This function can reach from multiple threads, which means OnEventGenericHook
+                // will require to acquire the Python GIL, make sure it tries to lock it using TryExecuteWithLock.
+                [[maybe_unused]] const bool executed = editorPythonEventsInterface->TryExecuteWithLock(
+                    [handler, eventName, callback = callbackEntry->second, eventIndex, result, numParameters, parameters]()
+                    {
+                        handler->OnEventGenericHook(eventName, callback, eventIndex, result, numParameters, parameters);
+                    });
+
+                AZ_Error("python", executed,
+                    "Ebus(%s) event(%s) could not be executed because it could not acquire the Python GIL. "
+                    "This occurs when there is already another thread executing python, which has the GIL locked, "
+                    "making it not possible for this thread to callback python at the same time. "
+                    "This is a limitation of python interpreter. Python scripts executions and event callbacks "
+                    "from EBuses need be designed to avoid this scenario.",
+                    handler->m_ebus->m_name.c_str(), eventName);
+            }
+
+            void OnEventGenericHook(const char* eventName, pybind11::function callback, [[maybe_unused]] int eventIndex, AZ::BehaviorValueParameter* result, int numParameters, AZ::BehaviorValueParameter* parameters)
+            {
                 // build the parameters to send to callback
                 Convert::StackVariableAllocator stackVariableAllocator;
                 pybind11::tuple pythonParamters(numParameters);

Gems/EditorPythonBindings/Code/Source/PythonSystemComponent.cpp

@@ -258,6 +258,85 @@ namespace EditorPythonBindings
         void Finalize() override {}
     };
 
+    // Manages the acquisition and release of the Python GIL (Global Interpreter Lock).
+    // Used by PythonSystemComponent to lock the GIL when executing python.
+    class PythonSystemComponent::PythonGILScopedLock final
+    {
+    public:
+        PythonGILScopedLock(AZStd::recursive_mutex& lock, int& lockRecursiveCounter, bool tryLock = false);
+        ~PythonGILScopedLock();
+
+        bool IsLocked() const;
+
+    protected:
+        void Lock(bool tryLock);
+        void Unlock();
+
+        AZStd::recursive_mutex& m_lock;
+        int& m_lockRecursiveCounter;
+        bool m_locked = false;
+        AZStd::unique_ptr<pybind11::gil_scoped_release> m_releaseGIL;
+        AZStd::unique_ptr<pybind11::gil_scoped_acquire> m_acquireGIL;
+    };
+
+    PythonSystemComponent::PythonGILScopedLock::PythonGILScopedLock(AZStd::recursive_mutex& lock, int& lockRecursiveCounter, bool tryLock)
+        : m_lock(lock)
+        , m_lockRecursiveCounter(lockRecursiveCounter)
+    {
+        Lock(tryLock);
+    }
+
+    PythonSystemComponent::PythonGILScopedLock::~PythonGILScopedLock()
+    {
+        Unlock();
+    }
+
+    bool PythonSystemComponent::PythonGILScopedLock::IsLocked() const
+    {
+        return m_locked;
+    }
+
+    void PythonSystemComponent::PythonGILScopedLock::Lock(bool tryLock)
+    {
+        if (tryLock)
+        {
+            if (!m_lock.try_lock())
+            {
+                return;
+            }
+        }
+        else
+        {
+            m_lock.lock();
+        }
+
+        m_locked = true;
+        m_lockRecursiveCounter++;
+
+        // Only Acquire the GIL when there is no recursion. If there is
+        // recursion that means it's the same thread (because the mutex was able
+        // to be locked) and therefore it's already got the GIL acquired.
+        if (m_lockRecursiveCounter == 1)
+        {
+            m_releaseGIL = AZStd::make_unique<pybind11::gil_scoped_release>();
+            m_acquireGIL = AZStd::make_unique<pybind11::gil_scoped_acquire>();
+        }
+    }
+
+    void PythonSystemComponent::PythonGILScopedLock::Unlock()
+    {
+        if (!m_locked)
+        {
+            return;
+        }
+
+        m_acquireGIL.reset();
+        m_releaseGIL.reset();
+        m_lockRecursiveCounter--;
+        m_locked = false;
+        m_lock.unlock();
+    }
+
     void PythonSystemComponent::Reflect(AZ::ReflectContext* context)
     {
         if (AZ::SerializeContext* serialize = azrtti_cast<AZ::SerializeContext*>(context))
@@ -296,9 +375,9 @@ namespace EditorPythonBindings
 
     void PythonSystemComponent::Deactivate()
     {
+        StopPython(true);
         AzToolsFramework::EditorPythonRunnerRequestBus::Handler::BusDisconnect();
         AZ::Interface<AzToolsFramework::EditorPythonEventsInterface>::Unregister(this);
-        StopPython(true);
     }
 
     bool PythonSystemComponent::StartPython([[maybe_unused]] bool silenceWarnings)
@@ -354,7 +433,6 @@ namespace EditorPythonBindings
 
         bool result = false;
         EditorPythonBindingsNotificationBus::Broadcast(&EditorPythonBindingsNotificationBus::Events::OnPreFinalize);
-        AzToolsFramework::EditorPythonRunnerRequestBus::Handler::BusDisconnect();
 
         result = StopPythonInterpreter();
         EditorPythonBindingsNotificationBus::Broadcast(&EditorPythonBindingsNotificationBus::Events::OnPostFinalize);
@@ -374,10 +452,19 @@ namespace EditorPythonBindings
 
     void PythonSystemComponent::ExecuteWithLock(AZStd::function<void()> executionCallback)
     {
-        AZStd::lock_guard<decltype(m_lock)> lock(m_lock);
+        PythonGILScopedLock lock(m_lock, m_lockRecursiveCounter);
-        pybind11::gil_scoped_release release;
+        executionCallback();
-        pybind11::gil_scoped_acquire acquire;
+    }
+
+    bool PythonSystemComponent::TryExecuteWithLock(AZStd::function<void()> executionCallback)
+    {
+        PythonGILScopedLock lock(m_lock, m_lockRecursiveCounter, true /*tryLock*/);
+        if (lock.IsLocked())
+        {
             executionCallback();
+            return true;
+        }
+        return false;
     }
 
     void PythonSystemComponent::DiscoverPythonPaths(PythonPathStack& pythonPathStack)
@@ -548,8 +635,7 @@ namespace EditorPythonBindings
             RedirectOutput::Intialize(PyImport_ImportModule("azlmbr_redirect"));
 
             // Acquire GIL before calling Python code
-            AZStd::lock_guard<decltype(m_lock)> lock(m_lock);
+            PythonGILScopedLock lock(m_lock, m_lockRecursiveCounter);
-            pybind11::gil_scoped_acquire acquire;
 
             if (EditorPythonBindings::PythonSymbolEventBus::GetTotalNumOfEventHandlers() == 0)
             {
@@ -622,8 +708,7 @@ namespace EditorPythonBindings
                 &AzToolsFramework::EditorPythonScriptNotificationsBus::Events::OnStartExecuteByString, script);
 
             // Acquire GIL before calling Python code
-            AZStd::lock_guard<decltype(m_lock)> lock(m_lock);
+            PythonGILScopedLock lock(m_lock, m_lockRecursiveCounter);
-            pybind11::gil_scoped_acquire acquire;
 
             // Acquire scope for __main__ for executing our script
             pybind11::object scope = pybind11::module::import("__main__").attr("__dict__");
@@ -741,8 +826,7 @@ namespace EditorPythonBindings
         try
         {
             // Acquire GIL before calling Python code
-            AZStd::lock_guard<decltype(m_lock)> lock(m_lock);
+            PythonGILScopedLock lock(m_lock, m_lockRecursiveCounter);
-            pybind11::gil_scoped_acquire acquire;
 
             // Create standard "argc" / "argv" command-line parameters to pass in to the Python script via sys.argv.
             // argc = number of parameters.  This will always be at least 1, since the first parameter is the script name.

Gems/EditorPythonBindings/Code/Source/PythonSystemComponent.h

@@ -48,6 +48,7 @@ namespace EditorPythonBindings
         bool IsPythonActive() override;
         void WaitForInitialization() override;
         void ExecuteWithLock(AZStd::function<void()> executionCallback) override;
+        bool TryExecuteWithLock(AZStd::function<void()> executionCallback) override;
         ////////////////////////////////////////////////////////////////////////
 
         ////////////////////////////////////////////////////////////////////////
@@ -60,11 +61,13 @@ namespace EditorPythonBindings
         
     private:
         class SymbolLogHelper;
+        class PythonGILScopedLock;
 
         // handle multiple Python initializers and threads
         AZStd::atomic_int m_initalizeWaiterCount {0};
         AZStd::semaphore m_initalizeWaiter;
         AZStd::recursive_mutex m_lock;
+        int m_lockRecursiveCounter = 0;
         AZStd::shared_ptr<SymbolLogHelper> m_symbolLogHelper;
     
         enum class Result

Gems/EditorPythonBindings/Code/Tests/EditorPythonBindingsTest.cpp

@@ -65,7 +65,7 @@ namespace UnitTest
         {
             // clearing up memory
             m_notificationSink = EditorPythonBindingsNotificationBusSink();
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
 
             // shutdown time!
             PythonTestingFixture::TearDown();
@@ -333,7 +333,7 @@ sys.version
         {
             // clearing up memory
             m_notificationSink = EditorPythonBindingsNotificationBusSink();
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
 
             // shutdown time!
             PythonTestingFixture::TearDown();

Gems/EditorPythonBindings/Code/Tests/PythonAssetTypesTests.cpp

@@ -366,7 +366,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
     };

Gems/EditorPythonBindings/Code/Tests/PythonAssociativeTests.cpp

@@ -95,7 +95,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
     };

Gems/EditorPythonBindings/Code/Tests/PythonContainerAnyTests.cpp

@@ -227,7 +227,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
     };

Gems/EditorPythonBindings/Code/Tests/PythonDictionaryTests.cpp

@@ -107,7 +107,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
     };

Gems/EditorPythonBindings/Code/Tests/PythonGlobalsTests.cpp

@@ -175,7 +175,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
 

Gems/EditorPythonBindings/Code/Tests/PythonPairTests.cpp

@@ -121,7 +121,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
     };

Gems/EditorPythonBindings/Code/Tests/PythonProxyBusTests.cpp

@@ -207,11 +207,13 @@ namespace UnitTest
         : AZ::EBusTraits
     {
         static const AZ::EBusAddressPolicy AddressPolicy = AZ::EBusAddressPolicy::Single;
+        using MutexType = AZStd::mutex;
         virtual ~PythonTestSingleAddressNotifications() = default;
         virtual void OnPing(AZ::u64 count) = 0;
         virtual void OnPong(AZ::u64 count) = 0;
         virtual void MultipleInputs(AZ::u64 one, AZ::s8 two, AZStd::string_view three) = 0;
         virtual AZStd::string OnAddFish(AZStd::string_view value) = 0;
+        virtual void OnFire() = 0;
     };
     using PythonTestSingleAddressNotificationBus = AZ::EBus<PythonTestSingleAddressNotifications>;
 
@@ -220,7 +222,7 @@ namespace UnitTest
         , public AZ::BehaviorEBusHandler
     {
         AZ_EBUS_BEHAVIOR_BINDER(PythonTestNotificationHandler, "{97052D15-A4E8-461B-B065-91D16E31C4F7}", AZ::SystemAllocator, 
-            OnPing, OnPong, MultipleInputs, OnAddFish);
+            OnPing, OnPong, MultipleInputs, OnAddFish, OnFire);
 
         virtual ~PythonTestNotificationHandler() = default;
 
@@ -246,6 +248,11 @@ namespace UnitTest
             return result;
         }
 
+        void OnFire() override
+        {
+            Call(FN_OnFire);
+        }
+
         static AZ::u64 s_pongCount;
         static AZ::u64 s_pingCount;
 
@@ -270,6 +277,25 @@ namespace UnitTest
             return result;
         }
 
+        static void DoFire()
+        {
+            PythonTestSingleAddressNotificationBus::Broadcast(&PythonTestSingleAddressNotificationBus::Events::OnFire);
+        }
+
+        static void DoFiresInParallel(int value)
+        {
+            AZStd::vector<AZStd::thread> threads;
+            threads.reserve(value);
+            for (size_t i = 0; i < value; ++i)
+            {
+                threads.emplace_back(&DoFire);
+            }
+            for (AZStd::thread& thread : threads)
+            {
+                thread.join();
+            }
+        }
+
         static void Reset()
         {
             s_pingCount = 0;
@@ -288,6 +314,7 @@ namespace UnitTest
                     ->Event("on_pong", &PythonTestSingleAddressNotificationBus::Events::OnPong)
                     ->Event("MultipleInputs", &PythonTestSingleAddressNotificationBus::Events::MultipleInputs)
                     ->Event("OnAddFish", &PythonTestSingleAddressNotificationBus::Events::OnAddFish)
+                    ->Event("OnFire", &PythonTestSingleAddressNotificationBus::Events::OnFire)
                     ;
 
                 // for testing from Python to send out the events
@@ -297,6 +324,8 @@ namespace UnitTest
                     ->Method("do_ping", &PythonTestNotificationHandler::DoPing)
                     ->Method("do_pong", &PythonTestNotificationHandler::DoPong)
                     ->Method("do_add_fish", &PythonTestNotificationHandler::DoAddFish)
+                    ->Method("do_fire", &PythonTestNotificationHandler::DoFire)
+                    ->Method("do_fires_in_parallel", &PythonTestNotificationHandler::DoFiresInParallel)
                     ;
             }
         }
@@ -358,7 +387,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
     };
@@ -809,6 +838,65 @@ namespace UnitTest
         EXPECT_EQ(2, m_testSink.m_evaluationMap[static_cast<int>(LogTypes::AtAddress_Match)]);
     }
 
+    TEST_F(PythonBusProxyTests, SingleAddressNotifications_InParallel_Errors)
+    {
+        PythonTestNotificationHandler pythonTestNotificationHandler;
+        pythonTestNotificationHandler.Reflect(m_app.GetBehaviorContext());
+
+        AZ::Entity e;
+        Activate(e);
+
+        SimulateEditorBecomingInitialized();
+
+        enum class LogTypes
+        {
+            Skip = 0,
+            Notifications_OnFire,
+        };
+
+        m_testSink.m_evaluateMessage = [](const char* window, const char* message) -> int
+        {
+            if (AzFramework::StringFunc::Equal(window, "python"))
+            {
+                if (AzFramework::StringFunc::Equal(message, "Notifications_OnFire"))
+                {
+                    return static_cast<int>(LogTypes::Notifications_OnFire);
+                }
+            }
+            return static_cast<int>(LogTypes::Skip);
+        };
+
+        UnitTest::PythonTestNotificationHandler::Reset();
+
+        const int numFiresInParallel = 220;
+
+        const AZStd::string script = AZStd::string::format(
+            "import azlmbr.bus\n"
+            "import azlmbr.test\n"
+            "\n"
+            "def OnFire(parameters) :\n"
+            "    print('Notifications_OnFire')\n"
+            "\n"
+            "handler = azlmbr.bus.NotificationHandler('PythonTestSingleAddressNotificationBus')\n"
+            "handler.connect(None)\n"
+            "handler.add_callback('OnFire', OnFire)\n"
+            "\n"
+            "azlmbr.test.PythonTestNotificationHandler_do_fire()\n"
+            "\n"
+            "azlmbr.test.PythonTestNotificationHandler_do_fires_in_parallel(%d)\n"
+            "\n"
+            "handler.disconnect()\n",
+            numFiresInParallel);
+
+        AZ_TEST_START_TRACE_SUPPRESSION;
+        AzToolsFramework::EditorPythonRunnerRequestBus::Broadcast(&AzToolsFramework::EditorPythonRunnerRequestBus::Events::ExecuteByString, script, false);
+        AZ_TEST_STOP_TRACE_SUPPRESSION(numFiresInParallel); // Expect numFiresInParallel errors
+
+        e.Deactivate();
+
+        EXPECT_EQ(1, m_testSink.m_evaluationMap[static_cast<int>(LogTypes::Notifications_OnFire)]);
+    }
+
     TEST_F(PythonBusProxyTests, NotificationsWithNoAddress)
     {
         PythonTestNotificationHandler pythonTestNotificationHandler;

Gems/EditorPythonBindings/Code/Tests/PythonProxyObjectTests.cpp

@@ -419,7 +419,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
     };

Gems/EditorPythonBindings/Code/Tests/PythonReflectionComponentTests.cpp

@@ -532,7 +532,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
     };

Gems/EditorPythonBindings/Code/Tests/PythonThreadingTests.cpp

@@ -75,7 +75,7 @@ namespace UnitTest
         void TearDown() override
         {
             // clearing up memory
-            m_testSink = PythonTraceMessageSink();
+            m_testSink.CleanUp();
             PythonTestingFixture::TearDown();
         }
     };
@@ -114,38 +114,29 @@ namespace UnitTest
         try
         {
             // prepare handler on this thread
-            pybind11::exec(R"(
+            const AZStd::string_view script =
-                import azlmbr.test
+                "import azlmbr.test\n"
-
+                "\n"
-                def on_notification(args):
+                "def on_notification(args) :\n"
-                    value = args[0] + 2
+                "    value = args[0] + 2\n"
-                    print ('RanInThread')
+                "    print('RanInThread')\n"
-                    return value
+                "    return value\n"
-
+                "\n"
-                handler = azlmbr.test.PythonThreadNotificationBusHandler()
+                "handler = azlmbr.test.PythonThreadNotificationBusHandler()\n"
-                handler.connect()
+                "handler.connect()\n"
-                handler.add_callback('OnNotification', on_notification)
+                "handler.add_callback('OnNotification', on_notification)\n";
-                )");
+
+            AzToolsFramework::EditorPythonRunnerRequestBus::Broadcast(&AzToolsFramework::EditorPythonRunnerRequestBus::Events::ExecuteByString, script, false /*printResult*/);
 
             // start thread; in thread issue notification
             auto threadCallback = []()
             {
                 AZ::s64 result = 0;
-                auto notificationCallback = [&result]()
-                {
                 PythonThreadNotificationBus::BroadcastResult(result, &PythonThreadNotificationBus::Events::OnNotification, 40);
-                };
-
-                auto editorPythonEventsInterface = AZ::Interface<AzToolsFramework::EditorPythonEventsInterface>::Get();
-                if (editorPythonEventsInterface)
-                {
-                    editorPythonEventsInterface->ExecuteWithLock(notificationCallback);
-                }
 
                 EXPECT_EQ(42, result);
             };
             AZStd::thread theThread(threadCallback);
-            AZStd::this_thread::sleep_for(AZStd::chrono::milliseconds(100));
             theThread.join();
         }
         catch ([[maybe_unused]] const std::exception& e)
@@ -187,21 +178,11 @@ namespace UnitTest
             auto threadCallback = []()
             {
                 AZ::s64 result = 0;
-                auto notificationCallback = [&result]()
-                {
                 PythonThreadNotificationBus::BroadcastResult(result, &PythonThreadNotificationBus::Events::OnNotification, 40);
-                };
-
-                auto editorPythonEventsInterface = AZ::Interface<AzToolsFramework::EditorPythonEventsInterface>::Get();
-                if (editorPythonEventsInterface)
-                {
-                    editorPythonEventsInterface->ExecuteWithLock(notificationCallback);
-                }
 
                 EXPECT_EQ(0, result);
             };
             AZStd::thread theThread(threadCallback);
-            AZStd::this_thread::sleep_for(AZStd::chrono::milliseconds(100));
             theThread.join();
 
             // the Python script above raises an exception which causes two AZ_Error() message lines:

Gems/EditorPythonBindings/Code/Tests/PythonTraceMessageSink.h

@@ -37,6 +37,15 @@ namespace UnitTest
         using EvaluationMap = AZStd::unordered_map<int, int>; // tag to count
         EvaluationMap m_evaluationMap;
 
+        AZStd::mutex m_lock;
+
+        void CleanUp()
+        {
+            AZStd::lock_guard<decltype(m_lock)> lock(m_lock);
+            m_evaluateMessage = {};
+            m_evaluationMap.clear();
+        }
+
         //////////////////////////////////////////////////////////////////////////
         // TraceMessageDrillerBus
         void OnPrintf(const char* window, const char* message) override
@@ -46,6 +55,8 @@ namespace UnitTest
 
         void OnOutput(const char* window, const char* message) override
         {
+            AZStd::lock_guard<decltype(m_lock)> lock(m_lock);
+
             if (m_evaluateMessage)
             {
                 int key = m_evaluateMessage(window, message);