libdpf/include/dpf/net/io_pool.hpp

/// @file dpf/net/io_pool.hpp
/// @brief Socket completion threads and a separate compute pool.
/// @details `context()` is run by `io_threads` workers (0 = hardware
///          concurrency). `post_compute` runs on `compute_threads` workers
///          (0 = same as io), so a DPF evaluation never occupies a thread that
///          should be completing reads and writes. Shutdown stops socket I/O
///          first, then joins compute.
#ifndef LIBDPF_INCLUDE_DPF_NET_IO_POOL_HPP__
#define LIBDPF_INCLUDE_DPF_NET_IO_POOL_HPP__

#include <cstddef>
#include <thread>
#include <utility>
#include <vector>

#include "dpf/net/asio_ns.hpp"
#include <asio/thread_pool.hpp>

namespace dpf
{
namespace net
{

class io_pool
{
public:
    explicit io_pool(std::size_t io_threads = 0, std::size_t compute_threads = 0)
      : work_(asio::make_work_guard(io_)),
        compute_n_(pick(compute_threads == 0 ? io_threads : compute_threads)),
        compute_(compute_n_)
    {
        const std::size_t n = pick(io_threads);
        threads_.reserve(n);
        for (std::size_t i = 0; i < n; ++i)
            threads_.emplace_back([this] { io_.run(); });
    }

    io_pool(const io_pool &) = delete;
    io_pool & operator=(const io_pool &) = delete;

    ~io_pool()
    {
        work_.reset();
        io_.stop();
        for (auto & t : threads_)
            if (t.joinable())
                t.join();
        compute_.join();
    }

    asio::io_context & context() noexcept { return io_; }
    std::size_t size() const noexcept { return threads_.size(); }
    std::size_t compute_size() const noexcept { return compute_n_; }

    /// @brief Run `fn` on a socket worker.
    template <typename Fn>
    void post(Fn && fn)
    {
        asio::post(io_, std::forward<Fn>(fn));
    }

    /// @brief Run `fn` on the compute pool, leaving socket workers free.
    template <typename Fn>
    void post_compute(Fn && fn)
    {
        asio::post(compute_, std::forward<Fn>(fn));
    }

private:
    static std::size_t pick(std::size_t n)
    {
        if (n != 0)
            return n;
        const auto hw = std::thread::hardware_concurrency();
        return hw == 0 ? 1 : hw;
    }

    asio::io_context io_;
    asio::executor_work_guard<asio::io_context::executor_type> work_;
    std::size_t compute_n_ = 1;
    asio::thread_pool compute_;
    std::vector<std::thread> threads_;
};

}  // namespace net
}  // namespace dpf

#endif