IntrusiveRefCounts.h

#pragma once
 
#include <xrpl/beast/utility/instrumentation.h>
 
#include <atomic>
#include <cstddef>
#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
    useCount() 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 kStrongCountNumBits = sizeof(CountType) * 8;
    static constexpr size_t kWeakCountNumBits = kStrongCountNumBits - 2;
    using FieldType = std::uint32_t;
    static constexpr size_t kFieldTypeBits = sizeof(FieldType) * 8;
    static constexpr FieldType kOne = 1;

 
    mutable std::atomic<FieldType> refCounts_{kStrongDelta};

    static constexpr FieldType kStrongDelta = 1;

    static constexpr FieldType kWeakDelta = (kOne << kStrongCountNumBits);

    static constexpr FieldType kPartialDestroyStartedMask = (kOne << (kFieldTypeBits - 1));

    static constexpr FieldType kPartialDestroyFinishedMask = (kOne << (kFieldTypeBits - 2));

    static constexpr FieldType kTagMask = kPartialDestroyStartedMask | kPartialDestroyFinishedMask;

    static constexpr FieldType kValueMask = ~kTagMask;

    static constexpr FieldType kStrongMask = ((kOne << kStrongCountNumBits) - 1) & kValueMask;

    static constexpr FieldType kWeakMask =
        (((kOne << kWeakCountNumBits) - 1) << kStrongCountNumBits) & kValueMask;

    struct RefCountPair
    {
        CountType strong;
        CountType weak;
        FieldType partialDestroyStartedBit{0};
        FieldType partialDestroyFinishedBit{0};
        RefCountPair(FieldType v) noexcept;
        RefCountPair(CountType s, CountType w) noexcept;

        [[nodiscard]] FieldType
        combinedValue() const noexcept;
 
        static constexpr CountType kMaxStrongValue =
            static_cast<CountType>((kOne << kStrongCountNumBits) - 1);
        static constexpr CountType kMaxWeakValue =
            static_cast<CountType>((kOne << kWeakCountNumBits) - 1);
        static constexpr CountType kCheckStrongMaxValue = kMaxStrongValue - 32;
        static constexpr CountType kCheckWeakMaxValue = kMaxWeakValue - 32;
    };
};
 
inline void
IntrusiveRefCounts::addStrongRef() const noexcept
{
    refCounts_.fetch_add(kStrongDelta, std::memory_order_acq_rel);
}
 
inline void
IntrusiveRefCounts::addWeakRef() const noexcept
{
    refCounts_.fetch_add(kWeakDelta, 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 (true)
    {
        RefCountPair const prevVal{prevIntVal};
        XRPL_ASSERT(
            (prevVal.strong >= kStrongDelta),
            "xrpl::IntrusiveRefCounts::releaseStrongRef : previous ref "
            "higher than new");
        auto nextIntVal = prevIntVal - kStrongDelta;
        ReleaseStrongRefAction action = NoOp;
        if (prevVal.strong == 1)
        {
            if (prevVal.weak == 0)
            {
                action = Destroy;
            }
            else
            {
                nextIntVal |= kPartialDestroyStartedMask;
                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 & kPartialDestroyStartedMask),
                "xrpl::IntrusiveRefCounts::releaseStrongRef : not in partial "
                "destroy");
            return action;
        }
    }
}
 
inline ReleaseStrongRefAction
IntrusiveRefCounts::addWeakReleaseStrongRef() const
{
    using enum ReleaseStrongRefAction;
 
    static_assert(kWeakDelta > kStrongDelta);
    static constexpr auto kDelta = kWeakDelta - kStrongDelta;
    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 (true)
    {
        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 + kDelta;
        ReleaseStrongRefAction action = NoOp;
        if (prevVal.strong == 1)
        {
            if (prevVal.weak == 0)
            {
                action = NoOp;
            }
            else
            {
                nextIntVal |= kPartialDestroyStartedMask;
                action = PartialDestroy;
            }
        }
        if (refCounts_.compare_exchange_weak(prevIntVal, nextIntVal, std::memory_order_acq_rel))
        {
            XRPL_ASSERT(
                (!(prevIntVal & kPartialDestroyStartedMask)),
                "xrpl::IntrusiveRefCounts::addWeakReleaseStrongRef : not "
                "started partial destroy");
            return action;
        }
    }
}
 
inline ReleaseWeakRefAction
IntrusiveRefCounts::releaseWeakRef() const
{
    auto prevIntVal = refCounts_.fetch_sub(kWeakDelta, std::memory_order_acq_rel);
    RefCountPair prev = prevIntVal;
    if (prev.weak == 1 && prev.strong == 0)
    {
        if (prev.partialDestroyStartedBit == 0u)
        {
            // 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 == 0u)
        {
            // partial destroy MUST finish before running a full destroy (when
            // using weak pointers)
            refCounts_.wait(prevIntVal - kWeakDelta, 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 == 0u)
            return false;
 
        desiredValue = curValue + kStrongDelta;
    }
    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::useCount() 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 & kValueMask)), "xrpl::IntrusiveRefCounts::~IntrusiveRefCounts : count must be zero");
    auto t = v & kTagMask;
    XRPL_ASSERT((!t || t == kTagMask), "xrpl::IntrusiveRefCounts::~IntrusiveRefCounts : valid tag");
#endif
}
 
//------------------------------------------------------------------------------
 
inline IntrusiveRefCounts::RefCountPair::RefCountPair(IntrusiveRefCounts::FieldType v) noexcept
    : strong{static_cast<CountType>(v & kStrongMask)}
    , weak{static_cast<CountType>((v & kWeakMask) >> kStrongCountNumBits)}
    , partialDestroyStartedBit{v & kPartialDestroyStartedMask}
    , partialDestroyFinishedBit{v & kPartialDestroyFinishedMask}
{
    XRPL_ASSERT(
        (strong < kCheckStrongMaxValue && weak < kCheckWeakMaxValue),
        "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 < kCheckStrongMaxValue && weak < kCheckWeakMaxValue),
        "xrpl::IntrusiveRefCounts::RefCountPair(CountType, CountType) : "
        "inputs inside range");
}
 
inline IntrusiveRefCounts::FieldType
IntrusiveRefCounts::RefCountPair::combinedValue() const noexcept
{
    XRPL_ASSERT(
        (strong < kCheckStrongMaxValue && weak < kCheckWeakMaxValue),
        "xrpl::IntrusiveRefCounts::RefCountPair::combinedValue : inputs "
        "inside range");
    return (static_cast<IntrusiveRefCounts::FieldType>(weak)
            << IntrusiveRefCounts::kStrongCountNumBits) |
        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::kPartialDestroyFinishedMask);
    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