AuroraRuntime/Source/Threading/Primitives/AuRWLock.cpp

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

    File: AuRWLock.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
//#define RWLOCK_VIEW_HAS_PARENT
#include "AuRWLock.hpp"
#include "SMTYield.hpp"
#include "../AuWakeInternal.hpp"

namespace Aurora::Threading::Primitives
{
#if defined(RWLOCK_VIEW_HAS_PARENT)
    #define ViewParent (&this->parent_)
#else
    #define ViewParent ((T *)(((char *)this) - (bIsReadView ? RWLockImpl<true>::kOffsetOfRead : RWLockImpl<true>::kOffsetOfWrite)))
#endif

#define RWLockAcquire           AU_LOCK_GUARD(this->mutex_);
#define RWLockBarrier(expMutex) { (expMutex).Lock(); (expMutex).Unlock(); }
#define RWLockBarrierSelf()     RWLockBarrier(this->mutex_)

    static const auto kRWThreadWriterHardContextSwitchBias = 15;

    template<bool bIsReadView, typename T>
    void RWLockAccessView<bIsReadView, T>::Unlock()
    {
        if constexpr (bIsReadView)
        {
            ViewParent->UnlockRead();
        }
        else
        {
            ViewParent->UnlockWrite();
        }
    }

    template<bool bIsReadView, typename T>
    bool RWLockAccessView<bIsReadView, T>::LockAbsMS(AuUInt64 timeout)
    {
        if constexpr (bIsReadView)
        {
            return ViewParent->LockReadNSAbs(AuMSToNS<AuUInt64>(timeout));
        }
        else
        {
            return ViewParent->LockWriteNSAbs(AuMSToNS<AuUInt64>(timeout));
        }
    }

    template<bool bIsReadView, typename T>
    bool RWLockAccessView<bIsReadView, T>::LockAbsNS(AuUInt64 timeout)
    {
        if constexpr (bIsReadView)
        {
            return ViewParent->LockReadNSAbs(timeout);
        }
        else
        {
            return ViewParent->LockWriteNSAbs(timeout);
        }
    }

    template<bool bIsReadView, typename T>
    bool RWLockAccessView<bIsReadView, T>::LockMS(AuUInt64 timeout)
    {
        if constexpr (bIsReadView)
        {
            return ViewParent->LockReadNS(AuMSToNS<AuUInt64>(timeout));
        }
        else
        {
            return ViewParent->LockWriteNS(AuMSToNS<AuUInt64>(timeout));
        }
    }

    template<bool bIsReadView, typename T>
    bool RWLockAccessView<bIsReadView, T>::LockNS(AuUInt64 timeout)
    {
        if constexpr (bIsReadView)
        {
            return ViewParent->LockReadNS(timeout);
        }
        else
        {
            return ViewParent->LockWriteNS(timeout);
        }
    }

    template<bool bIsReadView, typename T>
    bool RWLockAccessView<bIsReadView, T>::TryLock()
    {
        if constexpr (bIsReadView)
        {
            return ViewParent->TryLockRead();
        }
        else
        {
            return ViewParent->TryLockWriteMaybeSpin();
        }
    }

    template<bool bIsWriteRecursionAllowed>
    RWLockImpl<bIsWriteRecursionAllowed>::RWLockImpl()
    #if defined(RWLOCK_VIEW_HAS_PARENT)
        : read_(*this),
        write_(*this)
    #endif
    #if 0
        , condition_(AuUnsafeRaiiToShared(&this->mutex_)),
        conditionWriter_(AuUnsafeRaiiToShared(&this->mutex_))
    #endif
    {

    }

    template<bool bIsWriteRecursionAllowed>
    RWLockImpl<bIsWriteRecursionAllowed>::~RWLockImpl()
    {
    }

    template<bool bIsWriteRecursionAllowed>
    ConditionVariableInternal &RWLockImpl<bIsWriteRecursionAllowed>::GetCondition()
    {
    #if !defined(AURWLOCK_NO_SIZE_OPTIMIZED_CONDVAR)
        return *(ConditionVariableInternal *)this->conditionVariable_;
    #else
        return this->condition_;
    #endif
    }

    template<bool bIsWriteRecursionAllowed>
    ConditionVariableInternal &RWLockImpl<bIsWriteRecursionAllowed>::GetConditionWriter()
    {
    #if !defined(AURWLOCK_NO_SIZE_OPTIMIZED_CONDVAR)
        return *(ConditionVariableInternal *)this->conditionVariableWriter_;
    #else
        return this->conditionWriter_;
    #endif
    }

    template<bool bIsWriteRecursionAllowed>
    AuUInt32 *RWLockImpl<bIsWriteRecursionAllowed>::GetFutexCondition()
    {
        return (AuUInt32 *)&this->iState_;
    }

    template<bool bIsWriteRecursionAllowed>
    AuUInt32 *RWLockImpl<bIsWriteRecursionAllowed>::GetFutexConditionWriter()
    {
        return (AuUInt32 *)this->conditionVariableWriter_;
    }

    template<bool bIsWriteRecursionAllowed>
    AuUInt32 *RWLockImpl<bIsWriteRecursionAllowed>::GetReadSleepCounter()
    {
        return (AuUInt32 *)this->conditionVariable_;
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::LockReadNSAbs(AuUInt64 uTimeout)
    {
        AuInt32 iCurState {};
        bool bRet {};

        if (this->TryLockRead())
        {
            return true;
        }

        auto pCounter = this->GetReadSleepCounter();

        do
        {
            iCurState = this->iState_;

            if (iCurState < 0)
            {
                if (gUseFutexRWLock)
                {
                    AuAtomicAdd(pCounter, 1u);
                    bRet = InternalLTSWaitOnAddressHighRes((const void *)&this->iState_, &iCurState, sizeof(iCurState), uTimeout);
                    AuAtomicSub(pCounter, 1u);

                    if (!bRet)
                    {
                        return false;
                    }
                }
                else
                {
                    RWLockAcquire;

                    iCurState = this->iState_;
                    if (iCurState < 0)
                    {
                        AuInt64 iSecondTimeout {};

                        if (uTimeout)
                        {
                            iSecondTimeout = AuInt64(uTimeout) - AuTime::SteadyClockNS();

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

                    #if defined(AURWLOCK_NO_SIZE_OPTIMIZED_CONDVAR)
                        if (!this->GetCondition().WaitForSignalNS(iSecondTimeout))
                    #else
                        if (!this->GetCondition().WaitForSignalNsEx(&this->mutex_, iSecondTimeout))
                    #endif
                        {
                            return false;
                        }
                    }
                }
            }
        }
        while (iCurState < 0 ||
               AuAtomicCompareExchange(&this->iState_, iCurState + 1, iCurState) != iCurState);

        return true;
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::LockReadNS(AuUInt64 uTimeout)
    {
        AuInt32 iCurState {};
        bool bRet {};

        if (this->TryLockReadNoSpin<true>())
        {
            return true;
        }

        AuInt64 uEndTime = uTimeout ? AuTime::SteadyClockNS() + uTimeout : 0;

        if (this->TryLockRead())
        {
            return true;
        }

        auto pCounter = this->GetReadSleepCounter();

        do
        {
            iCurState = this->iState_;

            if (iCurState < 0)
            {
                if (gUseFutexRWLock)
                {
                    AuAtomicAdd(pCounter, 1u);
                    bRet = InternalLTSWaitOnAddressHighRes((const void *)&this->iState_, &iCurState, sizeof(iCurState), uEndTime);
                    AuAtomicSub(pCounter, 1u);

                    if (!bRet)
                    {
                        return false;
                    }
                }
                else
                {
                    RWLockAcquire;

                    iCurState = this->iState_;
                    if (iCurState < 0)
                    {
                        AuInt64 iSecondTimeout {};

                        if (uTimeout)
                        {
                            iSecondTimeout = uEndTime - AuTime::SteadyClockNS();

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

                    #if defined(AURWLOCK_NO_SIZE_OPTIMIZED_CONDVAR)
                        if (!this->GetCondition().WaitForSignalNS(iSecondTimeout))
                    #else
                        if (!this->GetCondition().WaitForSignalNsEx(&this->mutex_, iSecondTimeout))
                    #endif
                        {
                            return false;
                        }
                    }
                }
            }
        }
        while (iCurState < 0 ||
               AuAtomicCompareExchange(&this->iState_, iCurState + 1, iCurState) != iCurState);

        return true;
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::TryLockWriteMaybeSpin()
    {
        if (this->TryLockWriteNoSpin())
        {
            return true;
        }

        if (gUseFutexRWLock)
        {
            if (DoTryIf([=]()
            {
                return this->TryLockWriteNoSpin();
            }))
            {
                return true;
            }
        }

        return false;
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::LockWriteNSAbs(AuUInt64 uTimeout)
    {
        bool bRet {};

        if (this->TryLockWriteMaybeSpin())
        {
            return true;
        }

        AuAtomicAdd(&this->dwWritersPending_, 1);

        if (gUseFutexRWLock)
        {
            bRet = this->LockWriteNSAbsUnlocked(uTimeout);
        }
        else
        {
            RWLockAcquire;

            bRet = this->LockWriteNSAbsUnlocked(uTimeout);
        }

        AuAtomicSub(&this->dwWritersPending_, 1);

        return bRet;
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::TryLockWriteNoSpin()
    { 
        auto uOld = this->iState_;

        if constexpr (bIsWriteRecursionAllowed)
        {
            if (uOld < 0)
            {
                if (this->reentrantWriteLockHandle_ == GetThreadCookie())
                {
                    AuAtomicSub(&this->iState_, 1);
                    return true;
                }
            }
        }
        
        if (uOld == 0)
        {
            if (AuAtomicCompareExchange(&this->iState_, -1, uOld) == uOld)
            {
                this->reentrantWriteLockHandle_ = GetThreadCookie();
                return true;
            }
        }

        return false;
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::LockWriteNS(AuUInt64 uTimeout)
    {
        bool bRet {};
        AuInt64 uEndTime {};
 
        if (this->TryLockWriteMaybeSpin())
        {
            return true;
        }

        AuAtomicAdd(&this->dwWritersPending_, 1);

        if (gUseFutexRWLock)
        {
            uEndTime = uTimeout ? AuTime::SteadyClockNS() + uTimeout : 0;

            bRet = this->LockWriteNSAbsUnlocked(uEndTime);
        }
        else
        {
            uEndTime = uTimeout ? AuTime::SteadyClockNS() + uTimeout : 0;

            RWLockAcquire;

            bRet = this->LockWriteNSAbsUnlocked(uEndTime);
        }

        AuAtomicSub(&this->dwWritersPending_, 1);

        return bRet;
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::LockWriteNSAbsUnlocked(AuUInt64 qwTimeoutNS)
    {
        while (true)
        {
            if (this->WriterSleep(qwTimeoutNS))
            {
                if (this->WriterTryLock())
                {
                    return true;
                }
            }
            else
            {
                return this->WriterTryLock();
            }
        }
    }
    
    template<bool bIsWriteRecursionAllowed>
    void RWLockImpl<bIsWriteRecursionAllowed>::FutexWriterWake()
    {
        auto pSemaphore = this->GetFutexConditionWriter();
        while (true)
        {
            auto uState = *pSemaphore;

            if (uState == 0)
            {
                break;
            }

            if (AuAtomicCompareExchange(pSemaphore, uState - 1, uState) == uState)
            {
                break;
            }
        }
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::WriterTryLock()
    {
        if (AuAtomicCompareExchange(&this->iState_, -1, 0) == 0)
        {
            this->reentrantWriteLockHandle_ = GetThreadCookie();
            return true;
        }
        else
        {
            return false;
        }
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::WriterSleep(AuUInt64 qwTimeoutNS)
    {
        bool bStatus { true };
        AuInt32 iCurState;

        auto pSemaphore = this->GetFutexConditionWriter();

        while ((iCurState = AuAtomicLoad(&this->iState_)) != 0)
        {
            AuInt64 uSecondTimeout = 0;

            if (gUseFutexRWLock)
            {
                static const AuUInt32 kExpect { 0 };

                bStatus = InternalLTSWaitOnAddressHighRes(pSemaphore, &kExpect, sizeof(kExpect), qwTimeoutNS);

                this->WriterWake();

                if (!bStatus)
                {
                    return false;
                }

                this->FutexWriterWake();
            }
            else
            {
                if (qwTimeoutNS)
                {
                    uSecondTimeout = qwTimeoutNS - AuTime::SteadyClockNS();

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

            #if defined(AURWLOCK_NO_SIZE_OPTIMIZED_CONDVAR)
                bStatus = this->GetConditionWriter().WaitForSignalNS(uSecondTimeout);
            #else
                bStatus = this->GetConditionWriter().WaitForSignalNsEx(&this->mutex_, uSecondTimeout);
            #endif

                this->WriterWake();

                if (!bStatus)
                {
                    continue;
                }
            }
        }

        return true;
    }

    template<bool bIsWriteRecursionAllowed>
    void RWLockImpl<bIsWriteRecursionAllowed>::WriterWake()
    {
        if (AuAtomicLoad(&this->iState_) == 1 &&
            AuAtomicLoad(&this->dwWritersPending_) > 1)
        {
            this->SignalManyWriter(-1);
        }
    }

    template<bool bIsWriteRecursionAllowed>
    void RWLockImpl<bIsWriteRecursionAllowed>::SignalOneReader()
    {
        // Unused method!

        if (gUseFutexRWLock)
        {
            InternalLTSWakeOne((const void *)&this->iState_);
        }
        else
        {
            RWLockBarrierSelf();
            this->GetCondition().Signal();
        }
    }

    template<bool bIsWriteRecursionAllowed>
    void RWLockImpl<bIsWriteRecursionAllowed>::SignalOneWriter()
    {
        if (gUseFutexRWLock)
        {
            auto pThat = this->GetFutexConditionWriter();
            AuAtomicAdd(pThat, 1u);
            InternalLTSWakeOne(pThat); // Note: all paths check for a waiter ahead of time!
        }
        else
        {
            RWLockBarrierSelf();
            this->GetConditionWriter().Signal();
        }
    }

    template<bool bIsWriteRecursionAllowed>
    void RWLockImpl<bIsWriteRecursionAllowed>::SignalManyReader()
    {
        if (gUseFutexRWLock)
        {
            if (AuAtomicLoad(this->GetReadSleepCounter()))
            {
                InternalLTSWakeAll((const void *)&this->iState_);
            }
        }
        else
        {
            RWLockBarrierSelf();
            this->GetCondition().Broadcast();
        }
    }

    template<bool bIsWriteRecursionAllowed>
    void RWLockImpl<bIsWriteRecursionAllowed>::SignalManyWriter(int iBias)
    {
        if (gUseFutexRWLock)
        {
            auto pThat = this->GetFutexConditionWriter();
            if (auto uCount = AuUInt32(AuAtomicLoad(&this->dwWritersPending_) + iBias))
            {
                AuAtomicAdd(pThat, uCount);
                InternalLTSWakeCount(pThat, uCount);
            }
        }
        else
        {
            RWLockBarrierSelf();
            this->GetConditionWriter().Broadcast();
        }
    }
    
    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::TryLockRead()
    {
        if (this->TryLockReadNoSpin<true>())
        {
            return true;
        }

        if (ThrdCfg::gPreferRWLockReadLockSpin &&
            AuAtomicLoad(&this->dwWritersPending_) == 0)
        {
            return DoTryIf([=]()
            {
                return this->TryLockReadNoSpin<true>();
            });
        }

        return false;
    }

    template<bool bIsWriteRecursionAllowed>
    template<bool bCheckWrite>
    bool RWLockImpl<bIsWriteRecursionAllowed>::TryLockReadNoSpin()
    {
        auto iCurState = this->iState_;

        if (iCurState < 0)
        {
            return this->reentrantWriteLockHandle_ == GetThreadCookie();
        }

        if constexpr (bCheckWrite)
        {
            if ((AuAtomicLoad(&this->dwWritersPending_)) && 
                (iCurState > 0 || ThrdCfg::gAlwaysRWLockWriteBiasOnReadLock) &&
                (ThrdCfg::gEnableRWLockWriteBiasOnReadLock))
            {
                return false;
            }
        }

        return AuAtomicCompareExchange(&this->iState_, iCurState + 1, iCurState) == iCurState;
    }
    
    template<bool bIsWriteRecursionAllowed>
    void RWLockImpl<bIsWriteRecursionAllowed>::UnlockRead()
    {
        AuInt32 uVal {};

        if (this->iState_ < 0)
        {
            SysAssertDbg(this->reentrantWriteLockHandle_ == GetThreadCookie());
            return;
        }

        uVal = AuAtomicSub(&this->iState_, 1);

        bool bAlt {};
        if constexpr (bIsWriteRecursionAllowed)
        {
            bAlt = uVal == 1;
        }

        if (uVal == 0 || bAlt)
        {
            if (AuAtomicLoad(&this->dwWritersPending_) > 0)
            {
                this->SignalOneWriter();
            }
        }
    }

    template<bool bIsWriteRecursionAllowed>
    void RWLockImpl<bIsWriteRecursionAllowed>::UnlockWrite()
    {
        if constexpr (!bIsWriteRecursionAllowed)
        {
            this->reentrantWriteLockHandle_ = 0;

            AuAtomicStore(&this->iState_, 0);

            if (AuAtomicLoad(&this->dwWritersPending_) > 0)
            {
                this->SignalOneWriter();
            }
            else
            {
                this->SignalManyReader();
            }
        }
        else
        {
            AuInt32 val {};

            // love me cas
            {
                AuInt32 curVal {};
                do
                {
                    curVal = this->iState_;
                    if (curVal != -1)
                    {
                        continue;
                    }

                    this->reentrantWriteLockHandle_ = 0;
                }
                while (AuAtomicCompareExchange(&this->iState_, val = (curVal + 1), curVal) != curVal);
            }
            
            if (val == 0)
            {
                if (AuAtomicLoad(&this->dwWritersPending_) > 0)
                {
                    this->SignalOneWriter();
                }
                else
                {
                    this->SignalManyReader();
                }
            }
        }
    }
    
    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::UpgradeReadToWrite(AuUInt64 uTimeout)
    {
        if (this->iState_ == 1)
        {
            if (this->UpgradeReadToWriteDoUpgrade())
            {
                return true;
            }
        }

        auto uEndTime = uTimeout ? AuTime::SteadyClockNS() + uTimeout : 0;
        
        if (!gUseFutexRWLock)
        {
            AuAtomicAdd(&this->dwWritersPending_, 1);
            RWLockAcquire;
    
            while (this->iState_ != 1)
            {
                AuInt64 iSecondTimeout {};

                if (uTimeout)
                {
                    iSecondTimeout  = AuInt64(uEndTime) - AuTime::SteadyClockNS();

                    if (iSecondTimeout <= 0)
                    {
                        AuAtomicSub(&this->dwWritersPending_, 1);
                        return false;
                    }
                }

            #if defined(AURWLOCK_NO_SIZE_OPTIMIZED_CONDVAR)
                if (!this->GetConditionWriter().WaitForSignalNS(iSecondTimeout))
            #else
                if (!this->GetConditionWriter().WaitForSignalNsEx(&this->mutex_, iSecondTimeout))
            #endif
                {
                    AuAtomicSub(&this->dwWritersPending_, 1);
                    return false;
                }
            }
         
            AuAtomicSub(&this->dwWritersPending_, 1);
            return this->UpgradeReadToWriteDoUpgrade();
        }
        else
        {
            while (true)
            {
                auto pSemaphore = this->GetFutexConditionWriter();
                
                AuInt32 iCurState;
                while ((iCurState = AuAtomicLoad(&this->iState_)) != 1)
                {
                    bool bStatusTwo {};
                    bool bStatus {};

                    AuAtomicAdd(&this->dwWritersPending_, 1);
                    static const AuUInt32 kExpect { 0 };
                    if ((iCurState = AuAtomicLoad(&this->iState_)) == 1)
                    {
                        bStatus = true;
                        bStatusTwo = true;
                    }
                    else
                    {
                        bStatus = InternalLTSWaitOnAddressHighRes(pSemaphore, &kExpect, sizeof(kExpect), uEndTime);
                    }
                    AuAtomicSub(&this->dwWritersPending_, 1);

                    if (!bStatus)
                    {
                        return false;
                    }

                    if (!bStatusTwo)
                    {
                        while (true)
                        {
                            auto uState = *pSemaphore;

                            if (uState == 0)
                            {
                                break;
                            }

                            if (AuAtomicCompareExchange(pSemaphore, uState - 1, uState) == uState)
                            {
                                break;
                            }
                        }
                    }
                }

                if (this->UpgradeReadToWriteDoUpgrade())
                {
                    return true;
                }
            }
        }
    
        /* unreachable */
        return false;
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::UpgradeReadToWriteDoUpgrade()
    {
        if (AuAtomicCompareExchange(&this->iState_, -1, 1) == 1)
        {
            this->reentrantWriteLockHandle_ = GetThreadCookie();
            return true;
        }
        else
        {
            return false;
        }
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::DowngradeWriteToRead()
    {
        if (AuAtomicCompareExchange(&this->iState_, 1, -1) == -1)
        {
            this->SignalManyReader();
            return true;
        }
        else
        {
            return false;
        }
    }

    template<bool bIsWriteRecursionAllowed>
    IWaitable *RWLockImpl<bIsWriteRecursionAllowed>::AsReadable()
    {
        return &this->read_;
    }

    template<bool bIsWriteRecursionAllowed>
    IWaitable *RWLockImpl<bIsWriteRecursionAllowed>::AsWritable()
    {
        return &this->write_;
    }

    template<bool bIsWriteRecursionAllowed>
    bool RWLockImpl<bIsWriteRecursionAllowed>::CheckSelfThreadIsWriter()
    {
        return this->reentrantWriteLockHandle_ == GetThreadCookie();
    }

    AUKN_SYM IRWLock *RWLockNew()
    {
        return _new RWLockImpl<false>();
    }

    AUKN_SYM void RWLockRelease(IRWLock *pRwLock)
    {
        AuSafeDelete<RWLockImpl<false> *>(pRwLock);
    }

    AUKN_SYM IRWLock *RWRenterableLockNew()
    {
        return _new RWLockImpl<true>();
    }

    AUKN_SYM void RWRenterableLockRelease(IRWLock *pRwLock)
    {
        AuSafeDelete<RWLockImpl<true> *>(pRwLock);
    }

    AUROXTL_INTERFACE_SOO_SRC_EX(AURORA_SYMBOL_EXPORT, RWRenterableLock, RWLockImpl<true>)
    AUROXTL_INTERFACE_SOO_SRC_EX(AURORA_SYMBOL_EXPORT, RWLock, RWLockImpl<false>)
}