IntrusiveRefCounts.h

//------------------------------------------------------------------------------
/*
    This file is part of rippled: https://github.com/ripple/rippled
    Copyright (c) 2023 Ripple Labs Inc.
 
    Permission to use, copy, modify, and/or distribute this software for any
    purpose  with  or without fee is hereby granted, provided that the above
    copyright notice and this permission notice appear in all copies.
 
    THE  SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
    WITH  REGARD  TO  THIS  SOFTWARE  INCLUDING  ALL  IMPLIED  WARRANTIES  OF
    MERCHANTABILITY  AND  FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
    ANY  SPECIAL ,  DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
    WHATSOEVER  RESULTING  FROM  LOSS  OF USE, DATA OR PROFITS, WHETHER IN AN
    ACTION  OF  CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
    OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
*/
//==============================================================================
 
#ifndef RIPPLE_BASICS_INTRUSIVEREFCOUNTS_H_INCLUDED
#define RIPPLE_BASICS_INTRUSIVEREFCOUNTS_H_INCLUDED
 
#include <xrpl/beast/utility/instrumentation.h>
 
#include <atomic>
#include <cstdint>
 
namespace ripple {
 
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),
            "ripple::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),
                "ripple::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 `partialDestoryStartedMask`
    // 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),
            "ripple::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)),
                "ripple::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)),
        "ripple::IntrusiveRefCounts::~IntrusiveRefCounts : count must be zero");
    auto t = v & tagMask;
    XRPL_ASSERT(
        (!t || t == tagMask),
        "ripple::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),
        "ripple::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),
        "ripple::IntrusiveRefCounts::RefCountPair(CountType, CountType) : "
        "inputs inside range");
}
 
inline IntrusiveRefCounts::FieldType
IntrusiveRefCounts::RefCountPair::combinedValue() const noexcept
{
    XRPL_ASSERT(
        (strong < checkStrongMaxValue && weak < checkWeakMaxValue),
        "ripple::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 p =
        self.refCounts.fetch_or(IntrusiveRefCounts::partialDestroyFinishedMask);
    XRPL_ASSERT(
        (!p.partialDestroyFinishedBit && p.partialDestroyStartedBit &&
         !p.strong),
        "ripple::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 ripple
#endif