TaggedThreadPool.hxx

Go to the documentation of this file.

//
// Created by Zhiping Jiang on 9/12/22.
//
 
#ifndef PICOSCENES_PLATFORM_TAGGEDTHREADPOOL_HXX
#define PICOSCENES_PLATFORM_TAGGEDTHREADPOOL_HXX
 
#include <thread>
#include <iostream>
#include <any>
#include <numeric>
#include <vector>
#include <deque>
#include <functional>
#include <condition_variable>
#include <optional>
#include <map>
#include <shared_mutex>
#include <algorithm>
#include <chrono>
#include "Singleton.hxx"
#ifdef _WIN32
#include "pthread.h"
#endif
 
using namespace std::chrono_literals;
 
template<unsigned ThreadCountCompileTime = 1>
class TaggedThreadPool {
public:
    explicit TaggedThreadPool(const int threadCount = ThreadCountCompileTime, const std::string& threadNamePrefix = std::string{}) : poolState(true), threadNamePrefix(threadNamePrefix) {
        addThread(threadCount);
    }
 
    void addThread(const int numThreads2Add) {
        for (auto i = 0; i < numThreads2Add; i++) {
            threads.emplace_back(std::thread([this] {
                {
                    auto lock = std::unique_lock(poolMutex);
                    threadSequenceNumberMap[std::this_thread::get_id()] = threadSequenceNumberSource;
                    if (threadSpecificTasksMap.contains(threadSequenceNumberSource) || threadSpecificMutexes.contains(threadSequenceNumberSource))
                        throw std::runtime_error("Duplicate sequence number for threads");
                    threadSpecificTasksMap[threadSequenceNumberSource] = std::deque<std::pair<std::function<void(const std::any&)>, std::any>>();
                    threadSpecificMutexes[threadSequenceNumberSource] = std::make_unique<std::mutex>();
                    perThreadJobCount[threadSequenceNumberSource] = 0;
                    perThreadTimeOccupation[threadSequenceNumberSource] = std::make_pair(0.0, 0.0);
                    ++threadSequenceNumberSource; // 放在最后，是确保thread number是从0开始
                    setThreadNamePrefix();
                }
                serviceLoopPerThread();
            }));
        }
    }
 
    template<typename Task, typename Params, typename = std::enable_if_t<std::is_invocable_v<Task, const std::any &> && std::is_same_v<std::decay_t<Params>, std::any>>>
    void addJob(Task&& task, Params&& params, const std::optional<uint32_t>& threadSN = std::nullopt) {
        if (threadSN && threadSpecificTasksMap.contains(*threadSN)) {
            auto lock = std::lock_guard(*threadSpecificMutexes[*threadSN]);
            threadSpecificTasksMap[*threadSN].emplace_back(std::forward<Task>(task), std::forward<Params>(params));
        } else {
            auto lock = std::lock_guard(sharedPoolMutex);
            sharedTaskQueue.emplace_back(std::forward<Task>(task), std::forward<Params>(params));
        }
        taskWaitCV.notify_all(); // Must be notify_all to prevent thread hanging
    }
 
    void AddJob(const std::function<void()>& task) {
        addJob([task](const std::any&) {
            task();
        }, std::any{}, std::nullopt);
    }
 
    [[nodiscard]] size_t threadCount() const {
        return threads.size();
    }
 
    [[nodiscard]] size_t taskCount() const {
        auto taskCount = sharedTaskQueue.size();
        taskCount += std::accumulate(threadSpecificTasksMap.cbegin(), threadSpecificTasksMap.cend(), 0,
                                     [](size_t sum, const auto& pair) {
                                         return sum + pair.second.size();
                                     });
        return taskCount;
    }
 
    ~TaggedThreadPool() {
        shutdown();
    }
 
private:
    bool poolState{};
    const std::string threadNamePrefix{"PS"};
    std::atomic<uint32_t> threadSequenceNumberSource{0};
 
    // Threads and Thread-Ids
    std::vector<std::thread> threads;
    std::unordered_map<std::thread::id, uint32_t> threadSequenceNumberMap;
 
    // Shared and Thread-Specific queues
    std::deque<std::pair<std::function<void(const std::any&)>, std::any>> sharedTaskQueue;
    std::unordered_map<uint32_t, std::deque<std::pair<std::function<void(const std::any&)>, std::any>>> threadSpecificTasksMap;
 
    // queue mutexes
    std::mutex sharedPoolMutex;
    std::unordered_map<uint32_t, std::unique_ptr<std::mutex>> threadSpecificMutexes;
 
    // Statistics
    std::unordered_map<uint32_t, int64_t> perThreadJobCount;
    std::unordered_map<uint32_t, std::pair<double, double>> perThreadTimeOccupation;
 
    std::shared_mutex poolMutex;
    std::mutex taskWaitMutex;
    std::condition_variable taskWaitCV;
 
    void serviceLoopPerThread() {
        const auto threadOnlineTime = std::chrono::system_clock::now();
        while (poolState) {
            if (const auto jobPack = getNextJob()) {
                auto start = std::chrono::system_clock::now();
                jobPack->first(jobPack->second);
                auto end = std::chrono::system_clock::now();
                const auto threadId = std::this_thread::get_id();
                perThreadJobCount[threadSequenceNumberMap[threadId]]++;
                perThreadTimeOccupation[threadSequenceNumberMap[threadId]].first += (end - start).count();
                perThreadTimeOccupation[threadSequenceNumberMap[threadId]].second = (end - threadOnlineTime).count();
            }
        }
    }
 
    std::optional<std::pair<std::function<void(const std::any&)>, std::any>> getNextJob() {
        const auto threadId = std::this_thread::get_id();
 
        auto shouldWait = [this, &threadId]() {
            if (!sharedTaskQueue.empty())
                return true;
 
            auto lock = std::shared_lock(poolMutex);
            auto threadSpecificResult = threadSequenceNumberMap.contains(threadId) && !threadSpecificTasksMap[threadSequenceNumberMap[threadId]].empty();
            return threadSpecificResult;
        };
 
        auto taskWaitLock = std::unique_lock(taskWaitMutex);
        taskWaitCV.wait(taskWaitLock, shouldWait);
 
        if (!sharedTaskQueue.empty()) {
            auto lock = std::lock_guard(sharedPoolMutex);
            auto jobPack = sharedTaskQueue.front();
            sharedTaskQueue.pop_front();
            return std::move(jobPack);
        }
 
        if (const auto it = threadSequenceNumberMap.find(threadId); it != threadSequenceNumberMap.end() && !threadSpecificTasksMap[it->second].empty()) {
            auto taskQueueOpLock = std::lock_guard(*threadSpecificMutexes[it->second]);
            auto jobPack = threadSpecificTasksMap[it->second].front();
            threadSpecificTasksMap[it->second].pop_front();
            return jobPack;
        }
 
        return std::nullopt;
    }
 
    void setThreadNamePrefix() {
        if (threadNamePrefix.empty())
            return;
 
#if !defined(__APPLE__)
        if (threadNamePrefix.length() > 15)
            throw std::invalid_argument("Thread name prefix MUST be shorter than 15 characters.");
#endif
 
        const size_t threadOrder = threadSequenceNumberMap[std::this_thread::get_id()];
        std::string threadName = threadNamePrefix + std::to_string(threadOrder);
 
#if  !defined(__APPLE__)
        if (pthread_setname_np(pthread_self(), threadName.c_str()))
#else
        if (pthread_setname_np(threadName.c_str()))
#endif
            throw std::runtime_error("failed to set name for current thread...");
    }
 
    void shutdown() {
        auto currentThreadId = std::this_thread::get_id();
        if (std::any_of(threads.cbegin(), threads.cend(), [currentThreadId](const std::thread& t) { return t.get_id() == currentThreadId; })) {
            throw std::runtime_error("Cannot shutdown thread pool from within ...");
        }
 
        auto lock = std::lock_guard(poolMutex);
        poolState = false;
        taskWaitCV.notify_all();
        for (auto& thread: threads) {
            if (thread.joinable()) {
                thread.join();
            }
        }
 
        threadSequenceNumberMap.clear();
        threadSpecificTasksMap.clear();
        sharedTaskQueue.clear();
        threads.clear();
    }
};
 
typedef Singleton<TaggedThreadPool<300>> ThreadPoolSingleton;
 
#endif //PICOSCENES_PLATFORM_TAGGEDTHREADPOOL_HXX