IntrusiveRefCounts.h

#pragma once
 
#include <xrpl/beast/utility/instrumentation.h>
 
#include <atomic>
#include <cstdint>
 
namespace xrpl {
 
enum class ReleaseStrongRefAction { noop, partialDestroy, destroy };
 
enum class ReleaseWeakRefAction { noop, destroy };
 
struct IntrusiveRefCounts
{
    virtual ~IntrusiveRefCounts() noexcept;
 
    // This must be `noexcept` or the make_SharedIntrusive function could leak
    // memory.
    void
    addStrongRef() const noexcept;
 
    void
    addWeakRef() const noexcept;
 
    ReleaseStrongRefAction
    releaseStrongRef() const;
 
    // Same as:
    // {
    //   addWeakRef();
    //   return releaseStrongRef;
    // }
    // done as one atomic operation
    ReleaseStrongRefAction
    addWeakReleaseStrongRef() const;
 
    ReleaseWeakRefAction
    releaseWeakRef() const;
 
    // Returns true is able to checkout a strong ref. False otherwise
    bool
    checkoutStrongRefFromWeak() const noexcept;
 
    bool
    expired() const noexcept;
 
    std::size_t
    use_count() const noexcept;
 
    // This function MUST be called after a partial destructor finishes running.
    // Calling this function may cause other threads to delete the object
    // pointed to by `o`, so `o` should never be used after calling this
    // function. The parameter will be set to a `nullptr` after calling this
    // function to emphasize that it should not be used.
    // Note: This is intentionally NOT called at the end of `partialDestructor`.
    // The reason for this is if new classes are written to support this smart
    // pointer class, they need to write their own `partialDestructor` function
    // and ensure `partialDestructorFinished` is called at the end. Putting this
    // call inside the smart pointer class itself is expected to be less error
    // prone.
    // Note: The "two-star" programming is intentional. It emphasizes that `o`
    // may be deleted and the unergonomic API is meant to signal the special
    // nature of this function call to callers.
    // Note: This is a template to support incompletely defined classes.
    template <class T>
    friend void
    partialDestructorFinished(T** o);
 
private:
    // TODO: We may need to use a uint64_t for both counts. This will reduce the
    // memory savings. We need to audit the code to make sure 16 bit counts are
    // enough for strong pointers and 14 bit counts are enough for weak
    // pointers. Use type aliases to make it easy to switch types.
    using CountType = std::uint16_t;
    static constexpr size_t StrongCountNumBits = sizeof(CountType) * 8;
    static constexpr size_t WeakCountNumBits = StrongCountNumBits - 2;
    using FieldType = std::uint32_t;
    static constexpr size_t FieldTypeBits = sizeof(FieldType) * 8;
    static constexpr FieldType one = 1;
 
    mutable std::atomic<FieldType> refCounts{strongDelta};
 
    static constexpr FieldType strongDelta = 1;
 
    static constexpr FieldType weakDelta = (one << StrongCountNumBits);
 
    static constexpr FieldType partialDestroyStartedMask = (one << (FieldTypeBits - 1));
 
    static constexpr FieldType partialDestroyFinishedMask = (one << (FieldTypeBits - 2));
 
    static constexpr FieldType tagMask = partialDestroyStartedMask | partialDestroyFinishedMask;
 
    static constexpr FieldType valueMask = ~tagMask;
 
    static constexpr FieldType strongMask = ((one << StrongCountNumBits) - 1) & valueMask;
 
    static constexpr FieldType weakMask =
        (((one << WeakCountNumBits) - 1) << StrongCountNumBits) & valueMask;
 
    struct RefCountPair
    {
        CountType strong;
        CountType weak;
        FieldType partialDestroyStartedBit{0};
        FieldType partialDestroyFinishedBit{0};
        RefCountPair(FieldType v) noexcept;
        RefCountPair(CountType s, CountType w) noexcept;
 
        FieldType
        combinedValue() const noexcept;
 
        static constexpr CountType maxStrongValue =
            static_cast<CountType>((one << StrongCountNumBits) - 1);
        static constexpr CountType maxWeakValue =
            static_cast<CountType>((one << WeakCountNumBits) - 1);
        static constexpr CountType checkStrongMaxValue = maxStrongValue - 32;
        static constexpr CountType checkWeakMaxValue = maxWeakValue - 32;
    };
};
 
inline void
IntrusiveRefCounts::addStrongRef() const noexcept
{
    refCounts.fetch_add(strongDelta, std::memory_order_acq_rel);
}
 
inline void
IntrusiveRefCounts::addWeakRef() const noexcept
{
    refCounts.fetch_add(weakDelta, std::memory_order_acq_rel);
}
 
inline ReleaseStrongRefAction
IntrusiveRefCounts::releaseStrongRef() const
{
    // Subtract `strongDelta` from refCounts. If this releases the last strong
    // ref, set the `partialDestroyStarted` bit. It is important that the ref
    // count and the `partialDestroyStartedBit` are changed atomically (hence
    // the loop and `compare_exchange` op). If this didn't need to be done
    // atomically, the loop could be replaced with a `fetch_sub` and a
    // conditional `fetch_or`. This loop will almost always run once.
 
    using enum ReleaseStrongRefAction;
    auto prevIntVal = refCounts.load(std::memory_order_acquire);
    while (1)
    {
        RefCountPair const prevVal{prevIntVal};
        XRPL_ASSERT(
            (prevVal.strong >= strongDelta),
            "xrpl::IntrusiveRefCounts::releaseStrongRef : previous ref "
            "higher than new");
        auto nextIntVal = prevIntVal - strongDelta;
        ReleaseStrongRefAction action = noop;
        if (prevVal.strong == 1)
        {
            if (prevVal.weak == 0)
            {
                action = destroy;
            }
            else
            {
                nextIntVal |= partialDestroyStartedMask;
                action = partialDestroy;
            }
        }
 
        if (refCounts.compare_exchange_weak(prevIntVal, nextIntVal, std::memory_order_acq_rel))
        {
            // Can't be in partial destroy because only decrementing the strong
            // count to zero can start a partial destroy, and that can't happen
            // twice.
            XRPL_ASSERT(
                (action == noop) || !(prevIntVal & partialDestroyStartedMask),
                "xrpl::IntrusiveRefCounts::releaseStrongRef : not in partial "
                "destroy");
            return action;
        }
    }
}
 
inline ReleaseStrongRefAction
IntrusiveRefCounts::addWeakReleaseStrongRef() const
{
    using enum ReleaseStrongRefAction;
 
    static_assert(weakDelta > strongDelta);
    auto constexpr delta = weakDelta - strongDelta;
    auto prevIntVal = refCounts.load(std::memory_order_acquire);
    // This loop will almost always run once. The loop is needed to atomically
    // change the counts and flags (the count could be atomically changed, but
    // the flags depend on the current value of the counts).
    //
    // Note: If this becomes a perf bottleneck, the `partialDestroyStartedMask`
    // may be able to be set non-atomically. But it is easier to reason about
    // the code if the flag is set atomically.
    while (1)
    {
        RefCountPair const prevVal{prevIntVal};
        // Converted the last strong pointer to a weak pointer.
        //
        // Can't be in partial destroy because only decrementing the
        // strong count to zero can start a partial destroy, and that
        // can't happen twice.
        XRPL_ASSERT(
            (!prevVal.partialDestroyStartedBit),
            "xrpl::IntrusiveRefCounts::addWeakReleaseStrongRef : not in "
            "partial destroy");
 
        auto nextIntVal = prevIntVal + delta;
        ReleaseStrongRefAction action = noop;
        if (prevVal.strong == 1)
        {
            if (prevVal.weak == 0)
            {
                action = noop;
            }
            else
            {
                nextIntVal |= partialDestroyStartedMask;
                action = partialDestroy;
            }
        }
        if (refCounts.compare_exchange_weak(prevIntVal, nextIntVal, std::memory_order_acq_rel))
        {
            XRPL_ASSERT(
                (!(prevIntVal & partialDestroyStartedMask)),
                "xrpl::IntrusiveRefCounts::addWeakReleaseStrongRef : not "
                "started partial destroy");
            return action;
        }
    }
}
 
inline ReleaseWeakRefAction
IntrusiveRefCounts::releaseWeakRef() const
{
    auto prevIntVal = refCounts.fetch_sub(weakDelta, std::memory_order_acq_rel);
    RefCountPair prev = prevIntVal;
    if (prev.weak == 1 && prev.strong == 0)
    {
        if (!prev.partialDestroyStartedBit)
        {
            // This case should only be hit if the partialDestroyStartedBit is
            // set non-atomically (and even then very rarely). The code is kept
            // in case we need to set the flag non-atomically for perf reasons.
            refCounts.wait(prevIntVal, std::memory_order_acquire);
            prevIntVal = refCounts.load(std::memory_order_acquire);
            prev = RefCountPair{prevIntVal};
        }
        if (!prev.partialDestroyFinishedBit)
        {
            // partial destroy MUST finish before running a full destroy (when
            // using weak pointers)
            refCounts.wait(prevIntVal - weakDelta, std::memory_order_acquire);
        }
        return ReleaseWeakRefAction::destroy;
    }
    return ReleaseWeakRefAction::noop;
}
 
inline bool
IntrusiveRefCounts::checkoutStrongRefFromWeak() const noexcept
{
    auto curValue = RefCountPair{1, 1}.combinedValue();
    auto desiredValue = RefCountPair{2, 1}.combinedValue();
 
    while (!refCounts.compare_exchange_weak(curValue, desiredValue, std::memory_order_acq_rel))
    {
        RefCountPair const prev{curValue};
        if (!prev.strong)
            return false;
 
        desiredValue = curValue + strongDelta;
    }
    return true;
}
 
inline bool
IntrusiveRefCounts::expired() const noexcept
{
    RefCountPair const val = refCounts.load(std::memory_order_acquire);
    return val.strong == 0;
}
 
inline std::size_t
IntrusiveRefCounts::use_count() const noexcept
{
    RefCountPair const val = refCounts.load(std::memory_order_acquire);
    return val.strong;
}
 
inline IntrusiveRefCounts::~IntrusiveRefCounts() noexcept
{
#ifndef NDEBUG
    auto v = refCounts.load(std::memory_order_acquire);
    XRPL_ASSERT(
        (!(v & valueMask)), "xrpl::IntrusiveRefCounts::~IntrusiveRefCounts : count must be zero");
    auto t = v & tagMask;
    XRPL_ASSERT((!t || t == tagMask), "xrpl::IntrusiveRefCounts::~IntrusiveRefCounts : valid tag");
#endif
}
 
//------------------------------------------------------------------------------
 
inline IntrusiveRefCounts::RefCountPair::RefCountPair(IntrusiveRefCounts::FieldType v) noexcept
    : strong{static_cast<CountType>(v & strongMask)}
    , weak{static_cast<CountType>((v & weakMask) >> StrongCountNumBits)}
    , partialDestroyStartedBit{v & partialDestroyStartedMask}
    , partialDestroyFinishedBit{v & partialDestroyFinishedMask}
{
    XRPL_ASSERT(
        (strong < checkStrongMaxValue && weak < checkWeakMaxValue),
        "xrpl::IntrusiveRefCounts::RefCountPair(FieldType) : inputs inside "
        "range");
}
 
inline IntrusiveRefCounts::RefCountPair::RefCountPair(
    IntrusiveRefCounts::CountType s,
    IntrusiveRefCounts::CountType w) noexcept
    : strong{s}, weak{w}
{
    XRPL_ASSERT(
        (strong < checkStrongMaxValue && weak < checkWeakMaxValue),
        "xrpl::IntrusiveRefCounts::RefCountPair(CountType, CountType) : "
        "inputs inside range");
}
 
inline IntrusiveRefCounts::FieldType
IntrusiveRefCounts::RefCountPair::combinedValue() const noexcept
{
    XRPL_ASSERT(
        (strong < checkStrongMaxValue && weak < checkWeakMaxValue),
        "xrpl::IntrusiveRefCounts::RefCountPair::combinedValue : inputs "
        "inside range");
    return (static_cast<IntrusiveRefCounts::FieldType>(weak)
            << IntrusiveRefCounts::StrongCountNumBits) |
        static_cast<IntrusiveRefCounts::FieldType>(strong) | partialDestroyStartedBit |
        partialDestroyFinishedBit;
}
 
template <class T>
inline void
partialDestructorFinished(T** o)
{
    T& self = **o;
    IntrusiveRefCounts::RefCountPair const p =
        self.refCounts.fetch_or(IntrusiveRefCounts::partialDestroyFinishedMask);
    XRPL_ASSERT(
        (!p.partialDestroyFinishedBit && p.partialDestroyStartedBit && !p.strong),
        "xrpl::partialDestructorFinished : not a weak ref");
    if (!p.weak)
    {
        // There was a weak count before the partial destructor ran (or we would
        // have run the full destructor) and now there isn't a weak count. Some
        // thread is waiting to run the destructor.
        self.refCounts.notify_one();
    }
    // Set the pointer to null to emphasize that the object shouldn't be used
    // after calling this function as it may be destroyed in another thread.
    *o = nullptr;
}
//------------------------------------------------------------------------------
 
}  // namespace xrpl