thread-pool.hpp

#pragma once

#include <vector>
#include <queue>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <future>
#include <functional>
#include <stdexcept>
#include <type_traits>
#include <atomic>

class ThreadPool {
public:
    explicit ThreadPool(std::size_t thread_count = std::thread::hardware_concurrency())
        : stop(false), paused(false), active_tasks_(0)
    {
        if (thread_count == 0) thread_count = 1;
        spawn_workers(thread_count);
    }

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

    ~ThreadPool() { shutdown(); }

    /**
     * Adds a new task to the thread pool.
     * The task is defined by a callable `f` and its arguments `args`.
     */
    template<typename F, typename... Args>
    auto submit_task(F&& f, Args&&... args) -> std::future<std::invoke_result_t<F, Args...>>
    {
        using ReturnType = std::invoke_result_t<F, Args...>;

        {
            std::unique_lock<std::mutex> lock(m);
            if (stop) throw std::runtime_error("ThreadPool: submit on stopped pool");
        }

        auto task = std::make_shared<std::packaged_task<ReturnType()>>(
            [f = std::forward<F>(f),
             args_tuple = std::make_tuple(std::forward<Args>(args)...)]() mutable {
                return std::apply(std::move(f), std::move(args_tuple));
            }
        );

        std::future<ReturnType> future = task->get_future();

        {
            std::unique_lock<std::mutex> lock(m);
            tasks.emplace([task]() { (*task)(); });
        }

        cv_task.notify_one(); // Notify one worker that a new task is available
        return future;
    }

    /**
     * Waits for all currently queued and active tasks to complete.
     */
    void wait_all()
    {
        std::unique_lock<std::mutex> lock(m);
        cv_done.wait(lock, [this] {
            return tasks.empty() && active_tasks_ == 0;
        });
    }

    /**
     * Checks if there are no pending tasks and no active tasks in the thread pool.
     */
    [[nodiscard]] bool all_done() const
    {
        std::unique_lock<std::mutex> lock(m);
        return tasks.empty() && active_tasks_ == 0;
    }

    /**
     * Pauses the thread pool, preventing workers from executing tasks
     * until `resume()` is called.
     */
    void pause()
    {
        std::unique_lock<std::mutex> lock(m);
        paused = true;
    }

    /**
     * Resumes the thread pool, allowing workers to execute tasks again after being paused.
     */
    void resume()
    {
        {
            std::unique_lock<std::mutex> lock(m);
            paused = false;
        }
        cv_task.notify_all();
    }

    [[nodiscard]] bool is_paused() const
    {
        std::unique_lock<std::mutex> lock(m);
        return paused;
    }

    void resize(std::size_t new_count)
    {
        if (new_count == 0) new_count = 1;

        std::size_t current = workers.size();
        if (new_count == current) return;

        if (new_count > current) {
            spawn_workers(new_count - current);
        } else {
            {
                std::unique_lock<std::mutex> lock(m);
                workers_to_remove += (current - new_count);
            }
            cv_task.notify_all();

            workers.erase(
                std::remove_if(workers.begin(), workers.end(),
                    [](std::thread& t) {
                        if (t.joinable()) { t.join(); return true; }
                        return false;
                    }),
                workers.end()
            );
        }
    }

    void shutdown()
    {
        {
            std::unique_lock<std::mutex> lock(m);
            if (stop) return;
            stop = true;
        }
        cv_task.notify_all();
        for (auto& w : workers) if (w.joinable()) w.join();
        workers.clear();
    }

    void shutdown_now()
    {
        {
            std::unique_lock<std::mutex> lock(m);
            if (stop) return;
            stop = true;
            std::queue<std::function<void()>> empty;
            std::swap(tasks, empty);
        }
        cv_task.notify_all();
        for (auto& w : workers) if (w.joinable()) w.join();
        workers.clear();
    }

    [[nodiscard]] std::size_t pending_tasks() const
    {
        std::unique_lock<std::mutex> lock(m);
        return tasks.size();
    }

    [[nodiscard]] std::size_t active_tasks() const
    {
        return active_tasks_.load(std::memory_order_relaxed);
    }

    [[nodiscard]] std::size_t thread_count() const
    {
        std::unique_lock<std::mutex> lock(m);
        return workers.size();
    }

private:
    void spawn_workers(std::size_t n)
    {
        for (std::size_t i = 0; i < n; ++i)
            workers.emplace_back([this] { worker_loop(); });
    }

    void worker_loop()
    {
        for (;;) {
            std::function<void()> task;

            {
                std::unique_lock<std::mutex> lock(m);

                cv_task.wait(lock, [this] {
                    return stop || workers_to_remove > 0 || (!paused && !tasks.empty());
                });

                if (stop && tasks.empty()) return;
                if (workers_to_remove > 0) { --workers_to_remove; return; }
                if (paused) continue;

                task = std::move(tasks.front());
                tasks.pop();
            }

            ++active_tasks_;
            try { task(); } catch (...) {}
            --active_tasks_;

            {
                std::unique_lock<std::mutex> lock(m);
                if (tasks.empty() && active_tasks_ == 0)
                    cv_done.notify_all();
            }
        }
    }

    mutable std::mutex         m;
    std::condition_variable    cv_task;
    std::condition_variable    cv_done;

    std::vector<std::thread>          workers;
    std::queue<std::function<void()>> tasks;

    bool        stop;
    bool        paused;
    std::size_t workers_to_remove = 0;

    std::atomic<std::size_t> active_tasks_;
};