AuroraRuntime/Include/Aurora/Threading/Waitables/FutexCondWaitable.hpp

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

    File: FutexCondWaitable.hpp
    Date: 2023-08-19
    Author: Reece
***/
#pragma once

#include "../WakeOnAddress.hpp"
#include "../SpinTime.hpp"

namespace Aurora::Threading::Waitables
{
    struct FutexCondWaitable final 
    {
        inline constexpr FutexCondWaitable()
        {}
        AU_MOVE(FutexCondWaitable)
        inline constexpr FutexCondWaitable(const FutexCondWaitable &that)
        {
            // NOP
        }
        AU_OPERATOR_COPY(FutexCondWaitable)

        template <class T>
        inline bool WaitNS(T pWaitable, AuUInt64 uRelativeNanoseconds)
        {
            AuAtomicAdd(&this->uAtomicSleeping, 1u);

            if (pWaitable)
            {
                pWaitable->Unlock();
            }

            auto bSuccess = this->TryLock2() || 
                this->SleepOne(uRelativeNanoseconds ? Time::SteadyClockNS() + uRelativeNanoseconds : 0);

            if (!bSuccess)
            {
                auto uWaiters = this->uAtomicSleeping;
                auto uWaitCount = 1;

                while (uWaiters &&
                       AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - uWaitCount, uWaiters) != uWaiters)
                {
                    uWaiters = this->uAtomicSleeping;

                    if (uWaiters == 0)
                    {
                        break;
                    }
                }
            }

            if (pWaitable)
            {
                pWaitable->Lock();
            }

            return bSuccess;
        }

        template <class T>
        inline bool WaitAbsNS(T pWaitable, AuUInt64 uSteadyClockNanoseconds)
        {
            AuAtomicAdd(&this->uAtomicSleeping, 1u);

            if (pWaitable)
            {
                pWaitable->Unlock();
            }

            auto bSuccess = this->TryLock2() || 
                this->SleepOne(uSteadyClockNanoseconds);

            if (!bSuccess)
            {
                auto uWaiters = this->uAtomicSleeping;
                auto uWaitCount = 1;

                while (uWaiters &&
                       AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - uWaitCount, uWaiters) != uWaiters)
                {
                    uWaiters = this->uAtomicSleeping;

                    if (uWaiters == 0)
                    {
                        break;
                    }
                }
            }

            if (pWaitable)
            {
                pWaitable->Lock();
            }

            return bSuccess;
        }

        inline auline void NotifyOne()
        {
            AuUInt32 uWaiters = AuAtomicLoad(&this->uAtomicSleeping);
            if (uWaiters == 0)
            {
                return;
            }

            AuAtomicAdd(&this->uAtomicState, 1u);
            WakeOnAddress((const void *)&this->uAtomicState);

            while (AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - 1u, uWaiters) != uWaiters)
            {
                uWaiters = this->uAtomicSleeping;

                if (uWaiters == 0)
                {
                    return;
                }
            }
        }

        inline auline void Broadcast()
        {
            AuUInt32 uWaitCount {};
            AuUInt32 uWaiters {};

        #if defined(AURORA_RUNTIME_FUTEX_AGGRESSIVE_COND_WAKE)
            while ((uWaiters = AuAtomicLoad(&this->uAtomicSleeping)))
        #else
            if ((uWaiters = AuAtomicLoad(&this->uAtomicSleeping)))
        #endif
            {
                AuAtomicAdd(&this->uAtomicState, uWaiters);
                WakeNOnAddress((const void *)&this->uAtomicState, uWaiters);
                uWaitCount = uWaiters;

                while (uWaiters &&
                       AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - uWaitCount, uWaiters) != uWaiters)
                {
                    uWaiters = this->uAtomicSleeping;

                    if (uWaiters <= uWaitCount)
                    {
                        uWaitCount = uWaiters;
                    }
                }
            }
        }

        inline auline void NotifyN(AuUInt32 uThreads)
        {
            AuUInt32 uWaitCount {};
            AuUInt32 uWaiters {};

            uWaiters = AuAtomicLoad(&this->uAtomicSleeping);
            if (uWaiters > 0)
            {
                auto uMin = AuMin(uWaiters, uThreads);
                AuAtomicAdd(&this->uAtomicState, uMin);
                WakeNOnAddress((const void *)&this->uAtomicState, uMin);
                uWaitCount = uMin;
            }

            while (uWaiters &&
                   AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - uWaitCount, uWaiters) != uWaiters)
            {
                uWaiters = this->uAtomicSleeping;

                if (uWaiters <= uWaitCount)
                {
                    uWaitCount = uWaiters;
                }
            }
        }

        AuAUInt32 uAtomicState {};
        AuAUInt32 uAtomicSleeping {};
    private:

        inline auline bool TryLock3()
        {
            auto old = AuAtomicLoad(&this->uAtomicState);
            return ((old != 0 && AuAtomicCompareExchange(&this->uAtomicState, old - 1, old) == old));
        }

        inline auline bool TryLock2()
        {
            static const AuUInt32 kRef { 0 };

            if (TryLock3())
            {
                return true;
            }

            return TryWaitOnAddressEx((const void *)&this->uAtomicState,
                                      &kRef,
                                      sizeof(kRef),
                                      [&](const void *pTargetAddress,
                                          const void *pCompareAddress,
                                          AuUInt8 uWordSize)
            {
                return this->TryLock3();
            });
        }
        
        inline auline bool SleepOne(AuUInt64 qwTimeout)
        {
            static const AuUInt32 kRef { 0 };

            while (!TryLock2())
            {
                bool bStatus {};

                bStatus = WaitOnAddressSteady((const void *)&this->uAtomicState, &kRef, sizeof(kRef), qwTimeout, true);

                if (!bStatus)
                {
                    return TryLock2();
                }
            }

            return true;
        }
    };
}