From 2943ffdbc2c92331dee348c00bc4ee9453c6de10 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Sat, 27 Jan 2024 11:46:18 +0000
Subject: [PATCH] [*] Harden Windows XP - 7 condvar; where signalers can lock
 up pending scheduling n threads whilst a spurious wake is occurring and
 trying to preemptively require the lock under the signal

---
 .../Primitives/AuConditionVariable.NT.cpp     | 102 ++++++++----------
 .../Primitives/AuConditionVariable.NT.hpp     |   3 +-
 2 files changed, 49 insertions(+), 56 deletions(-)

diff --git a/Source/Threading/Primitives/AuConditionVariable.NT.cpp b/Source/Threading/Primitives/AuConditionVariable.NT.cpp
index 30412a77..63f25a44 100644
--- a/Source/Threading/Primitives/AuConditionVariable.NT.cpp
+++ b/Source/Threading/Primitives/AuConditionVariable.NT.cpp
@@ -26,6 +26,21 @@ namespace Aurora::Threading::Primitives
     #endif
     }
   
+    void ConditionVariableNT::AddWaiter()
+    {
+        while (true)
+        {
+            auto uNow = this->wlist;
+            auto waiting = uNow >> kShiftCountByBits;
+            auto uNext = ((waiting + 1) << kShiftCountByBits) | (!bool(waiting)) | (uNow & 1);
+
+            if (AuAtomicCompareExchange(&this->wlist, uNext, uNow) == uNow)
+            {
+                break;
+            }
+        }
+    }
+
     bool ConditionVariableNT::WaitForSignalNsEx(Win32ConditionMutex *pMutex, AuUInt64 qwTimeout, bool bSpin)
     {
     #if !defined(AURORA_FORCE_SRW_LOCKS)
@@ -39,6 +54,10 @@ namespace Aurora::Threading::Primitives
             auto uTargetTimeNt  = AuTime::ConvertTimestampNs(uEndTimeWall);
             bool bIOU {};
 
+            this->AddWaiter();
+
+            pMutex->Unlock();
+
             while (true)
             {
                 LARGE_INTEGER word;
@@ -46,20 +65,6 @@ namespace Aurora::Threading::Primitives
 
                 if (bRet)
                 {
-                    while (true)
-                    {
-                        auto uNow = this->wlist;
-                        auto waiting = uNow >> kShiftCountByBits;
-                        auto uNext = ((waiting + 1) << kShiftCountByBits) | (!bool(waiting)) | (uNow & 1);
-
-                        if (AuAtomicCompareExchange(&this->wlist, uNext, uNow) == uNow)
-                        {
-                            break;
-                        }
-                    }
-                    
-                    pMutex->Unlock();
-
                     if (gUseNativeWaitCondvar)
                     {
                         // Reverted: 5b495f7fd9495aa55395666e166ac499955215dc
@@ -93,10 +98,8 @@ namespace Aurora::Threading::Primitives
                             }
                         }
 
-                        bRet = pNtWaitForKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, 0, &word) != NTSTATUS_TIMEOUT;
+                        bRet = pNtWaitForKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, 0, &word) == 0;
                     }
-
-                    pMutex->Lock();
                 }
                 else /* unblock NtReleaseKeyedEvent after an atomic this->wlist change <-> NtReleaseKeyedEvent race condition. */
                 {    /* this->wlist waiters should still be accounting for us, leading to a NtReleaseKeyedEvent block condition*/
@@ -105,8 +108,6 @@ namespace Aurora::Threading::Primitives
 
                     if (gUseNativeWaitCondvar)
                     {
-                        pMutex->Unlock();
-
                         // do not do a syscall (worst case) under ::unlock(...) barrier/pMutex's lock
                         // best case: rdtsc extrapolation, if nts hal trusts our system
                         if (uEndTimeSteady <= AuTime::SteadyClockNS())
@@ -121,35 +122,31 @@ namespace Aurora::Threading::Primitives
                     }
                     else
                     {
-                        pMutex->Unlock();
-
                         // Obligatory Windows XP+ resched
-                        bRet = pNtWaitForKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, 0, nullptr) != NTSTATUS_TIMEOUT;
+                        bRet = pNtWaitForKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, 0, nullptr) == 0;
                     }
-
-                    pMutex->Lock();
                 }
 
                 if (!bRet)
                 {
-                    auto uNow = this->wlist;
-                    auto uOld = (uNow >> kShiftCountByBits);
-
-                    if (uOld == 0)
+                    while (true)
                     {
-                        continue;
-                    }
+                        auto uNow = this->wlist;
+                        auto uOld = (uNow >> kShiftCountByBits);
 
-                    auto waiting = uOld - 1u;
-                    auto uNext   = waiting << kShiftCountByBits;
+                        if (uOld == 0)
+                        {
+                            break;
+                        }
 
-                    if (AuAtomicCompareExchange(&this->wlist, uNext, uNow) == uNow)
-                    {
-                        return bIOU;
-                    }
-                    else
-                    {
-                        continue;
+                        auto waiting = uOld - 1u;
+                        auto uNext   = waiting << kShiftCountByBits;
+
+                        if (AuAtomicCompareExchange(&this->wlist, uNext, uNow) == uNow)
+                        {
+                            pMutex->Lock();
+                            return bIOU;
+                        }
                     }
                 }
                 else
@@ -157,30 +154,24 @@ namespace Aurora::Threading::Primitives
                     if (bIOU ||
                         (bSpin ? this->CheckOut() : this->CheckOutNoSpin()))
                     {
+                        pMutex->Lock();
                         return true;
                     }
+
+                    this->AddWaiter();
                 }
             }
         }
         else
         {
             bool bIOU {};
+
+            this->AddWaiter();
+
+            pMutex->Unlock();
+
             while (true)
             {
-                while (true)
-                {
-                    auto uNow = this->wlist;
-                    auto waiting = uNow >> kShiftCountByBits;
-                    auto uNext = ((waiting + 1) << kShiftCountByBits) | (!bool(waiting)) | (uNow & 1);
-
-                    if (AuAtomicCompareExchange(&this->wlist, uNext, uNow) == uNow)
-                    {
-                        break;
-                    }
-                }
-
-                pMutex->Unlock();
-
                 if (gUseNativeWaitCondvar)
                 {
                     if (bSpin &&
@@ -215,13 +206,14 @@ namespace Aurora::Threading::Primitives
                     pNtWaitForKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, 0, nullptr);
                 }
 
-                pMutex->Lock();
-
                 if (bIOU ||
                     (bSpin ? this->CheckOut() : this->CheckOutNoSpin()))
                 {
+                    pMutex->Lock();
                     return bRet;
                 }
+
+                this->AddWaiter();
             }
         }
     #else
diff --git a/Source/Threading/Primitives/AuConditionVariable.NT.hpp b/Source/Threading/Primitives/AuConditionVariable.NT.hpp
index 6e29b5b7..efba6e3a 100644
--- a/Source/Threading/Primitives/AuConditionVariable.NT.hpp
+++ b/Source/Threading/Primitives/AuConditionVariable.NT.hpp
@@ -21,7 +21,8 @@ namespace Aurora::Threading::Primitives
         void Signal();
         void Broadcast();
         void BroadcastN(AuUInt32 nBroadcast);
-
+        
+        auline void AddWaiter();
         auline bool CheckOut();
         auline bool CheckOutNoSpin();