rippled/include/xrpl/server/detail/io_list.h

#pragma once

#include <boost/container/flat_map.hpp>

#include <condition_variable>
#include <cstddef>
#include <functional>
#include <memory>
#include <mutex>
#include <type_traits>
#include <utility>

namespace xrpl {

/**
 * Manages a set of objects performing asynchronous I/O.
 */
class IOList final
{
public:
    class Work
    {
        template <class = void>
        void
        destroy();

        friend class IOList;
        IOList* ios_ = nullptr;

    public:
        virtual ~Work()
        {
            destroy();
        }

        /**
         * Return the IOList associated with the work.
         *
         * Requirements:
         *     The call to IOList::emplace to
         *     create the work has already returned.
         */
        IOList&
        ios()
        {
            return *ios_;
        }

        virtual void
        close() = 0;
    };

private:
    template <class = void>
    void
    destroy();

    std::mutex m_;
    std::size_t n_ = 0;
    bool closed_ = false;
    std::condition_variable cv_;
    boost::container::flat_map<Work*, std::weak_ptr<Work>> map_;
    std::function<void(void)> f_;

public:
    IOList() = default;

    /**
     * Destroy the list.
     *
     * Effects:
     *     Closes the IOList if it was not previously
     *         closed. No finisher is invoked in this case.
     *
     *     Blocks until all work is destroyed.
     */
    ~IOList()
    {
        destroy();
    }

    /**
     * Return `true` if the list is closed.
     *
     * Thread Safety:
     *     Undefined result if called concurrently
     *     with close().
     */
    [[nodiscard]] bool
    closed() const
    {
        return closed_;
    }

    /**
     * Create associated work if not closed.
     *
     * Requirements:
     *     `std::is_base_of_v<Work, T> == true`
     *
     * Thread Safety:
     *     May be called concurrently.
     *
     * Effects:
     *     Atomically creates, inserts, and returns new
     *     work T, or returns nullptr if the io_list is
     *     closed,
     *
     * If the call succeeds and returns a new object,
     * it is guaranteed that a subsequent call to close
     * will invoke Work::close on the object.
     */
    template <class T, class... Args>
    std::shared_ptr<T>
    emplace(Args&&... args);

    /**
     * Cancel active I/O.
     *
     * Thread Safety:
     *     May not be called concurrently.
     *
     * Effects:
     *     Associated work is closed.
     *
     *     Finisher if provided, will be called when
     *     all associated work is destroyed. The finisher
     *     may be called from a foreign thread, or within
     *     the call to this function.
     *
     *     Only the first call to close will set the
     *     finisher.
     *
     *     No effect after the first call.
     */
    template <class Finisher>
    void
    close(Finisher&& f);

    void
    close()
    {
        close([] {});
    }

    /**
     * Block until the io_list stops.
     *
     * Effects:
     *     The caller is blocked until the io_list is
     *     closed and all associated work is destroyed.
     *
     * Thread safety:
     *     May be called concurrently.
     *
     * Preconditions:
     *     No call to io_context::run on any io_context
     *     used by work objects associated with this io_list
     *     exists in the caller's call stack.
     */
    template <class = void>
    void
    join();
};

//------------------------------------------------------------------------------

template <class>
void
IOList::Work::destroy()
{
    if (!ios_)
        return;
    std::function<void(void)> f;
    {
        std::scoped_lock const lock(ios_->m_);
        ios_->map_.erase(this);
        if (--ios_->n_ == 0 && ios_->closed_)
        {
            std::swap(f, ios_->f_);
            ios_->cv_.notify_all();
        }
    }
    if (f)
        f();
}

template <class>
void
IOList::destroy()
{
    close();
    join();
}

template <class T, class... Args>
std::shared_ptr<T>
IOList::emplace(Args&&... args)
{
    static_assert(std::is_base_of_v<Work, T>, "T must derive from IOList::Work");
    if (closed_)
        return nullptr;
    auto sp = std::make_shared<T>(std::forward<Args>(args)...);
    decltype(sp) dead;

    std::scoped_lock const lock(m_);
    if (!closed_)
    {
        ++n_;
        sp->Work::ios_ = this;
        map_.emplace(sp.get(), sp);
    }
    else
    {
        std::swap(sp, dead);
    }
    return sp;
}

template <class Finisher>
void
IOList::close(Finisher&& f)
{
    std::unique_lock<std::mutex> lock(m_);
    if (closed_)
        return;
    closed_ = true;
    auto map = std::move(map_);
    if (!map.empty())
    {
        f_ = std::forward<Finisher>(f);
        lock.unlock();
        for (auto const& p : map)
        {
            if (auto sp = p.second.lock())
                sp->close();
        }
    }
    else
    {
        lock.unlock();
        f();
    }
}

template <class>
void
IOList::join()
{
    std::unique_lock<std::mutex> lock(m_);
    cv_.wait(lock, [&] { return closed_ && n_ == 0; });
}

}  // namespace xrpl