WorkQueue.hpp

#pragma once
 
#include "util/Mutex.hpp"
#include "util/config/ConfigDefinition.hpp"
#include "util/log/Logger.hpp"
#include "util/prometheus/Counter.hpp"
#include "util/prometheus/Gauge.hpp"
 
#include <boost/asio/spawn.hpp>
#include <boost/asio/thread_pool.hpp>
#include <boost/json/object.hpp>
 
#include <atomic>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <limits>
#include <optional>
#include <queue>
#include <utility>
 
namespace rpc {

struct Reportable {
    virtual ~Reportable() = default;

    [[nodiscard]] virtual boost::json::object
    report() const = 0;
};

class WorkQueue : public Reportable {
    using TaskType = std::function<void(boost::asio::yield_context)>;
 
    struct TaskWithTimestamp {
        TaskType task;
        std::chrono::system_clock::time_point queuedAt;
    };
 
    using QueueType = std::queue<TaskWithTimestamp>;
 
public:
    enum class Priority : uint8_t {
        High,
        Default,
    };
 
private:
    struct QueueState {
        QueueType high;
        QueueType normal;
 
        size_t highPriorityCounter = 0;
 
        void
        push(Priority priority, TaskType&& task)
        {
            auto& queue = [this, priority] -> QueueType& {
                if (priority == Priority::High)
                    return high;
                return normal;
            }();
            queue.push(
                TaskWithTimestamp{
                    .task = std::move(task), .queuedAt = std::chrono::system_clock::now()
                }
            );
        }
 
        [[nodiscard]] bool
        empty() const
        {
            return high.empty() and normal.empty();
        }
 
        [[nodiscard]] std::optional<TaskWithTimestamp>
        popNext()
        {
            if (not high.empty() and (highPriorityCounter < kTAKE_HIGH_PRIO or normal.empty())) {
                auto taskWithTimestamp = std::move(high.front());
                high.pop();
                ++highPriorityCounter;
                return taskWithTimestamp;
            }
 
            if (not normal.empty()) {
                auto taskWithTimestamp = std::move(normal.front());
                normal.pop();
                highPriorityCounter = 0;
                return taskWithTimestamp;
            }
 
            return std::nullopt;
        }
    };
 
private:
    static constexpr auto kTAKE_HIGH_PRIO = 4uz;
 
    // these are cumulative for the lifetime of the process
    std::reference_wrapper<util::prometheus::CounterInt> queued_;
    std::reference_wrapper<util::prometheus::CounterInt> durationUs_;
 
    std::reference_wrapper<util::prometheus::GaugeInt> curSize_;
    uint32_t maxSize_ = std::numeric_limits<uint32_t>::max();
 
    util::Logger log_{"RPC"};
    boost::asio::thread_pool ioc_;
 
    std::atomic_bool stopping_;
    std::atomic_bool processingStarted_{false};
 
    class OneTimeCallable {
        std::function<void()> func_;
        bool called_{false};
 
    public:
        void
        setCallable(std::function<void()> func);
 
        void
        operator()();
 
        explicit
        operator bool() const;
    };
    util::Mutex<OneTimeCallable> onQueueEmpty_;
    util::Mutex<QueueState> queueState_;
 
public:
    struct DontStartProcessingTag {};
    static constexpr DontStartProcessingTag kDONT_START_PROCESSING_TAG = {};

    WorkQueue(std::uint32_t numWorkers, uint32_t maxSize = 0);

    WorkQueue(DontStartProcessingTag, std::uint32_t numWorkers, uint32_t maxSize = 0);
 
    ~WorkQueue() override;

    void
    startProcessing();

    void
    requestStop(std::function<void()> onQueueEmpty = [] {});

    void
    stop();

    [[nodiscard]] static WorkQueue
    makeWorkQueue(util::config::ClioConfigDefinition const& config);

    bool
    postCoro(TaskType func, bool isWhiteListed, Priority priority = Priority::Default);

    [[nodiscard]] boost::json::object
    report() const override;

    void
    join();

    [[nodiscard]] size_t
    size() const;
 
private:
    void
    executeTask(boost::asio::yield_context yield);
};
 
}  // namespace rpc