ClosureCounter.h

#pragma once
 
#include <xrpl/basics/Log.h>
#include <xrpl/beast/utility/Journal.h>
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <mutex>
#include <optional>
#include <type_traits>
 
namespace xrpl {

template <typename Ret, typename... Args>
class ClosureCounter
{
private:
    std::mutex mutable mutex_;
    std::condition_variable allClosuresDoneCond_;  // guard with mutex_
    bool waitForClosures_{false};                  // guard with mutex_
    std::atomic<int> closureCount_{0};
 
    // Increment the count.
    ClosureCounter&
    operator++()
    {
        ++closureCount_;
        return *this;
    }
 
    // Decrement the count.  If we're stopping and the count drops to zero
    // notify allClosuresDoneCond_.
    ClosureCounter&
    operator--()
    {
        // Even though closureCount_ is atomic, we decrement its value under
        // a lock.  This removes a small timing window that occurs if the
        // waiting thread is handling a spurious wakeup when closureCount_
        // drops to zero.
        std::scoped_lock const lock{mutex_};
 
        // Update closureCount_.  Notify if stopping and closureCount_ == 0.
        if ((--closureCount_ == 0) && waitForClosures_)
            allClosuresDoneCond_.notify_all();
        return *this;
    }
 
    // A private template class that helps count the number of closures
    // in flight. This allows callers to block until all their postponed
    // closures are dispatched.
    template <typename Closure>
    class Substitute
    {
    private:
        ClosureCounter& counter_;
        std::remove_reference_t<Closure> closure_{};
 
        static_assert(
            std::is_same_v<decltype(closure_(std::declval<Args>()...)), Ret>,
            "Closure arguments don't match ClosureCounter Ret or Args");
 
    public:
        Substitute() = delete;
 
        Substitute(Substitute const& rhs) : counter_(rhs.counter_), closure_(rhs.closure_)
        {
            ++counter_;
        }
 
        Substitute(Substitute&& rhs) noexcept(std::is_nothrow_move_constructible_v<Closure>)
            : counter_(rhs.counter_), closure_(std::move(rhs.closure_))
        {
            ++counter_;
        }
 
        Substitute(
            ClosureCounter& counter,
            Closure&& closure)  // NOLINT(cppcoreguidelines-rvalue-reference-param-not-moved)
            : counter_(counter), closure_(std::forward<Closure>(closure))
        {
            ++counter_;
        }
 
        Substitute&
        operator=(Substitute const& rhs) = delete;
        Substitute&
        operator=(Substitute&& rhs) = delete;
 
        ~Substitute()
        {
            --counter_;
        }
 
        // Note that Args is not deduced, it is explicit.  So Args&&
        // would be an rvalue reference, not a forwarding reference.  We
        // want to forward exactly what the user declared.
        Ret
        operator()(Args... args)
        {
            return closure_(std::forward<Args>(args)...);
        }
    };
 
public:
    ClosureCounter() = default;
    // Not copyable or movable.  Outstanding counts would be hard to sort out.
    ClosureCounter(ClosureCounter const&) = delete;
 
    ClosureCounter&
    operator=(ClosureCounter const&) = delete;

    ~ClosureCounter()
    {
        using namespace std::chrono_literals;
        join("ClosureCounter", 1s, debugLog());
    }

    void
    join(char const* name, std::chrono::milliseconds wait, beast::Journal j)
    {
        std::unique_lock<std::mutex> lock{mutex_};
        waitForClosures_ = true;
        if (closureCount_ > 0)
        {
            if (!allClosuresDoneCond_.wait_for(lock, wait, [this] { return closureCount_ == 0; }))
            {
                if (auto stream = j.error())
                    stream << name << " waiting for ClosureCounter::join().";
                allClosuresDoneCond_.wait(lock, [this] { return closureCount_ == 0; });
            }
        }
    }

    template <class Closure>
    std::optional<Substitute<Closure>>
    wrap(Closure&& closure)
    {
        std::optional<Substitute<Closure>> ret;
 
        std::scoped_lock const lock{mutex_};
        if (!waitForClosures_)
            ret.emplace(*this, std::forward<Closure>(closure));
 
        return ret;
    }

    int
    count() const
    {
        return closureCount_;
    }

    bool
    joined() const
    {
        std::scoped_lock const lock{mutex_};
        return waitForClosures_;
    }
};
 
}  // namespace xrpl