From 53215e57b152e71d841f4e6d8809d2c7516ca04c Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Sat, 23 Sep 2023 20:32:01 +0100
Subject: [PATCH] [*] Harden thread entrypoints and post-call deinit

---
 Source/Threading/Threads/AuOSThread.cpp | 43 +++++++++++++++++++++----
 Source/Threading/Threads/AuOSThread.hpp |  7 ++++
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/Source/Threading/Threads/AuOSThread.cpp b/Source/Threading/Threads/AuOSThread.cpp
index ba51498f..ac715944 100644
--- a/Source/Threading/Threads/AuOSThread.cpp
+++ b/Source/Threading/Threads/AuOSThread.cpp
@@ -77,7 +77,8 @@ namespace Aurora::Threading::Threads
 
 #endif
     
-    OSThread::OSThread(const ThreadInfo &info) : info_(info)
+    OSThread::OSThread(const ThreadInfo &info) : info_(info),
+        epExecEvent(false, false, true)
     {
         this->name_ = info.name.value_or("Aurora Thread");
 
@@ -91,7 +92,8 @@ namespace Aurora::Threading::Threads
         SysAssert(this->terminatedSignalLs_ && this->terminateSignalLs_ && this->terminated_ && this->terminateSignal_);
     }
 
-    OSThread::OSThread() : info_(gDummyThreadInfo)
+    OSThread::OSThread() : info_(gDummyThreadInfo),
+        epExecEvent(false, false, true)
     {
         this->name_ = "Main Thread";
 
@@ -102,7 +104,8 @@ namespace Aurora::Threading::Threads
         this->bNotOwned = true;
     }
 
-    OSThread::OSThread(AuUInt64 id) : info_(gDummyThreadInfo)
+    OSThread::OSThread(AuUInt64 id) : info_(gDummyThreadInfo),
+        epExecEvent(false, false, true)
     {
         this->name_ = "System Thread";
         this->handle_ = reinterpret_cast<decltype(handle_)>(id);
@@ -138,6 +141,13 @@ namespace Aurora::Threading::Threads
 #endif
         if (this->contextUsed_)
         {
+            this->epExecEvent->Wait();
+
+            {
+                AU_LOCK_GUARD(this->pFlag->mutex);
+                this->pFlag->bLock = true;
+            }
+
             if (this->detached_)
             {
                 bDetached = true;
@@ -569,19 +579,28 @@ namespace Aurora::Threading::Threads
         this->bSupportsAltKill = true;
     #endif
 
+        auto pFlag = this->pFlag;
+
         try
         {
-            if (task_)
+            if (auto task = task_)
             {
-                task_();
+                this->epExecEvent->Set();
+                task();
             }
         }
-
 #if !defined(AURORA_IS_POSIX_DERIVED)
         catch (...)
         {
             SysPushErrorHAL("OS Thread Aborted");
         }
+
+        AU_LOCK_GUARD(pFlag->mutex);
+        if (pFlag->bLock)
+        {
+            return;
+        }
+
         Exit(true);
 #else
     #if defined(AURORA_COMPILER_GCC)
@@ -592,6 +611,12 @@ namespace Aurora::Threading::Threads
     #endif
         catch (...)
         {
+            AU_LOCK_GUARD(pFlag->mutex);
+            if (pFlag->bLock)
+            {
+                return;
+            }
+
             bFailing = true;
             if (!Aurora::kIsDebugBuild)
             {
@@ -617,6 +642,12 @@ namespace Aurora::Threading::Threads
             }
         }
 
+        AU_LOCK_GUARD(pFlag->mutex);
+        if (pFlag->bLock)
+        {
+            return;
+        }
+
         if (!bFailing)
         {
             Exit(true);
diff --git a/Source/Threading/Threads/AuOSThread.hpp b/Source/Threading/Threads/AuOSThread.hpp
index 2675db27..f7cb7bd5 100644
--- a/Source/Threading/Threads/AuOSThread.hpp
+++ b/Source/Threading/Threads/AuOSThread.hpp
@@ -93,6 +93,12 @@ namespace Aurora::Threading::Threads
         bool                      contextUsed_{}; // can this thread instance execute code again?
         Primitives::EventShared_t terminated_;
         Primitives::EventShared_t terminateSignal_;
+        struct Lock
+        {
+            bool bLock {};
+            AuMutex mutex;
+        };
+        AuSPtr<Lock>              pFlag = AuMakeSharedPanic<Lock>();
         AuSPtr<AuLoop::ILSEvent>  terminateSignalLs_;
         AuSPtr<AuLoop::ILSEvent>  terminatedSignalLs_;
         bool                      bLongJmpOnce {};
@@ -100,6 +106,7 @@ namespace Aurora::Threading::Threads
         bool                      bNotOwned {};
         Primitives::CriticalSection exitOnlyOnce_;
         AuList<AuSPtr<IThreadFeature>> threadFeatures_;
+        AuBinarySemaphore         epExecEvent;
 
         AuFunction<void()> task_;