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

/***
    Copyright (C) 2021 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 "SMPYield.hpp"

#if !defined(_AURUNTIME_GENERICCV)

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

namespace Aurora::Threading::Primitives
{
    ConditionVariableImpl::ConditionVariableImpl(const AuSPtr<IConditionMutex> &pMutex) :
        mutex_(AuStaticCast<Win32ConditionMutex>(pMutex))
    {
    }

    AuSPtr<IConditionMutex> ConditionVariableImpl::GetMutex()
    {
        return mutex_;
    }

    bool ConditionVariableImpl::WaitForSignal(AuUInt32 uTimeout)
    {
        return WaitForSignalNS(AuMSToNS<AuUInt64>(uTimeout));
    }

    bool ConditionVariableImpl::WaitForSignalNS(AuUInt64 qwTimeout)
    {
        bool bRet { true };
        auto pThatMutex =  reinterpret_cast<NT4Mutex *>(&this->mutex_->lock_);

        this->mutex_->Unlock();
  
        if (qwTimeout)
        {
            auto uEndTimeSteady = AuTime::SteadyClockNS() + qwTimeout;
            auto uEndTimeWall   = AuTime::CurrentClockNS() + qwTimeout;
            auto uTargetTimeNt  = AuTime::ConvertTimestampNs(uEndTimeWall);

            while (true)
            {
                LARGE_INTEGER word;
                word.QuadPart = uTargetTimeNt;


                // forced smp stall
                // see the "hopefully nt is smart enough" comment
                bool bIOU = DoTryIf([=]()
                {
                    bool b = true;
                    return CheckOut(b);
                });

                if (bRet)
                {
                    while (true)
                    {
                        auto uNow = this->wlist;
                        auto waiting = uNow >> 2u;
                        auto uNext = ((waiting + 1) << 2u) | 1;

                        if (AuAtomicCompareExchange(&this->wlist, uNext, uNow) == uNow)
                        {
                            break;
                        }
                    }

                    bRet = pNtWaitForKeyedEvent(gKeyedEventHandle, &this->wlist, 0, &word) != NTSTATUS_TIMEOUT;
                }
                else
                {
                    LARGE_INTEGER word;
                    word.QuadPart = 0;
                    bRet = pNtWaitForKeyedEvent(gKeyedEventHandle, &this->wlist, 0, &word) != NTSTATUS_TIMEOUT;
                }

            #if !defined(AU_TRUST_NT_KERNEL_SCHED_TIMEOUT)
                if (!bRet)
            #else
                if (!bRet && uEndTimeSteady <= AuTime::SteadyClockNS())
            #endif
                {
                    auto uNow = this->wlist;
                    auto uOld = (uNow >> 2u);

                    if (uOld == 0)
                    {
                        // broadcast has woken everyone up

                        if (bIOU || CheckOut(bRet)) // the cope acquire
                        {
                            // in which case we're good
                            this->mutex_->Lock();
                            return true;
                        }
                        else
                        {
                            // in which case we're still dealaing with out overflowed waitlist
                            bRet = false;
                            continue;
                        }
                    }

                    // go for atomic decrement
                    auto waiting = uOld - 1u;
                    auto uNext   = waiting << 2u;

                    if (AuAtomicCompareExchange(&this->wlist, uNext, uNow) == uNow)
                    {
                        // break if successful
                        this->mutex_->Lock();
                        return bIOU;
                    }
                    else
                    {
                        // block again because we couldn't decrement the counter
                        // broadcast still thinks we're asleep
                        // ...and we still owe NtReleaseKeyedEvent 1 therad
                        continue;
                    }
                }
                else
                {
                    // we good?
                    if (bIOU || CheckOut(bRet))
                    {
                        this->mutex_->Lock();
                        return true;
                    }
                }
            }
        }
        else
        {
            while (true)
            {
                while (true)
                {
                    auto uNow = this->wlist;
                    auto waiting = uNow >> 2u;
                    auto uNext = ((waiting + 1) << 2u) | 1;

                    if (AuAtomicCompareExchange(&this->wlist, uNext, uNow) == uNow)
                    {
                        break;
                    }
                }

                // forced smp stall
                bool bIOU = DoTryIf([=]()
                {
                    bool b = true;
                    return CheckOut(b);
                });

                // then into the kernel no matter what (hopefully nt is smart enough to have a fast path)
                pNtWaitForKeyedEvent(gKeyedEventHandle, &this->wlist, 0, nullptr);

                if (bIOU || CheckOut(bRet))
                {
                    this->mutex_->Lock();
                    return bRet;
                }
            }
        }
    }

    bool ConditionVariableImpl::CheckOut(bool &bRet)
    {
        return DoTryIf([&]()
        {
            auto uSignalNow = this->signalCount;

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

            return AuAtomicCompareExchange(&this->signalCount, uSignalNow - 1, uSignalNow) == uSignalNow;
        });
    }

    void ConditionVariableImpl::Signal()
    {
        auto original = this->wlist;
        auto expected = original;
        expected = expected >> 2;

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

                original = this->wlist;
                expected = original >> 2;
            }
        }
    }

    void ConditionVariableImpl::Broadcast()
    {
        auto original = this->wlist;
        auto expected = original;
        expected = expected >> 2;

        auto uBroadcastIterations = expected;

        while (expected && uBroadcastIterations)
        {
            if (AuAtomicCompareExchange(&this->wlist, ((expected - 1) << 2)  /*intentional clear*/, original) == original)
            {
                AuAtomicAdd(&this->signalCount, 1u);

                pNtReleaseKeyedEvent(gKeyedEventHandle, &this->wlist, FALSE, nullptr);

                uBroadcastIterations--;
            }

            original = this->wlist;
            expected = original >> 2;
        }
    }

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

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

#endif