threadpool.cc

/**
 * @file threadpool.cc
 * @brief 高性能线程池实现
 * @author meow
 * @date 2025
 * 
 * 这是一个功能完整的C++线程池实现，支持：
 * - 核心线程和缓存线程管理
 * - 任务队列限制
 * - 异常处理
 * - 优雅关闭
 * - 线程安全操作
 */

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <functional>
#include <future>
#include <iostream>
#include <list>
#include <memory>
#include <mutex>
#include <queue>
#include <thread>
#include <exception>
#include <vector>
#include <random>
#include <cassert>

namespace meow {

/**
 * @class ThreadPool
 * @brief 高性能线程池实现
 * 
 * 特性：
 * - 支持核心线程（常驻）和缓存线程（可回收）
 * - 任务队列大小限制
 * - 线程安全的任务提交和执行
 * - 异常安全的任务执行
 * - 支持返回值的异步任务
 */
class ThreadPool {
   public:
    using PoolSeconds = std::chrono::seconds;

    /**
     * @struct ThreadPoolConfig
     * @brief 线程池配置参数
     */
    struct ThreadPoolConfig {
        int core_threads;    ///< 核心线程数量，始终保持活跃
        int max_threads;     ///< 最大线程数量
        int max_task_size;   ///< 任务队列最大容量
        PoolSeconds time_out; ///< 缓存线程超时时间
        
        /**
         * @brief 构造函数，提供默认配置
         * @param core 核心线程数（默认4）
         * @param max 最大线程数（默认8）
         * @param max_tasks 最大任务数（默认1000）
         * @param timeout 超时时间（默认60秒）
         */
        ThreadPoolConfig(int core = 4, int max = 8, int max_tasks = 1000, PoolSeconds timeout = PoolSeconds(60))
            : core_threads(core), max_threads(max), max_task_size(max_tasks), time_out(timeout) {}
    };

    /**
     * @enum ThreadState
     * @brief 线程状态枚举
     */
    enum class ThreadState { 
        kInit = 0,     ///< 初始状态
        kWaiting = 1,  ///< 等待任务
        kRunning = 2,  ///< 执行任务
        kStop = 3      ///< 停止状态
    };

    /**
     * @enum ThreadFlag
     * @brief 线程类型标识
     */
    enum class ThreadFlag { 
        kInit = 0,   ///< 初始状态
        kCore = 1,   ///< 核心线程
        kCache = 2   ///< 缓存线程
    };

    using ThreadPtr = std::shared_ptr<std::thread>;
    using ThreadId = std::atomic<int>;
    using ThreadStateAtomic = std::atomic<ThreadState>;
    using ThreadFlagAtomic = std::atomic<ThreadFlag>;

    /**
     * @struct ThreadWrapper
     * @brief 线程包装器，封装线程对象及其元数据
     */
    struct ThreadWrapper {
        ThreadPtr ptr;              ///< 线程对象指针
        ThreadId id;                ///< 线程ID
        ThreadFlagAtomic flag;      ///< 线程类型标识
        ThreadStateAtomic state;    ///< 线程状态

        ThreadWrapper() {
            ptr = nullptr;
            id = 0;
            state.store(ThreadState::kInit);
        }
    };
    using ThreadWrapperPtr = std::shared_ptr<ThreadWrapper>;
    using ThreadPoolLock = std::unique_lock<std::mutex>;

    /**
     * @brief 构造函数
     * @param config 线程池配置
     */
    ThreadPool(ThreadPoolConfig config) : config_(config) {
        // 初始化所有原子变量
        this->total_function_num_.store(0);
        this->waiting_thread_num_.store(0);
        this->running_function_num_.store(0);
        this->thread_id_.store(0);
        this->is_shutdown_.store(false);
        this->is_shutdown_now_.store(false);

        // 验证配置有效性
        if (IsValidConfig(config_)) {
            is_available_.store(true);
            std::cout << "ThreadPool initialized successfully with " 
                      << config_.core_threads << " core threads, max " 
                      << config_.max_threads << " threads" << std::endl;
        } else {
            is_available_.store(false);
            std::cerr << "ThreadPool initialization failed: invalid configuration" << std::endl;
        }
    }

    /**
     * @brief 析构函数，确保线程池优雅关闭
     */
    ~ThreadPool() { 
        std::cout << "ThreadPool destructor called" << std::endl;
        ShutDown(); 
        WaitForAllTasks();
        std::cout << "ThreadPool destroyed" << std::endl;
    }

    // 禁用拷贝构造和赋值操作
    ThreadPool(const ThreadPool&) = delete;
    ThreadPool& operator=(const ThreadPool&) = delete;

    /**
     * @brief 启动线程池
     * @return 成功返回true，失败返回false
     */
    bool Start() {
        if (!IsAvailable()) {
            std::cerr << "ThreadPool is not available for start" << std::endl;
            return false;
        }
        
        int core_thread_num = config_.core_threads;
        std::cout << "Starting ThreadPool with " << core_thread_num << " core threads" << std::endl;
        
        // 创建核心线程
        while (core_thread_num-- > 0) {
            AddThread(GetNextThreadId());
        }
        
        std::cout << "ThreadPool started successfully" << std::endl;
        return true;
    }

    /**
     * @brief 获取等待状态的线程数量
     * @return 等待线程数
     */
    int GetWaitingThreadSize() { return this->waiting_thread_num_.load(); }
    
    /**
     * @brief 获取正在执行任务的线程数量
     * @return 运行线程数
     */
    int GetRunningThreadSize() { return this->running_function_num_.load(); }

    /**
     * @brief 获取线程池中总线程数量（线程安全）
     * @return 总线程数
     */
    int GetTotalThreadSize() { 
        std::lock_guard<std::mutex> lock(threads_mutex_);
        return this->worker_threads_.size(); 
    }

    /**
     * @brief 获取待处理任务数量
     * @return 待处理任务数
     */
    int GetPendingTaskSize() {
        std::lock_guard<std::mutex> lock(task_mutex_);
        return tasks_.size();
    }

    /**
     * @brief 提交任务到线程池
     * @tparam F 函数类型
     * @tparam Args 参数类型
     * @param f 要执行的函数
     * @param args 函数参数
     * @return 指向future的智能指针，用于获取执行结果
     */
    template <typename F, typename... Args>
    auto Run(F &&f, Args &&... args) -> std::shared_ptr<std::future<std::invoke_result_t<F, Args...>>> {
        // 检查线程池状态
        if (this->is_shutdown_.load() || this->is_shutdown_now_.load() || !IsAvailable()) {
            std::cerr << "ThreadPool is not available for new tasks" << std::endl;
            return nullptr;
        }
        
        // 检查任务队列是否已满
        {
            std::lock_guard<std::mutex> lock(task_mutex_);
            if (tasks_.size() >= static_cast<size_t>(config_.max_task_size)) {
                std::cerr << "Task queue is full, rejecting new task" << std::endl;
                return nullptr;
            }
        }
        
        // 如果没有等待线程且未达到最大线程数，创建新的缓存线程
        if (GetWaitingThreadSize() == 0 && GetTotalThreadSize() < config_.max_threads) {
            std::cout << "Creating cache thread for task processing" << std::endl;
            AddThread(GetNextThreadId(), ThreadFlag::kCache);
        }

        // 创建packaged_task包装任务
        using return_type = std::invoke_result_t<F, Args...>;
        auto task = std::make_shared<std::packaged_task<return_type()>>(
            std::bind(std::forward<F>(f), std::forward<Args>(args)...));
        
        total_function_num_++;

        // 获取future用于返回结果
        std::future<return_type> res = task->get_future();
        
        // 将任务添加到队列中，包含异常处理
        {
            ThreadPoolLock lock(this->task_mutex_);
            this->tasks_.emplace([task]() { 
                try {
                    (*task)(); 
                } catch (const std::exception& e) {
                    std::cerr << "Task execution failed with exception: " << e.what() << std::endl;
                } catch (...) {
                    std::cerr << "Task execution failed with unknown exception" << std::endl;
                }
            });
        }
        
        // 通知等待的线程有新任务
        this->task_cv_.notify_one();
        
        return std::make_shared<std::future<std::invoke_result_t<F, Args...>>>(std::move(res));
    }

    /**
     * @brief 获取已执行的任务总数
     * @return 已执行任务数
     */
    int GetRunnedFuncNum() { return total_function_num_.load(); }

    /**
     * @brief 等待所有任务完成
     */
    void WaitForAllTasks() {
        std::cout << "Waiting for all tasks to complete..." << std::endl;
        std::unique_lock<std::mutex> lock(task_mutex_);
        task_cv_.wait(lock, [this] { 
            return tasks_.empty() && running_function_num_.load() == 0; 
        });
        std::cout << "All tasks completed" << std::endl;
    }

    /**
     * @brief 优雅关闭线程池（等待任务完成）
     */
    void ShutDown() {
        ShutDown(false);
        std::cout << "ThreadPool shutdown initiated" << std::endl;
    }

    /**
     * @brief 立即关闭线程池（不等待任务完成）
     */
    void ShutDownNow() {
        ShutDown(true);
        std::cout << "ThreadPool immediate shutdown initiated" << std::endl;
    }

    /**
     * @brief 检查线程池是否可用
     * @return 可用返回true
     */
    bool IsAvailable() { return is_available_.load(); }
    
    /**
     * @brief 检查线程池是否已关闭
     * @return 已关闭返回true
     */
    bool IsShutdown() { return is_shutdown_.load() || is_shutdown_now_.load(); }

   private:
    /**
     * @brief 内部关闭方法
     * @param is_now 是否立即关闭
     */
    void ShutDown(bool is_now) {
        if (is_available_.load()) {
            if (is_now) {
                this->is_shutdown_now_.store(true);
            } else {
                this->is_shutdown_.store(true);
            }
            // 通知所有等待的线程
            this->task_cv_.notify_all();
            is_available_.store(false);
        }
    }

    /**
     * @brief 添加核心线程
     * @param id 线程ID
     */
    void AddThread(int id) { AddThread(id, ThreadFlag::kCore); }

    /**
     * @brief 添加线程到线程池
     * @param id 线程ID
     * @param thread_flag 线程类型
     */
    void AddThread(int id, ThreadFlag thread_flag) {
        std::cout << "Adding thread " << id << " with flag " << static_cast<int>(thread_flag) << std::endl;
        
        ThreadWrapperPtr thread_ptr = std::make_shared<ThreadWrapper>();
        thread_ptr->id.store(id);
        thread_ptr->flag.store(thread_flag);
        
        // 定义线程执行函数
        auto func = [this, thread_ptr]() {
            std::cout << "Thread " << thread_ptr->id.load() << " started" << std::endl;
            
            for (;;) {
                std::function<void()> task;
                {
                    // 获取任务锁
                    ThreadPoolLock lock(this->task_mutex_);
                    
                    // 检查是否应该停止
                    if (thread_ptr->state.load() == ThreadState::kStop) {
                        break;
                    }
                    
                    // 设置为等待状态
                    thread_ptr->state.store(ThreadState::kWaiting);
                    ++this->waiting_thread_num_;
                    
                    bool is_timeout = false;
                    
                    // 根据线程类型选择等待策略
                    if (thread_ptr->flag.load() == ThreadFlag::kCore) {
                        // 核心线程无限等待
                        this->task_cv_.wait(lock, [this, thread_ptr] {
                            return (this->is_shutdown_ || this->is_shutdown_now_ || !this->tasks_.empty() ||
                                    thread_ptr->state.load() == ThreadState::kStop);
                        });
                    } else {
                        // 缓存线程带超时等待
                        this->task_cv_.wait_for(lock, this->config_.time_out, [this, thread_ptr] {
                            return (this->is_shutdown_ || this->is_shutdown_now_ || !this->tasks_.empty() ||
                                    thread_ptr->state.load() == ThreadState::kStop);
                        });
                        // 检查是否超时
                        is_timeout = !(this->is_shutdown_ || this->is_shutdown_now_ || !this->tasks_.empty() ||
                                       thread_ptr->state.load() == ThreadState::kStop);
                    }
                    
                    --this->waiting_thread_num_;

                    // 超时则标记为停止
                    if (is_timeout) {
                        std::cout << "Cache thread " << thread_ptr->id.load() << " timeout, stopping" << std::endl;
                        thread_ptr->state.store(ThreadState::kStop);
                    }

                    // 各种退出条件检查
                    if (thread_ptr->state.load() == ThreadState::kStop) {
                        break;
                    }
                    if (this->is_shutdown_ && this->tasks_.empty()) {
                        std::cout << "Thread " << thread_ptr->id.load() << " exiting: shutdown and no tasks" << std::endl;
                        break;
                    }
                    if (this->is_shutdown_now_) {
                        std::cout << "Thread " << thread_ptr->id.load() << " exiting: immediate shutdown" << std::endl;
                        break;
                    }
                    
                    // 获取任务并开始执行
                    thread_ptr->state.store(ThreadState::kRunning);
                    task = std::move(this->tasks_.front());
                    this->tasks_.pop();
                    ++this->running_function_num_;
                }
                
                // 执行任务（在锁外执行）
                task();
                --this->running_function_num_;
                
                // 通知可能等待的线程
                task_cv_.notify_one();
            }
            
            std::cout << "Thread " << thread_ptr->id.load() << " ended" << std::endl;
        };
        
        // 创建并启动线程
        thread_ptr->ptr = std::make_shared<std::thread>(std::move(func));
        if (thread_ptr->ptr->joinable()) {
            thread_ptr->ptr->detach();
        }
        
        // 将线程添加到线程列表（线程安全）
        {
            std::lock_guard<std::mutex> lock(threads_mutex_);
            this->worker_threads_.emplace_back(std::move(thread_ptr));
        }
    }

    /**
     * @brief 获取下一个线程ID
     * @return 线程ID
     */
    int GetNextThreadId() { return this->thread_id_++; }

    /**
     * @brief 验证配置参数的有效性
     * @param config 配置参数
     * @return 有效返回true
     */
    bool IsValidConfig(ThreadPoolConfig config) {
        if (config.core_threads < 1 || 
            config.max_threads < config.core_threads || 
            config.time_out.count() < 1 || 
            config.max_task_size < 1) {
            return false;
        }
        return true;
    }

   private:
    ThreadPoolConfig config_;                           ///< 线程池配置
    std::list<ThreadWrapperPtr> worker_threads_;        ///< 工作线程列表
    std::mutex threads_mutex_;                          ///< 保护线程列表的互斥锁
    std::queue<std::function<void()>> tasks_;           ///< 任务队列
    std::mutex task_mutex_;                             ///< 任务队列互斥锁
    std::condition_variable task_cv_;                   ///< 任务条件变量
    std::atomic<int> total_function_num_;               ///< 总执行任务数
    std::atomic<int> waiting_thread_num_;               ///< 等待线程数
    std::atomic<int> running_function_num_;             ///< 正在执行任务的线程数
    std::atomic<int> thread_id_;                        ///< 线程ID计数器
    std::atomic<bool> is_shutdown_now_;                 ///< 立即关闭标志
    std::atomic<bool> is_shutdown_;                     ///< 优雅关闭标志
    std::atomic<bool> is_available_;                    ///< 线程池可用标志
};

} // namespace meow

// ==================== 测试代码 ====================

/**
 * @brief 简单的计算任务
 * @param n 计算参数
 * @return 计算结果
 */
int fibonacci(int n) {
    if (n <= 1) return n;
    return fibonacci(n - 1) + fibonacci(n - 2);
}

/**
 * @brief 可能抛出异常的任务
 * @param should_throw 是否抛出异常
 * @return 成功返回42
 */
int potentially_throwing_task(bool should_throw) {
    if (should_throw) {
        throw std::runtime_error("Test exception");
    }
    return 42;
}

/**
 * @brief 模拟IO操作的任务
 * @param duration_ms 睡眠时间（毫秒）
 * @param task_id 任务ID
 * @return 任务ID
 */
int simulate_io_task(int duration_ms, int task_id) {
    std::this_thread::sleep_for(std::chrono::milliseconds(duration_ms));
    std::cout << "IO Task " << task_id << " completed after " << duration_ms << "ms" << std::endl;
    return task_id;
}

/**
 * @brief 测试基本功能
 */
void test_basic_functionality() {
    std::cout << "\n=== Testing Basic Functionality ===" << std::endl;
    
    // 创建线程池配置
    meow::ThreadPool::ThreadPoolConfig config(2, 4, 50, std::chrono::seconds(3));
    meow::ThreadPool pool(config);
    
    // 启动线程池
    assert(pool.Start());
    
    // 提交简单任务
    std::vector<std::shared_ptr<std::future<int>>> results;
    
    for (int i = 0; i < 5; ++i) {
        auto future = pool.Run([i]() {
            std::cout << "Task " << i << " executing in thread " << std::this_thread::get_id() << std::endl;
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            return i * i;
        });
        
        if (future) {
            results.push_back(future);
        }
    }
    
    // 获取结果
    std::cout << "Results: ";
    for (auto& future : results) {
        if (future) {
            std::cout << future->get() << " ";
        }
    }
    std::cout << std::endl;
    
    std::cout << "Total threads: " << pool.GetTotalThreadSize() << std::endl;
    std::cout << "Executed tasks: " << pool.GetRunnedFuncNum() << std::endl;
}

/**
 * @brief 测试异常处理
 */
void test_exception_handling() {
    std::cout << "\n=== Testing Exception Handling ===" << std::endl;
    
    meow::ThreadPool::ThreadPoolConfig config(1, 2, 10);
    meow::ThreadPool pool(config);
    pool.Start();
    
    // 提交会抛出异常的任务
    auto future1 = pool.Run([]() { return potentially_throwing_task(true); });
    auto future2 = pool.Run([]() { return potentially_throwing_task(false); });
    
    try {
        if (future1) {
            int result = future1->get();
            std::cout << "Task 1 result: " << result << std::endl;
        }
    } catch (const std::exception& e) {
        std::cout << "Task 1 exception caught: " << e.what() << std::endl;
    }
    
    if (future2) {
        std::cout << "Task 2 result: " << future2->get() << std::endl;
    }
}

/**
 * @brief 测试大量任务处理
 */
void test_heavy_load() {
    std::cout << "\n=== Testing Heavy Load ===" << std::endl;
    
    meow::ThreadPool::ThreadPoolConfig config(3, 6, 100);
    meow::ThreadPool pool(config);
    pool.Start();
    
    std::vector<std::shared_ptr<std::future<int>>> futures;
    
    // 提交大量任务
    const int task_count = 20;
    for (int i = 0; i < task_count; ++i) {
        auto future = pool.Run([i]() {
            return simulate_io_task(50 + (i % 5) * 10, i);
        });
        
        if (future) {
            futures.push_back(future);
        }
    }
    
    // 等待所有任务完成
    int completed = 0;
    for (auto& future : futures) {
        if (future) {
            future->get();
            completed++;
        }
    }
    
    std::cout << "Completed " << completed << " out of " << task_count << " tasks" << std::endl;
    std::cout << "Thread pool stats:" << std::endl;
    std::cout << "  Total threads: " << pool.GetTotalThreadSize() << std::endl;
    std::cout << "  Waiting threads: " << pool.GetWaitingThreadSize() << std::endl;
    std::cout << "  Running threads: " << pool.GetRunningThreadSize() << std::endl;
    std::cout << "  Pending tasks: " << pool.GetPendingTaskSize() << std::endl;
}

/**
 * @brief 测试任务队列限制
 */
void test_task_queue_limit() {
    std::cout << "\n=== Testing Task Queue Limit ===" << std::endl;
    
    // 创建容量很小的线程池
    meow::ThreadPool::ThreadPoolConfig config(1, 2, 5); // 最多5个任务
    meow::ThreadPool pool(config);
    pool.Start();
    
    // 提交超过容量的任务
    int submitted = 0;
    for (int i = 0; i < 10; ++i) {
        auto future = pool.Run([i]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(200));
            return i;
        });
        
        if (future) {
            submitted++;
        } else {
            std::cout << "Task " << i << " rejected (queue full)" << std::endl;
        }
    }
    
    std::cout << "Successfully submitted " << submitted << " tasks" << std::endl;
}

/**
 * @brief 测试缓存线程超时
 */
void test_cache_thread_timeout() {
    std::cout << "\n=== Testing Cache Thread Timeout ===" << std::endl;
    
    meow::ThreadPool::ThreadPoolConfig config(1, 3, 20, std::chrono::seconds(2));
    meow::ThreadPool pool(config);
    pool.Start();
    
    std::cout << "Initial thread count: " << pool.GetTotalThreadSize() << std::endl;
    
    // 提交任务触发缓存线程创建
    std::vector<std::shared_ptr<std::future<int>>> futures;
    for (int i = 0; i < 5; ++i) {
        auto future = pool.Run([i]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            return i;
        });
        if (future) futures.push_back(future);
    }
    
    // 等待任务完成
    for (auto& future : futures) {
        if (future) future->get();
    }
    
    std::cout << "Thread count after tasks: " << pool.GetTotalThreadSize() << std::endl;
    
    // 等待缓存线程超时
    std::cout << "Waiting for cache threads to timeout..." << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(3));
    
    std::cout << "Final thread count: " << pool.GetTotalThreadSize() << std::endl;
}

/**
 * @brief 主测试函数
 */
int main() {
    std::cout << "ThreadPool Test Suite Started" << std::endl;
    std::cout << "==============================" << std::endl;
    
    try {
        // 运行各项测试
        test_basic_functionality();
        test_exception_handling();
        test_heavy_load();
        test_task_queue_limit();
        test_cache_thread_timeout();
        
        std::cout << "\n==============================" << std::endl;
        std::cout << "All tests completed successfully!" << std::endl;
        
    } catch (const std::exception& e) {
        std::cerr << "Test failed with exception: " << e.what() << std::endl;
        return 1;
    } catch (...) {
        std::cerr << "Test failed with unknown exception" << std::endl;
        return 1;
    }
    
    return 0;
}