io_latency_probe.h

#pragma once
 
#include <xrpl/beast/utility/instrumentation.h>
 
#include <boost/asio/basic_waitable_timer.hpp>
#include <boost/asio/io_context.hpp>
#include <boost/asio/post.hpp>
 
#include <chrono>
#include <condition_variable>
#include <mutex>
#include <stdexcept>
 
namespace beast {
 
template <class Clock>
class io_latency_probe
{
private:
    using duration = typename Clock::duration;
    using time_point = typename Clock::time_point;
 
    std::recursive_mutex m_mutex;
    std::condition_variable_any m_cond;
    std::size_t m_count{1};
    duration const m_period;
    boost::asio::io_context& m_ios;
    boost::asio::basic_waitable_timer<std::chrono::steady_clock> m_timer;
    bool m_cancel{false};
 
public:
    io_latency_probe(duration const& period, boost::asio::io_context& ios)
        : m_period(period), m_ios(ios), m_timer(m_ios)
    {
    }
 
    ~io_latency_probe()
    {
        std::unique_lock<decltype(m_mutex)> lock(m_mutex);
        cancel(lock, true);
    }
 
    boost::asio::io_context&
    get_io_context()
    {
        return m_ios;
    }
 
    boost::asio::io_context const&
    get_io_context() const
    {
        return m_ios;
    }
    void
    cancel()
    {
        std::unique_lock<decltype(m_mutex)> lock(m_mutex);
        cancel(lock, true);
    }
 
    void
    cancel_async()
    {
        std::unique_lock<decltype(m_mutex)> lock(m_mutex);
        cancel(lock, false);
    }
    template <class Handler>
    void
    sample_one(Handler&& handler)
    {
        std::lock_guard const lock(m_mutex);
        if (m_cancel)
            throw std::logic_error("io_latency_probe is canceled");
        boost::asio::post(
            m_ios, sample_op<Handler>(std::forward<Handler>(handler), Clock::now(), false, this));
    }
 
    template <class Handler>
    void
    sample(Handler&& handler)
    {
        std::lock_guard const lock(m_mutex);
        if (m_cancel)
            throw std::logic_error("io_latency_probe is canceled");
        boost::asio::post(
            m_ios, sample_op<Handler>(std::forward<Handler>(handler), Clock::now(), true, this));
    }
 
private:
    void
    cancel(std::unique_lock<decltype(m_mutex)>& lock, bool wait)
    {
        if (!m_cancel)
        {
            --m_count;
            m_cancel = true;
        }
 
        if (wait)
            m_cond.wait(lock, [this] { return this->m_count == 0; });
    }
 
    void
    addref()
    {
        std::lock_guard const lock(m_mutex);
        ++m_count;
    }
 
    void
    release()
    {
        std::lock_guard const lock(m_mutex);
        if (--m_count == 0)
            m_cond.notify_all();
    }
 
    template <class Handler>
    struct sample_op
    {
        Handler m_handler;
        time_point m_start;
        bool m_repeat;
        io_latency_probe* m_probe;
 
        sample_op(
            Handler const& handler,
            time_point const& start,
            bool repeat,
            io_latency_probe* probe)
            : m_handler(handler), m_start(start), m_repeat(repeat), m_probe(probe)
        {
            XRPL_ASSERT(
                m_probe,
                "beast::io_latency_probe::sample_op::sample_op : non-null "
                "probe input");
            m_probe->addref();
        }
 
        sample_op(sample_op&& from) noexcept
            : m_handler(std::move(from.m_handler))
            , m_start(from.m_start)
            , m_repeat(from.m_repeat)
            , m_probe(from.m_probe)
        {
            XRPL_ASSERT(
                m_probe,
                "beast::io_latency_probe::sample_op::sample_op(sample_op&&) : "
                "non-null probe input");
            from.m_probe = nullptr;
        }
 
        sample_op(sample_op const&) = delete;
        sample_op
        operator=(sample_op const&) = delete;
        sample_op&
        operator=(sample_op&&) = delete;
 
        ~sample_op()
        {
            if (m_probe)
                m_probe->release();
        }
 
        void
        operator()() const
        {
            if (!m_probe)
                return;
            typename Clock::time_point const now(Clock::now());
            typename Clock::duration const elapsed(now - m_start);
 
            m_handler(elapsed);
 
            {
                std::lock_guard const lock(m_probe->m_mutex);
                if (m_probe->m_cancel)
                    return;
            }
 
            if (m_repeat)
            {
                // Calculate when we want to sample again, and
                // adjust for the expected latency.
                //
                typename Clock::time_point const when(now + m_probe->m_period - 2 * elapsed);
 
                if (when <= now)
                {
                    // The latency is too high to maintain the desired
                    // period so don't bother with a timer.
                    //
                    boost::asio::post(
                        m_probe->m_ios, sample_op<Handler>(m_handler, now, m_repeat, m_probe));
                }
                else
                {
                    m_probe->m_timer.expires_after(when - now);
                    m_probe->m_timer.async_wait(
                        sample_op<Handler>(m_handler, now, m_repeat, m_probe));
                }
            }
        }
 
        void
        operator()(boost::system::error_code const& ec)
        {
            if (!m_probe)
                return;
            typename Clock::time_point const now(Clock::now());
            boost::asio::post(
                m_probe->m_ios, sample_op<Handler>(m_handler, now, m_repeat, m_probe));
        }
    };
};
 
}  // namespace beast