GrabBag/Module/ThreadPool/Src/VrThreadPool.cpp

#include "VrThreadPool.h"
#include <thread>
#include <chrono>

// ============================================================================
// IVrThreadPool 静态工厂方法
// ============================================================================

bool IVrThreadPool::CreateThreadPool(IVrThreadPool** ppThreadPool, int maxThreadCount)
{
    if (nullptr == ppThreadPool) return false;
    *ppThreadPool = new VrThreadPool(maxThreadCount);
    return true;
}

// ============================================================================
// VrThreadPool 构造 / 析构
// ============================================================================

VrThreadPool::VrThreadPool(int maxThreadCount)
    : m_nMaxThreadCount(maxThreadCount)
    , m_bStop(false)
{
    // 0: 使用硬件并发数；负数（如 -1）: 不限制线程数
    if (m_nMaxThreadCount == 0)
    {
        unsigned int hwThreads = std::thread::hardware_concurrency();
        m_nMaxThreadCount = static_cast<int>(hwThreads > 0 ? hwThreads : 4);
    }
}

VrThreadPool::~VrThreadPool()
{
    Stop();

    // 等待所有线程结束并释放资源
    for (VrThread* pThread : m_vetThread)
    {
        if (pThread)
        {
            delete pThread;
        }
    }
    m_vetThread.clear();
    m_deqRecoveryThread.clear();
    m_deqPendingTasks.clear();
}

// ============================================================================
// 公开接口
// ============================================================================

int VrThreadPool::ExecTask(ThreadExecFunc fExecFunc, void* pParam)
{
    if (nullptr == fExecFunc) return VrThreadPoolError::InvalidParam;

    VrThread* pThread   = nullptr;
    bool      needInit  = false;

    {
        std::lock_guard<std::mutex> lck(m_mutexPool);

        // 1. 优先从回收队列取空闲线程
        if (!m_deqRecoveryThread.empty())
        {
            pThread = m_deqRecoveryThread.front();
            m_deqRecoveryThread.pop_front();
        }
        // 2. 未达上限，创建新线程
        else if (m_nMaxThreadCount < 0
              || static_cast<int>(m_vetThread.size()) < m_nMaxThreadCount)
        {
            pThread = new VrThread(static_cast<int>(m_vetThread.size()));
            m_vetThread.push_back(pThread);
            needInit = true;
        }
        // 3. 已达上限，放入等待队列
        else
        {
            m_deqPendingTasks.emplace_back(fExecFunc, pParam);
            return VrThreadPoolError::Ok;
        }
    }

    // 新线程在锁外初始化（Init 会阻塞直到线程启动）
    if (needInit)
    {
        pThread->Init();
    }

    return pThread->ExecTask(
        fExecFunc,
        pParam,
        std::bind(&VrThreadPool::_ExecFinish, this, std::placeholders::_1),
        pThread
    );
}

bool VrThreadPool::WaitAllTaskFinish(int timeoutMs)
{
    std::unique_lock<std::mutex> lock(m_mutexPool);

    auto allIdle = [this] {
        // 无线程 或 (所有线程空闲 且 无等待任务)
        return m_vetThread.empty()
            || (m_deqRecoveryThread.size() == m_vetThread.size()
                && m_deqPendingTasks.empty());
    };

    if (timeoutMs <= 0)
    {
        m_cvIdle.wait(lock, allIdle);
        return true;
    }
    else
    {
        return m_cvIdle.wait_for(lock, std::chrono::milliseconds(timeoutMs), allIdle);
    }
}

void VrThreadPool::Stop()
{
    m_bStop = true;

    // 通知所有线程停止
    for (VrThread* pThread : m_vetThread)
    {
        if (pThread)
        {
            pThread->Stop();
        }
    }
}

int VrThreadPool::GetThreadCount() const
{
    std::lock_guard<std::mutex> lck(m_mutexPool);
    return static_cast<int>(m_vetThread.size());
}

int VrThreadPool::GetIdleThreadCount() const
{
    std::lock_guard<std::mutex> lck(m_mutexPool);
    return static_cast<int>(m_deqRecoveryThread.size());
}

int VrThreadPool::GetPendingTaskCount() const
{
    std::lock_guard<std::mutex> lck(m_mutexPool);
    return static_cast<int>(m_deqPendingTasks.size());
}

// ============================================================================
// 内部方法
// ============================================================================

void VrThreadPool::_ExecFinish(void* pObject)
{
    if (m_bStop) return;

    VrThread* pThread = reinterpret_cast<VrThread*>(pObject);
    if (!pThread) return;

    ThreadExecFunc  nextFunc  = nullptr;
    void*           nextParam = nullptr;

    {
        std::lock_guard<std::mutex> lck(m_mutexPool);

        // 有排队任务 → 当前线程直接复用，不进入回收队列
        if (!m_deqPendingTasks.empty())
        {
            auto& task = m_deqPendingTasks.front();
            nextFunc  = task.first;
            nextParam = task.second;
            m_deqPendingTasks.pop_front();
        }
        else
        {
            // 无排队任务 → 放回回收队列
            m_deqRecoveryThread.push_front(pThread);
        }
    }

    if (nextFunc)
    {
        // 在当前线程上直接执行下一个排队任务（跳过回收→取出环节）
        pThread->ExecTask(
            nextFunc,
            nextParam,
            std::bind(&VrThreadPool::_ExecFinish, this, std::placeholders::_1),
            pThread
        );
    }

    m_cvIdle.notify_one();  // 通知 WaitAllTaskFinish
}