EventCount.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "Async/ParkingLot.h"
#include <atomic>
#include <type_traits>
 
namespace UE
{
 
template <typename CounterType>
class TEventCount;
 
template <typename CounterType>
class TEventCountToken
{
    static_assert(std::is_unsigned_v<CounterType>);
 
public:
    inline explicit operator bool() const { return !(Value & 1); }
 
private:
    CounterType Value = 1;
 
    friend TEventCount<CounterType>;
};
 
template <typename CounterType>
class TEventCount
{
    static_assert(std::is_unsigned_v<CounterType>);
 
public:
    constexpr TEventCount() = default;
 
    TEventCount(const TEventCount&) = delete;
    TEventCount& operator=(const TEventCount&) = delete;
 
    inline TEventCountToken<CounterType> PrepareWait()
    {
        TEventCountToken<CounterType> Token;
 
#if PLATFORM_CPU_X86_FAMILY
        // Doing a relaxed load here because fetch_or on x86 cannot return the previous value
        // so when we use the return value of fetch_or, it gets compiled into a compare_exchange.
        // The worst that can happen here is that we get a stale token value and end up 
        // not waiting for an iteration. We would then get the proper value on the next iteration.
        Token.Value = Count.load(std::memory_order_relaxed) & ~CounterType(1);
        Count.fetch_or(1, std::memory_order_acq_rel);
#else
        Token.Value = Count.fetch_or(1, std::memory_order_acq_rel) & ~CounterType(1);
#endif
        return Token;
    }
 
    inline void Wait(TEventCountToken<CounterType> Compare)
    {
        if ((Count.load(std::memory_order_acquire) & ~CounterType(1)) == Compare.Value)
        {
            ParkingLot::Wait(&Count, [this, Compare] { return (Count.load(std::memory_order_acquire) & ~CounterType(1)) == Compare.Value; }, nullptr);
        }
    }
 
    inline bool WaitFor(TEventCountToken<CounterType> Compare, FMonotonicTimeSpan WaitTime)
    {
        if ((Count.load(std::memory_order_acquire) & ~CounterType(1)) == Compare.Value)
        {
            const ParkingLot::FWaitState WaitState = ParkingLot::WaitFor(&Count, [this, Compare] { return (Count.load(std::memory_order_acquire) & ~CounterType(1)) == Compare.Value; }, nullptr, WaitTime);
 
            // Make sure we return true if we did wait but also if the wait was skipped because the value actually changed before we had the opportunity to wait.
            return WaitState.bDidWake || !WaitState.bDidWait;
        }
        return true;
    }
 
    inline bool WaitUntil(TEventCountToken<CounterType> Compare, FMonotonicTimePoint WaitTime)
    {
        if ((Count.load(std::memory_order_acquire) & ~CounterType(1)) == Compare.Value)
        {
            const ParkingLot::FWaitState WaitState = ParkingLot::WaitUntil(&Count, [this, Compare] { return (Count.load(std::memory_order_acquire) & ~CounterType(1)) == Compare.Value; }, nullptr, WaitTime);
 
            // Make sure we return true if we did wait but also if the wait was skipped because the value actually changed before we had the opportunity to wait.
            return WaitState.bDidWake || !WaitState.bDidWait;
        }
        return true;
    }
 
    inline void Notify()
    {
        //
        // .fetch_add(0, acq_rel) is used to have a StoreLoad barrier,
        // which we can't express in C++. That works by making the load
        // also be store (via RMW) and relying on a StoreStore barrier to
        // get the desired ordering.
        //
        // Previously, this code was:
        //   CounterType Value = Count.load(std::memory_order_relaxed);
        //
        // which had a memory re-ordering and stale values being read,
        // leading to a missed Wake and dead-locked waiter, as a result.
        //
        CounterType Value = Count.fetch_add(0, std::memory_order_acq_rel);
        if ((Value & 1) && Count.compare_exchange_strong(Value, Value + 1, std::memory_order_release))
        {
            ParkingLot::WakeAll(&Count);
        }
    }
 
    inline void NotifyWeak()
    {
        // We probably don't need the weakly consistent case anymore
        // for this weak version of notify, but will need more tests before removing.
#if PLATFORM_WEAKLY_CONSISTENT_MEMORY
        //
        // .fetch_add(0, acq_rel) is used to have a StoreLoad barrier,
        // which we can't express in C++. That works by making the load
        // also be store (via RMW) and relying on a StoreStore barrier to
        // get the desired ordering.
        //
        // Previously, this code was:
        //   CounterType Value = Count.load(std::memory_order_relaxed);
        //
        // which had a memory re-ordering and stale values being read,
        // leading to a missed Wake and dead-locked waiter, as a result.
        //
        CounterType Value = Count.fetch_add(0, std::memory_order_acq_rel);
#else
        // On x86 and other non weak memory model, the fetch_or inside PrepareWait
        // is a serializing instruction that will flush the store buffer. 
        // We can omit the expensive locked op here and just do a relaxed read.
        CounterType Value = Count.load(std::memory_order_relaxed);
#endif
        if ((Value & 1) && Count.compare_exchange_strong(Value, Value + 1, std::memory_order_release))
        {
            ParkingLot::WakeAll(&Count);
        }
    }
 
private:
    std::atomic<CounterType> Count = 0;
};
 
using FEventCount = TEventCount<uint32>;
using FEventCountToken = TEventCountToken<uint32>;
 
} // UE