Thread_Pool.cpp

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

// ============================================================================
// ThreadPool — 生产级线程池
// ============================================================================
class ThreadPool {
public:
    struct Config {
        size_t min_threads   = 2;
        size_t max_threads   = 8;
        size_t idle_timeout_ms = 30'000;  // 空闲线程回收
    };

    // ---------- 构造 / 析构 ----------
    explicit ThreadPool(size_t num_threads)
        : config_{.min_threads = num_threads,
                 .max_threads = num_threads} {
        Start(num_threads);
    }

    explicit ThreadPool(Config cfg) : config_(std::move(cfg)) {
        Start(cfg.min_threads);
    }

    ~ThreadPool() {
        Shutdown();
    }

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

    // ---------- 任务提交 ----------
    // 返回 std::future<Ret>，完美转发
    template <typename F, typename... Args>
    auto Enqueue(F&& f, Args&&... args)
        -> std::future<std::invoke_result_t<F, Args...>>;

    // ---------- 批量等待 ----------
    void WaitAll();

    // ---------- 控制 ----------
    void Pause();
    void Resume();
    void Resize(size_t n);

    // ---------- 统计 ----------
    size_t Active()   const noexcept { return active_.load(); }
    size_t Queued()   const noexcept { return queued_.load(); }
    size_t Threads()  const noexcept { return workers_.size(); }
    size_t Completed()const noexcept { return completed_.load(); }

private:
    struct TaskBase {
        virtual ~TaskBase() = default;
        virtual void Run() = 0;
    };

    template <typename F>
    struct Task final : TaskBase {
        F func;
        explicit Task(F&& f) : func(std::forward<F>(f)) {}
        void Run() override { func(); }
    };

    void Start(size_t n);
    void WorkerLoop();
    void Shutdown();

    Config config_;
    std::vector<std::thread> workers_;

    std::mutex              mtx_;
    std::condition_variable cv_;
    std::condition_variable done_cv_;  // WaitAll

    std::queue<std::unique_ptr<TaskBase>> tasks_;

    std::atomic<bool> stop_{false};
    std::atomic<bool> pause_{false};
    std::atomic<size_t> active_{0};
    std::atomic<size_t> queued_{0};
    std::atomic<size_t> completed_{0};
};

// ============================================================================
// 实现
// ============================================================================

void ThreadPool::Start(size_t n) {
    for (size_t i = 0; i < n; ++i) {
        workers_.emplace_back(&ThreadPool::WorkerLoop, this);
    }
}

void ThreadPool::WorkerLoop() {
    while (true) {
        std::unique_ptr<TaskBase> task;
        {
            std::unique_lock lock(mtx_);
            cv_.wait(lock, [this] {
                return stop_.load(std::memory_order_acquire) ||
                       (!pause_.load(std::memory_order_acquire) && !tasks_.empty());
            });
            if (stop_ && tasks_.empty()) return;
            if (pause_ && tasks_.empty()) continue;

            task = std::move(tasks_.front());
            tasks_.pop();
            queued_.fetch_sub(1, std::memory_order_relaxed);
            active_.fetch_add(1, std::memory_order_relaxed);
        }

        try {
            task->Run();
        } catch (const std::exception& e) {
            std::cerr << "[ThreadPool] unhandled exception: " << e.what() << std::endl;
        } catch (...) {
            std::cerr << "[ThreadPool] unhandled unknown exception" << std::endl;
        }

        active_.fetch_sub(1, std::memory_order_relaxed);
        completed_.fetch_add(1, std::memory_order_relaxed);
        done_cv_.notify_all();  // 通知 WaitAll
    }
}

template <typename F, typename... Args>
auto ThreadPool::Enqueue(F&& f, Args&&... args)
    -> std::future<std::invoke_result_t<F, Args...>> {

    using Ret = std::invoke_result_t<F, Args...>;

    auto pkg = std::make_shared<std::packaged_task<Ret()>>(
        std::bind(std::forward<F>(f), std::forward<Args>(args)...));
    auto fut = pkg->get_future();

    auto task = std::make_unique<Task<std::function<void()>>>(
        [pkg = std::move(pkg)]() { (*pkg)(); });

    {
        std::lock_guard lock(mtx_);
        if (stop_) throw std::runtime_error("ThreadPool: stopped");
        tasks_.push(std::move(task));
        queued_.fetch_add(1, std::memory_order_relaxed);
    }
    cv_.notify_one();
    return fut;
}

void ThreadPool::WaitAll() {
    std::unique_lock lock(mtx_);
    done_cv_.wait(lock, [this] {
        return tasks_.empty() && active_.load() == 0;
    });
}

void ThreadPool::Pause() {
    pause_ = true;
}

void ThreadPool::Resume() {
    pause_ = false;
    cv_.notify_all();
}

void ThreadPool::Resize(size_t n) {
    std::lock_guard lock(mtx_);
    if (n <= workers_.size()) return;
    for (size_t i = workers_.size(); i < n; ++i) {
        workers_.emplace_back(&ThreadPool::WorkerLoop, this);
    }
}

void ThreadPool::Shutdown() {
    stop_ = true;
    cv_.notify_all();
    for (auto& t : workers_) {
        if (t.joinable()) t.join();
    }
}

// ============================================================================
// 演示
// ============================================================================
int main() {
    ThreadPool pool(4);

    // ---- 基础用法 ----
    std::cout << "=== ThreadPool Basic ===" << std::endl;
    std::vector<std::future<int>> results;
    for (int i = 0; i < 20; ++i) {
        auto fut = pool.Enqueue([i]() -> int {
            std::cout << "  task " << i << " on " << std::this_thread::get_id() << std::endl;
            return i * i;
        });
        results.push_back(std::move(fut));
    }
    for (auto& f : results) {
        f.get();
    }
    std::cout << "Completed: " << pool.Completed() << std::endl;

    // ---- WaitAll ----
    std::cout << "\n=== WaitAll ===" << std::endl;
    for (int i = 0; i < 10; ++i) {
        pool.Enqueue([i]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        });
    }
    pool.WaitAll();
    std::cout << "All done. Queued=" << pool.Queued()
              << " Active=" << pool.Active() << std::endl;

    // ---- Pause / Resume ----
    std::cout << "\n=== Pause / Resume ===" << std::endl;
    pool.Pause();
    auto f1 = pool.Enqueue([]() { return 1; });
    auto f2 = pool.Enqueue([]() { return 2; });
    std::cout << "Paused — " << pool.Queued() << " tasks queued" << std::endl;
    pool.Resume();
    std::cout << "Resumed — results: " << f1.get() << ", " << f2.get() << std::endl;

    // ---- Resize ----
    std::cout << "\n=== Resize ===" << std::endl;
    pool.Resize(6);
    std::cout << "Threads after resize: " << pool.Threads() << std::endl;

    std::cout << "\nFinal stats: completed=" << pool.Completed()
              << " threads=" << pool.Threads() << std::endl;
    return 0;
}