From 100964ac87bda5707805cb1677410b5f24a2a0a1 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Thu, 24 Aug 2023 13:55:31 +0100
Subject: [PATCH] [*] NT Semaphore optimization

---
 .../Threading/Primitives/AuSemaphore.NT.cpp   | 93 +++++++++++++++++--
 .../Threading/Primitives/AuSemaphore.NT.hpp   |  2 +
 2 files changed, 86 insertions(+), 9 deletions(-)

diff --git a/Source/Threading/Primitives/AuSemaphore.NT.cpp b/Source/Threading/Primitives/AuSemaphore.NT.cpp
index 6a1b9849..332ceb43 100644
--- a/Source/Threading/Primitives/AuSemaphore.NT.cpp
+++ b/Source/Threading/Primitives/AuSemaphore.NT.cpp
@@ -60,6 +60,11 @@ namespace Aurora::Threading::Primitives
         }
     }
     
+    AuUInt32 *SemaphoreImpl::GetSleepCounter()
+    {
+        return (AuUInt32 *)&this->var;
+    }
+
     bool SemaphoreImpl::LockMS(AuUInt64 uTimeout)
     {
         return LockNS(AuMSToNS<AuUInt64>(uTimeout));
@@ -84,11 +89,17 @@ namespace Aurora::Threading::Primitives
         {
             AuUInt32 uYieldCounter {};
             auto old = this->dwState_;
-                 //!tryLock (with old in a scope we can access)
             while (!((old != 0) &&
                      (AuAtomicCompareExchange(&this->dwState_, old - 1, old) == old)))
             {
-                if (!InternalLTSWaitOnAddressHighRes(&this->dwState_, &old, sizeof(this->dwState_), uEnd))
+                static const AuUInt32 kExpect { 0 };
+
+                auto pCounter = GetSleepCounter();
+                AuAtomicAdd(pCounter, 1u);
+                bool bStatus = InternalLTSWaitOnAddressHighRes(&this->dwState_, (void *)&kExpect, sizeof(kExpect), uEnd);
+                AuAtomicSub(pCounter, 1u);
+
+                if (!bStatus)
                 {
                     return false;
                 }
@@ -122,9 +133,69 @@ namespace Aurora::Threading::Primitives
             }
 
             this->mutex.Unlock();
+
+            return true;
+        }
+    }
+
+    bool SemaphoreImpl::LockAbsNS(AuUInt64 qwTimeoutAbs)
+    {
+        if (this->TryLockHeavy())
+        {
+            return true;
         }
 
-        return true;
+        if (gUseNativeWaitSemapahore)
+        {
+            AuUInt32 uYieldCounter {};
+            auto old = this->dwState_;
+            while (!((old != 0) &&
+                     (AuAtomicCompareExchange(&this->dwState_, old - 1, old) == old)))
+            {
+                static const AuUInt32 kExpect { 0 };
+
+                auto pCounter = GetSleepCounter();
+                AuAtomicAdd(pCounter, 1u);
+                bool bStatus = InternalLTSWaitOnAddressHighRes(&this->dwState_, (void *)&kExpect, sizeof(kExpect), qwTimeoutAbs);
+                AuAtomicSub(pCounter, 1u);
+
+                if (!bStatus)
+                {
+                    return false;
+                }
+
+                old = this->dwState_;
+            }
+
+            return true;
+        }
+        else
+        {
+            this->mutex.Lock();
+
+            while (!TryLockNoSpin())
+            {
+                if (qwTimeoutAbs != 0)
+                {
+                    auto uStart = Time::SteadyClockNS();
+                    if (uStart >= qwTimeoutAbs)
+                    {
+                        this->mutex.Unlock();
+                        return false;
+                    }
+
+                    var.WaitForSignalNsEx(&this->mutex, qwTimeoutAbs - uStart);
+                }
+                else
+                {
+                    var.WaitForSignalNsEx(&this->mutex, 0);
+                }
+            }
+
+            this->mutex.Unlock();
+
+            return true;
+        }
     }
     
     void SemaphoreImpl::Lock()
@@ -138,13 +209,17 @@ namespace Aurora::Threading::Primitives
         if (gUseNativeWaitSemapahore)
         {
             AuAtomicAdd<AuUInt32>(&this->dwState_, count);
-            if (count == 1)
+
+            if (AuAtomicLoad(GetSleepCounter()))
             {
-                pWakeByAddressSingle(&this->dwState_);
-            }
-            else
-            {
-                pWakeByAddressAll(&this->dwState_);
+                if (count == 1)
+                {
+                    pWakeByAddressSingle(&this->dwState_);
+                }
+                else
+                {
+                    pWakeByAddressAll(&this->dwState_);
+                }
             }
         }
         else
diff --git a/Source/Threading/Primitives/AuSemaphore.NT.hpp b/Source/Threading/Primitives/AuSemaphore.NT.hpp
index 0ef51e8e..f0ae7821 100644
--- a/Source/Threading/Primitives/AuSemaphore.NT.hpp
+++ b/Source/Threading/Primitives/AuSemaphore.NT.hpp
@@ -22,6 +22,7 @@ namespace Aurora::Threading::Primitives
         bool TryLock() override;
         bool LockMS(AuUInt64 timeout) override;
         bool LockNS(AuUInt64 timeout) override;
+        bool LockAbsNS(AuUInt64 qwTimeoutAbs) override;
         void Lock() override;
         void Unlock(AuUInt16 count) override;
         void Unlock() override;
@@ -29,6 +30,7 @@ namespace Aurora::Threading::Primitives
         auline bool TryLockNoSpin();
         auline bool TryLockHeavy();
 
+        auline AuUInt32 *GetSleepCounter();
     private:
         AuUInt32 dwState_ {};
         ConditionMutexInternal mutex;