ThreadPool.h

#ifndef THREAD_POOL_H
#define THREAD_POOL_H


/* heavily modified */


#include <vector>
#include <queue>
#include <memory>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <future>
#include <functional>
#include <stdexcept>
#include "error.h"


class ThreadPool {
public:
    ThreadPool(size_t, std::unordered_map<std::thread::id, std::vector<std::string>>*);
    template<class F, class... Args>
    auto enqueue(F&& f, Args&&... args) 
        -> std::future<typename std::result_of<F(Args...)>::type>;
    ~ThreadPool();

    void set_title(const std::string& s)
    {
        std::lock_guard<std::mutex> lock(title_mutex);
        title = s;
    }

    const std::string& get_title()
    {
        std::lock_guard<std::mutex> lock(title_mutex);
        return title;
    }

private:
    // need to keep track of threads so we can join them
    std::vector< std::thread > workers;
    // the task queue
    std::queue<std::function<void()>> tasks;
    
    // synchronization
    std::mutex queue_mutex;
    std::condition_variable condition;
	std::atomic_bool stop;

	std::atomic_ulong total_queued;
	std::atomic_ulong queued;
	std::atomic_ulong done;

    std::string title;
    std::mutex title_mutex;
    std::unordered_map<std::thread::id, std::vector<std::string>>* thread_stacks;
    std::thread::id master_id;

	void display_progress();
};

inline void ThreadPool::display_progress()
{
	if (!msg_handler_t::color()) {
		return;
    } else {
        static constexpr const char* const SAVE_CURSOR = "\033[s";
        static constexpr const char* const RESTORE_CURSOR = "\033[u";
        static constexpr const char* const ERASE_TO_EOL = "\033[K";
        static constexpr const char* const GOTO_0_0 = "\033[0;0H";
        /*static constexpr const char* const HIDE_CURSOR = "\033?25l";*/
        /*static constexpr const char* const SHOW_CURSOR = "\033?25h";*/

        /*std::unique_lock<msg_handler_t::lock_type> lock(msg_handler_t::mutex);*/
        if (msg_handler_t::instance().queue.m_stop) {
            return;
        }
        std::stringstream msg;
        msg << SAVE_CURSOR;
        msg /*<< HIDE_CURSOR*/ << GOTO_0_0;
        msg << msg_handler_t::i();
        msg << "[SPELL-QTL] Current: " << msg_handler_t::w();
        {
            std::lock_guard<std::mutex> lock(title_mutex);
            msg <<  title;
        }
        msg << msg_handler_t::i() << ERASE_TO_EOL << std::endl;
        msg << "[SPELL-QTL] " << msg_handler_t::n() << workers.size() << " threads, " << total_queued << " tasks queued in previous batches" << msg_handler_t::i() << ERASE_TO_EOL << std::endl;
        if (queued) {
            msg << "[SPELL-QTL] Task info: ";
            msg << msg_handler_t::n();
            msg << queued << " queued, " << done << " done, ";
            msg << msg_handler_t::w();
            msg << (100 * done / queued) << "% progress";
        }
        msg << ERASE_TO_EOL;
        msg << msg_handler_t::n() << std::endl;
        int ti = 0;
        std::stringstream master;
        std::stringstream workers;
        std::stringstream* tmp;
        for (const auto& kv: *thread_stacks) {
            std::thread::id id = kv.first;
            if (id == master_id) {
                master << "[MASTER] ";
                tmp = &master;
            } else {
                workers << '[' << (ti++) << "] ";
                tmp = &workers;
            }
            if (kv.second.size()) {
                (*tmp) << kv.second[0];
                for (size_t i = 1; i < kv.second.size(); ++i) {
                    (*tmp) << " | " << kv.second[i];
                }
            } else {
                (*tmp) << "<idle>";
            }
            (*tmp) << ERASE_TO_EOL << std::endl;
        }
        msg << master.str() << workers.str();
        msg << std::endl << "----------------------------------------------------------";
        msg << ERASE_TO_EOL << std::endl;
        msg << RESTORE_CURSOR /*<< SHOW_CURSOR*/;
        /*CREATE_MESSAGE(msg_channel::Out, msg.str());*/
        if (!msg_handler_t::instance().queue.m_stop) {
            msg_handler_t::instance().queue.cout << msg.str() << std::flush;
        }
        if (done == queued) {
            total_queued += queued;
            queued = 0;
            done = 0;
        }
    }
}
 
// the constructor just launches some amount of workers
inline ThreadPool::ThreadPool(size_t threads, std::unordered_map<std::thread::id, std::vector<std::string>>* tt)
    : stop(false), total_queued(0), queued(0), done(0), title(), thread_stacks(tt)
    , master_id(std::this_thread::get_id())
{
    for(size_t i = 0;i<threads;++i) {
        workers.emplace_back(
            [this]
            {
                while(true)
                {
                    {
                        std::unique_lock<std::mutex> lock(this->queue_mutex);
                        while(!this->stop && this->tasks.empty()) {
                            this->condition.wait(lock);
                        }
                        if(this->stop && this->tasks.empty()) {
                            return;
                        }
                        std::function<void()> task(this->tasks.front());
                        this->tasks.pop();
                        lock.unlock();
                        task();
                        ++done;
                    }
                    if (msg_handler_t::color()) {
                        msg_handler_t::instance().queue.lock_stream();
    					display_progress();
                        msg_handler_t::instance().queue.unlock_stream();
                    }
                }
            }
        );
        /*MSG_DEBUG("NEW THREAD [" << workers.back().get_id() << ']');*/
        thread_stacks->insert({workers.back().get_id(), {}});
    }
    thread_stacks->insert({master_id, {}});
    msg_handler_t::hook([this] () { display_progress(); });
}

// add new work item to the pool
template<class F, class... Args>
auto ThreadPool::enqueue(F&& f, Args&&... args) 
    -> std::future<typename std::result_of<F(Args...)>::type>
{
    typedef typename std::result_of<F(Args...)>::type return_type;
    
    // don't allow enqueueing after stopping the pool
    if(stop)
        throw std::runtime_error("enqueue on stopped ThreadPool");

    auto task = std::make_shared< std::packaged_task<return_type()> >(
            std::bind(std::forward<F>(f), std::forward<Args>(args)...)
        );
        
    std::future<return_type> res = task->get_future();
    {
        std::unique_lock<std::mutex> lock(queue_mutex);
        tasks.push([task](){ (*task)(); });
    }
	++queued;
    condition.notify_one();
    return res;
}

// the destructor joins all threads
inline ThreadPool::~ThreadPool()
{
    {
        std::unique_lock<std::mutex> lock(queue_mutex);
        stop = true;
    }
    condition.notify_all();
    for(size_t i = 0;i<workers.size();++i)
        workers[i].join();
}

#endif