AuroraRuntime/Source/Threading/Primitives/AuConditionVariable.NT.cpp

/***
    Copyright (C) 2021-2024 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: AuConditionVariable.Win32.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuConditionVariable.Generic.hpp"
#include <Time/Time.hpp>
#include "SMTYield.hpp"
#include "../AuWakeInternal.hpp"

#if !defined(_AURUNTIME_GENERICCV)

#if !defined(NTSTATUS_TIMEOUT)
    #define NTSTATUS_TIMEOUT 0x102
#endif

namespace Aurora::Threading::Primitives
{
    ConditionVariableNT::ConditionVariableNT()
    {
    #if defined(AURORA_FORCE_SRW_LOCKS)
        ::InitializeConditionVariable(&this->winCond_);
    #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)
        bool bRet { true };
        bool bIOU {};

        auto pThatMutex = reinterpret_cast<NT4Mutex *>(&pMutex->lock_);

        this->AddWaiter();

        pMutex->Unlock();

        if (qwTimeout)
        {
        #if defined(_AURORA_WANT_STRICT_NT_WAKEUP_TIME_NO_WALL_CLOCK_SHIFT_ALLOWED)
            auto uEndTimeSteady = gUseNativeWaitCondvar ? AuTime::SteadyClockNS() + qwTimeout : 0;
            auto uEndTimeWall   = AuTime::CurrentClockNS() + qwTimeout;
            auto uTargetTimeNt  = AuTime::ConvertTimestampNs(uEndTimeWall);
        #else
            auto uEndTimeSteady = AuTime::SteadyClockNS() + qwTimeout;
        #endif

            while (true)
            {
                if (bRet)
                {
                    if (gUseNativeWaitCondvar)
                    {
                        // Reverted: 5b495f7fd9495aa55395666e166ac499955215dc
                        if (bSpin &&
                            ThrdCfg::gPreferNtCondvarModernWinSpin)
                        {
                            if (this->CheckOut())
                            {
                                pMutex->Lock();
                                return true;
                            }
                        }
                        else if (this->CheckOutNoSpin())
                        {
                            pMutex->Lock();
                            return true;
                        }
                       
                        AuUInt8 uBlockBit { 1 };
                        bRet = InternalLTSWaitOnAddressHighRes(&this->wlist, &uBlockBit, 1, uEndTimeSteady);
                    }
                    else
                    {
                        // Reverted: 5b495f7fd9495aa55395666e166ac499955215dc
                        if (bSpin &&
                            ThrdCfg::gPreferNtCondvarOlderWinSpin)
                        {
                            if (!bIOU)
                            {
                                bIOU = this->CheckOut();
                            }
                        }

                        {
                            LARGE_INTEGER word;
                        #if defined(_AURORA_WANT_STRICT_NT_WAKEUP_TIME_NO_WALL_CLOCK_SHIFT_ALLOWED)
                            word.QuadPart = uTargetTimeNt;
                        #else
                            auto uNow = AuTime::SteadyClockNS();

                            if (uEndTimeSteady > uNow)
                            {
                                word.QuadPart = -((uEndTimeSteady - uNow) / 100ull);
                            }
                            else
                            {
                                word.QuadPart = 0;
                            }
                        #endif

                            bRet = pNtWaitForKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, 0, &word) == 0;
                        }
                    }
                }
                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*/
                    LARGE_INTEGER word;
                    word.QuadPart = 0;

                    if (gUseNativeWaitCondvar)
                    {
                        // 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())
                        {
                            pMutex->Lock();
                            return bIOU;
                        }

                        AuUInt8 uBlockBit { 1 };
                        InternalLTSWakeAll((void *)&this->wlist); // this is kinda sad
                        bRet = InternalLTSWaitOnAddressHighRes(&this->wlist, &uBlockBit, 1, uEndTimeSteady); // why?
                    }
                    else
                    {
                        // Obligatory Windows XP+ resched
                        bRet = pNtWaitForKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, 0, nullptr) == 0;
                    }
                }

                if (!bRet)
                {
                    while (true)
                    {
                        auto uNow = this->wlist;
                        auto uOld = (uNow >> kShiftCountByBits);

                        if (uOld == 0)
                        {
                            break;
                        }

                        auto waiting = uOld - 1u;
                        auto uNext   = waiting << kShiftCountByBits;

                        if (AuAtomicCompareExchange(&this->wlist, uNext, uNow) == uNow)
                        {
                            pMutex->Lock();
                            return bIOU;
                        }
                    }
                }
                else
                {
                    // Acquire our signal before spinning and potentially giving up on our timeslice.
                    // If we were awoken, try to acquire before yielding to a contested mutex.
                    if (bIOU || this->CheckOutNoSpin())
                    {
                        pMutex->Lock();
                        return true;
                    }

                    if (!gUseNativeWaitCondvar)
                    {
                        // (NT 5-6.1)   Unblocks one race condition, where another thread checks out our signal, and we return with nothing.
                        //              Normal execution of the blocked thread can continue, so long as we trigger the keyed event as though this thread was signaled.
                        //              This is only relevant in an edge case of multi-waiters spinning, rapid signaling, and the IOU being hit once. Do this instead of ticketing.
                        pNtReleaseKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, FALSE, nullptr);
                    }
                }
            }
        }
        else
        {
            while (true)
            {
                if (gUseNativeWaitCondvar)
                {
                    if (bSpin &&
                        ThrdCfg::gPreferNtCondvarModernWinSpin)
                    {
                        if (this->CheckOut())
                        {
                            pMutex->Lock();
                            return true;
                        }
                    }
                    else if (this->CheckOutNoSpin())
                    {
                        pMutex->Lock();
                        return true;
                    }

                    AuUInt8 uBlockBit { 1 };
                    bRet = InternalLTSWaitOnAddressHighRes(&this->wlist, &uBlockBit, 1, 0);
                }
                else
                {
                    if (bSpin &&
                        ThrdCfg::gPreferNtCondvarOlderWinSpin)
                    {
                        if (!bIOU)
                        {
                            bIOU = this->CheckOut();
                        }
                    }

                    pNtWaitForKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, 0, nullptr);
                }

                if (bIOU || this->CheckOutNoSpin())
                {
                    pMutex->Lock();
                    return bRet;
                }

                if (!gUseNativeWaitCondvar)
                {
                    pNtReleaseKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, FALSE, nullptr);
                }
            }
        }
    #else
        AuInt64 uEndTime = qwTimeout ? AuTime::SteadyClockNS() + qwTimeout : 0;

        while (true)
        {
            AuInt64 uSecondTimeout = INFINITE;

            if (qwTimeout)
            {
                uSecondTimeout = uEndTime - AuTime::SteadyClockNS();

                if (uSecondTimeout <= 0)
                {
                    return false;
                }

                uSecondTimeout = AuNSToMS<AuUInt64>(uSecondTimeout);
                if (!uSecondTimeout)
                {
                    SMPPause();
                    AuThreading::ContextYield();
                    continue;
                }
            }

            auto bOk = ::SleepConditionVariableSRW(&this->winCond_, reinterpret_cast<PSRWLOCK>(pMutex->GetOSHandle()), uSecondTimeout, 0);

            if (bOk)
            {
                return true;
            }
            else
            {
                SysAssert(GetLastError() == ERROR_TIMEOUT, "SleepConditionVariable failure");
            }
        }
    #endif
    }

    bool ConditionVariableNT::CheckOutNoSpin()
    {
    #if defined(AURORA_FORCE_SRW_LOCKS)
        return false;
    #else
        AuUInt32 uSignalNow {};
        AuUInt32 uSignalNext {};

        while (true)
        {
            uSignalNow = AuAtomicLoad(&this->signalCount);

            if (uSignalNow == 0)
            {
                return false;
            }

            uSignalNext = uSignalNow - 1;

            if (AuAtomicCompareExchange(&this->signalCount, uSignalNext, uSignalNow) == uSignalNow)
            {
                break;
            }
        }

        if (uSignalNext == 0)
        {
            // Last thread sets the wait-bit (not atomic)
            InterlockedOr((volatile LONG *)&this->wlist, 1);

            AuUInt32 uSignalCount {};
            if ((uSignalCount = AuAtomicLoad(&this->signalCount)) != 0) [[unlikely]]
            {
                // Our CAS was wrong, and we potentially cock-blocked a waking thread, unset the sleep bit
                // (note, i dont see this happening in real world applications, but lets continue)
                // (also note, this could've been avoided with stict 8-byte atomic ops across all archs. no thank you.)
                AuAtomicUnset(&this->wlist, 0);

                // WaitOnAddress-based paranoia (only keyed events wont be lost throughout this potential nonatomic race condition):
                if (gUseNativeWaitCondvar)
                {
                    // Whatever, NT kernelbase's monitor address signaling is almost free
                    if (uSignalCount > 1)
                    {
                        InternalLTSWakeAll((const void *)&this->wlist);
                    }
                    else
                    {
                        InternalLTSWakeOne((const void *)&this->wlist);
                    }
                }
            }
        }

        return true;
    #endif
    }

    bool ConditionVariableNT::CheckOut()
    {
    #if defined(AURORA_FORCE_SRW_LOCKS)
        return false;
    #else
        return DoTryIf([&]()
        {
            return this->CheckOutNoSpin();
        });
    #endif
    }
    
    void ConditionVariableNT::Signal()
    {
    #if !defined(AURORA_FORCE_SRW_LOCKS)
        auto original = this->wlist;
        auto expected = original;
        expected = expected >> kShiftCountByBits;

        if (expected)
        {
            AuAtomicAdd(&this->signalCount, 1u);
            
            while (expected)
            {
                if (AuAtomicCompareExchange(&this->wlist, ((expected - 1) << kShiftCountByBits)  /*intentional clear*/, original) == original)
                {
                    if (gUseNativeWaitCondvar)
                    {
                        InternalLTSWakeOne((void *)&this->wlist);
                    }
                    else
                    {
                        pNtReleaseKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, FALSE, nullptr);
                    }

                    return;
                }

                original = this->wlist;
                expected = original >> kShiftCountByBits;
            }
        }
    #else
        ::WakeConditionVariable(&this->winCond_);
    #endif
    }

    void ConditionVariableNT::Broadcast()
    {
    #if !defined(AURORA_FORCE_SRW_LOCKS)
        if (gUseNativeWaitCondvar)
        {
            auto original = this->wlist;
            auto expected = original;
            expected      = expected >> kShiftCountByBits;

            if (!expected)
            {
                return;
            }

            AuAtomicAdd(&this->signalCount, expected);
            
            auto uAwoken = expected;

            while (true)
            {
                if (AuAtomicCompareExchange(&this->wlist,
                                            (expected - uAwoken) << kShiftCountByBits,
                                            original) == original)
                {
                    InternalLTSWakeAll((void *)&this->wlist);
                    return;
                }
                else
                {
                    original = this->wlist;
                    expected = original >> kShiftCountByBits;
                    uAwoken = AuMin<AuUInt32>(uAwoken, expected);
                }
            }
        }
        else
        {
            auto original = this->wlist;
            auto expected = original;
            expected = expected >> kShiftCountByBits;

            auto uBroadcastIterations = expected;

            while (expected && uBroadcastIterations)
            {
                AuAtomicAdd(&this->signalCount, 1u);

                while (expected && uBroadcastIterations)
                {
                    bool bBreak {};
                    if (AuAtomicCompareExchange(&this->wlist, ((expected - 1) << kShiftCountByBits)  /*intentional clear*/, original) == original)
                    {
                        pNtReleaseKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, FALSE, nullptr);

                        uBroadcastIterations--;
                        bBreak = true;
                    }

                    original = this->wlist;
                    expected = original >> kShiftCountByBits;

                    if (bBreak)
                    {
                        break;
                    }
                }
            }
        }
    #else
        ::WakeAllConditionVariable(&this->winCond_);
    #endif
    }

    void ConditionVariableNT::BroadcastN(AuUInt32 nBroadcast)
    {
    #if !defined(AURORA_FORCE_SRW_LOCKS)
        if (gUseNativeWaitCondvar)
        {
            auto original = this->wlist;
            auto expected = original;
            expected      = AuMin(nBroadcast, expected >> kShiftCountByBits);

            if (!expected)
            {
                return;
            }

            AuAtomicAdd(&this->signalCount, expected);
            
            auto uAwoken = expected;

            while (true)
            {
                auto uCount = expected - uAwoken;

                if (AuAtomicCompareExchange(&this->wlist,
                                            uCount << kShiftCountByBits,
                                            original) == original)
                {
                    InternalLTSWakeCount((void *)&this->wlist, uAwoken);
                    return;
                }
                else
                {
                    original = this->wlist;
                    expected = original >> kShiftCountByBits;
                    uAwoken = AuMin<AuUInt32>(uAwoken, expected);
                }
            }
        }
        else
        {
            auto original = this->wlist;
            auto expected = original;
            expected = expected >> kShiftCountByBits;

            auto uBroadcastIterations = AuMin(nBroadcast, expected);

            while (expected && uBroadcastIterations)
            {
                AuAtomicAdd(&this->signalCount, 1u);

                while (expected && uBroadcastIterations)
                {
                    bool bBreak {};
                    if (AuAtomicCompareExchange(&this->wlist, ((expected - 1) << kShiftCountByBits)  /*intentional clear*/, original) == original)
                    {
                        pNtReleaseKeyedEvent(gKeyedEventHandle, (void *)&this->wlist, FALSE, nullptr);

                        uBroadcastIterations--;
                        bBreak = true;
                    }

                    original = this->wlist;
                    expected = original >> kShiftCountByBits;

                    if (bBreak)
                    {
                        break;
                    }
                }
            }
        }
    #else
        ::WakeAllConditionVariable(&this->winCond_);
    #endif
    }

    AUKN_SYM IConditionVariable *ConditionVariableNew(const AuSPtr<IConditionMutex> &pMutex)
    {
        return _new ConditionVariableImpl(pMutex);
    }

    AUKN_SYM void ConditionVariableRelease(IConditionVariable *pCV)
    {
        AuSafeDelete<ConditionVariableImpl *>(pCV);
    }

    AUROXTL_INTERFACE_SOO_SRC(ConditionVariable, ConditionVariableImpl, (const AuSPtr<IConditionMutex> &, pMutex))
}

#endif