DispatchPool.hpp

/**
 * @file DispatchPool.hpp
 * @brief Fixed-size thread pool. Offloads all blocking UDS connect/send/recv
 *        from the epoll reactor thread, keeping the event loop non-blocking.
 *
 * CSCI 599: Network Systems for Cloud Computing
 * University of Southern California
 */

#ifndef __DISPATCH_POOL_HPP__
#define __DISPATCH_POOL_HPP__

#include <functional>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <queue>
#include <vector>
#include <atomic>
#include "log.hpp"

// DispatchPool: Fixed-size thread pool for handling blocking UDS I/O.
// The Reactor thread pushes tasks here and returns immediately.
// Pool threads do all the blocking connect/send/recv work.
class DispatchPool {
public:
    static const int DEFAULT_THREADS = 64;

private:
    std::vector<std::thread>              _threads;
    std::queue<std::function<void()>>     _tasks;
    std::mutex                            _mtx;
    std::condition_variable               _cv;
    std::atomic<bool>                     _stop;

public:
    explicit DispatchPool(int nthreads = DEFAULT_THREADS) : _stop(false) {
        for (int i = 0; i < nthreads; ++i) {
            _threads.emplace_back([this]() {
                while (true) {
                    std::function<void()> task;
                    {
                        std::unique_lock<std::mutex> lock(_mtx);
                        _cv.wait(lock, [this]() {
                            return _stop.load() || !_tasks.empty();
                        });
                        if (_stop.load() && _tasks.empty()) return;
                        task = std::move(_tasks.front());
                        _tasks.pop();
                    }
                    task();
                }
            });
        }
        logMessage(NORMAL, "[DispatchPool] Started with %d threads.", nthreads);
    }

    ~DispatchPool() {
        _stop.store(true);
        _cv.notify_all();
        for (auto& t : _threads) {
            if (t.joinable()) t.join();
        }
    }

    void push(std::function<void()> task) {
        {
            std::lock_guard<std::mutex> lock(_mtx);
            _tasks.push(std::move(task));
        }
        _cv.notify_one();
    }

    size_t pending_count() {
        std::lock_guard<std::mutex> lock(_mtx);
        return _tasks.size();
    }
};

#endif // __DISPATCH_POOL_HPP__