ThreadPool.hpp

// SPDX-FileCopyrightText: 2025 Vincent Leroy
// SPDX-License-Identifier: MIT
//
// This file is part of dfx.
//
// Licensed under the MIT License. See the LICENSE file in the project root
// for full license information.
 
#pragma once
 
// Standard includes
#include <condition_variable>
#include <list>
#include <mutex>
#include <optional>
#include <shared_mutex>
#include <thread>
#include <vector>
 
// Project includes
#include "tasks/Task.hpp"
#include "ThreadData.hpp"
#include <dfx-utilities/MonotonicIdAllocator.hpp>
 
namespace dfx::Runtime
{
class ThreadPool
{
    struct QueueItem
    {
        TaskPtr task;
        uint64_t seq;
    };
 
    struct QueueItemCompare
    {
        bool operator()(QueueItem const & a, QueueItem const & b) const noexcept
        {
            // higher priority first
            auto const ap = a.task->priority();
            auto const bp = b.task->priority();
            if (ap != bp)
                return ap < bp;
 
            // FIFO among equal priorities: smaller seq first
            return a.seq > b.seq;
        }
    };
 
public:
    ThreadPool(std::size_t threadCount = 1);

    ~ThreadPool();

    void addThread(std::size_t count = 1);

    void removeThread(std::size_t count = 1);

    std::size_t threadCount() const noexcept { std::lock_guard lock(_mutexThreads); return _threads.size(); }
 
public:
    void pushTask(TaskPtr task);

    void pushTasks(std::vector<TaskPtr> tasks);

    template<DerivedFromTask T, typename ... Args>
    requires std::is_constructible_v<T, Args...>
    void emplaceTask(Args && ... args)
    { pushTask(std::make_unique<T>(std::forward<Args>(args)...)); }
 
public:
    using MaybeThreadData = std::optional<std::reference_wrapper<ThreadData>>;
    using MaybeCThreadData = std::optional<std::reference_wrapper<ThreadData const>>;

    MaybeThreadData threadData();

    MaybeCThreadData threadData() const;
 
private:
    void _work(std::stop_token stopToken);
 
private:
    std::vector<QueueItem>                _queue;
    std::mutex                            _mutexQueue;
    Utils::MonotonicIdAllocator<uint64_t> _seqAllocator;
    std::condition_variable_any           _cvQueue;
 
private:
    // Use a list instead of a map so interators and references
    // aren't invalidated when inserted / removing elements
    std::list<ThreadData>     _threadDatas;
    mutable std::shared_mutex _mutexThreadData;
 
private:
    mutable std::mutex        _mutexThreads;
    // The threads have to be delete first in the dtor
    std::vector<std::jthread> _threads;
};
} // !namespace dfx::Runtime