quickpool.hpp

// Copyright 2021 Thomas Nagler (MIT License)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
 
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
 
#include <atomic>
#include <condition_variable>
#include <exception>
#include <functional>
#include <future>
#include <mutex>
#include <thread>
#include <vector>
 
namespace quickpool {
 
class TodoList
{
  public:
    TodoList(size_t num_tasks = 0) noexcept
      : num_tasks_(num_tasks)
    {}
 
    void add(size_t num_tasks = 1) noexcept { num_tasks_.fetch_add(num_tasks); }
 
    void cross(size_t num_tasks = 1)
    {
        num_tasks_.fetch_sub(num_tasks);
        if (num_tasks_ <= 0) {
            {
                std::lock_guard<std::mutex> lk(mtx_); // must lock before signal
            }
            cv_.notify_all();
        }
    }
 
    bool empty() const noexcept { return num_tasks_ <= 0; }
 
    void wait(size_t millis = 0)
    {
        std::this_thread::yield();
        auto wake_up = [this] { return (num_tasks_ <= 0) || exception_ptr_; };
        std::unique_lock<std::mutex> lk(mtx_);
        if (millis == 0) {
            cv_.wait(lk, wake_up);
        } else {
            cv_.wait_for(lk, std::chrono::milliseconds(millis), wake_up);
        }
        if (exception_ptr_)
            std::rethrow_exception(exception_ptr_);
    }
 
    void stop(std::exception_ptr eptr = nullptr) noexcept
    {
        {
            std::lock_guard<std::mutex> lk(mtx_);
            // Some threads may add() or cross() after we stop. The large
            // negative number prevents num_tasks_ from becoming positive again.
            num_tasks_ = std::numeric_limits<int>::min() / 2;
            exception_ptr_ = eptr;
        }
        cv_.notify_all();
    }
 
  private:
    alignas(64) std::atomic_int num_tasks_{ 0 };
    std::mutex mtx_;
    std::condition_variable cv_;
    std::exception_ptr exception_ptr_{ nullptr };
};
 
namespace detail {
 
template<typename T>
class RingBuffer
{
  public:
    explicit RingBuffer(size_t capacity)
      : buffer_{ std::unique_ptr<T[]>(new T[capacity]) }
      , capacity_{ capacity }
      , mask_{ capacity - 1 }
 
    {}
 
    size_t capacity() const { return capacity_; }
 
    void set_entry(size_t i, T val) { buffer_[i & mask_] = val; }
 
    T get_entry(size_t i) const { return buffer_[i & mask_]; }
 
    RingBuffer<T>* enlarged_copy(size_t bottom, size_t top) const
    {
        RingBuffer<T>* new_buffer = new RingBuffer{ 2 * capacity_ };
        for (size_t i = top; i != bottom; ++i)
            new_buffer->set_entry(i, this->get_entry(i));
        return new_buffer;
    }
 
  private:
    std::unique_ptr<T[]> buffer_;
    size_t capacity_;
    size_t mask_;
};
 
// exchange is not available in C++11, use implementatino from
// https://en.cppreference.com/w/cpp/utility/exchange
template<class T>
T
exchange(T& obj, T&& new_value) noexcept
{
    T old_value = std::move(obj);
    obj = std::forward<T>(new_value);
    return old_value;
}
 
class TaskQueue
{
    using Task = std::function<void()>;
 
  public:
    TaskQueue(size_t capacity = 256)
      : buffer_{ new RingBuffer<Task*>(capacity) }
    {}
 
    ~TaskQueue() noexcept
    {
        // must free memory allocated by push(), but not deallocated by pop()
        auto buf_ptr = buffer_.load();
        for (int i = top_; i < bottom_.load(m_relaxed); ++i)
            delete buf_ptr->get_entry(i);
        delete buf_ptr;
    }
 
    TaskQueue(TaskQueue const& other) = delete;
    TaskQueue& operator=(TaskQueue const& other) = delete;
 
    bool empty() const
    {
        return (bottom_.load(m_relaxed) <= top_.load(m_relaxed));
    }
 
    bool try_push(Task&& task)
    {
        {
            // must hold lock in case of multiple producers, abort if already
            // taken, so we can check out next queue
            std::unique_lock<std::mutex> lk(mutex_, std::try_to_lock);
            if (!lk)
                return false;
            this->push_unsafe(std::forward<Task>(task));
        }
        cv_.notify_one();
        return true;
    }
 
    void force_push(Task&& task)
    {
        {
            // must hold lock in case of multiple producers
            std::lock_guard<std::mutex> lk(mutex_);
            this->push_unsafe(std::forward<Task>(task));
        }
        cv_.notify_one();
    }
 
    void push_unsafe(Task&& task)
    {
        auto b = bottom_.load(m_relaxed);
        auto t = top_.load(m_acquire);
        RingBuffer<Task*>* buf_ptr = buffer_.load(m_relaxed);
 
        if (static_cast<int>(buf_ptr->capacity()) < (b - t) + 1) {
            // buffer is full, create enlarged copy before continuing
            old_buffers_.emplace_back(
              exchange(buf_ptr, buf_ptr->enlarged_copy(b, t)));
            buffer_.store(buf_ptr, m_relaxed);
        }
 
        buf_ptr->set_entry(b, new Task{ std::forward<Task>(task) });
        bottom_.store(b + 1, m_release);
    }
 
    bool try_pop(Task& task)
    {
        auto t = top_.load(m_acquire);
        std::atomic_thread_fence(m_seq_cst);
        auto b = bottom_.load(m_acquire);
 
        if (t < b) {
            // must load task pointer before acquiring the slot, because it
            // could be overwritten immediately after
            auto task_ptr = buffer_.load(m_acquire)->get_entry(t);
 
            if (top_.compare_exchange_strong(t, t + 1, m_seq_cst, m_relaxed)) {
                task = std::move(*task_ptr); // won race, get task
                delete task_ptr; // fre memory allocated in push_unsafe()
                return true;
            }
        }
        return false; // queue is empty or lost race
    }
 
    void wait()
    {
        std::unique_lock<std::mutex> lk(mutex_);
        cv_.wait(lk, [this] { return !this->empty() || stopped_; });
    }
 
    void stop()
    {
        {
            std::lock_guard<std::mutex> lk(mutex_);
            stopped_ = true;
        }
        cv_.notify_all();
    }
 
  private:
    alignas(64) std::atomic_int top_{ 0 };
    alignas(64) std::atomic_int bottom_{ 0 };
    alignas(64) std::atomic<RingBuffer<Task*>*> buffer_{ nullptr };
    std::vector<std::unique_ptr<RingBuffer<Task*>>> old_buffers_;
    std::mutex mutex_;
    std::condition_variable cv_;
    std::atomic<bool> stopped_;
 
    // convenience aliases
    static constexpr std::memory_order m_relaxed = std::memory_order_relaxed;
    static constexpr std::memory_order m_acquire = std::memory_order_acquire;
    static constexpr std::memory_order m_release = std::memory_order_release;
    static constexpr std::memory_order m_seq_cst = std::memory_order_seq_cst;
};
 
struct TaskManager
{
    std::vector<TaskQueue> queues_;
    size_t num_queues_;
    alignas(64) std::atomic_size_t push_idx_{ 0 };
    std::atomic_bool stopped_{ false };
    std::atomic_size_t todo_list_{ 0 };
 
    explicit TaskManager(size_t num_queues)
      : queues_{ std::vector<TaskQueue>(num_queues) }
      , num_queues_{ num_queues }
    {}
 
    template<typename Task>
    void push(Task&& task)
    {
        if (stopped_)
            return;
        for (size_t count = 0; count < num_queues_ * 20; count++) {
            if (queues_[push_idx_++ % num_queues_].try_push(task))
                return;
        }
        queues_[push_idx_++ % num_queues_].force_push(task);
    }
 
    template<typename Task>
    bool try_pop(Task& task, size_t worker_id = 0)
    {
        if (stopped_)
            return false;
        for (size_t k = 0; k <= num_queues_; k++) {
            if (queues_[(worker_id + k) % num_queues_].try_pop(task))
                return true;
        }
        return false;
    }
 
    void wait_for_jobs(size_t id) { queues_[id].wait(); }
 
    void stop()
    {
        for (auto& q : queues_)
            q.stop();
        stopped_ = true;
    }
 
    bool stopped() { return stopped_; }
};
 
} // end namespace detail
 
class ThreadPool
{
  public:
    ThreadPool()
      : ThreadPool(std::thread::hardware_concurrency())
    {}
 
    explicit ThreadPool(size_t n_workers)
      : task_manager_{ n_workers }
    {
        for (size_t id = 0; id < n_workers; ++id) {
            workers_.emplace_back([this, id] {
                std::function<void()> task;
                while (!task_manager_.stopped()) {
                    task_manager_.wait_for_jobs(id);
                    do {
                        // inner while to save a few cash misses calling empty()
                        if (task_manager_.try_pop(task, id))
                            this->execute_safely(task);
                    } while (!todo_list_.empty());
                }
            });
        }
    }
 
    ~ThreadPool()
    {
        task_manager_.stop();
        for (auto& worker : workers_) {
            if (worker.joinable())
                worker.join();
        }
    }
 
    ThreadPool(ThreadPool&&) = delete;
    ThreadPool(const ThreadPool&) = delete;
    ThreadPool& operator=(const ThreadPool&) = delete;
    ThreadPool& operator=(ThreadPool&& other) = delete;
 
    static ThreadPool& global_instance()
    {
        static ThreadPool instance_;
        return instance_;
    }
 
    template<class Function, class... Args>
    void push(Function&& f, Args&&... args)
    {
        if (workers_.size() == 0)
            return f(args...);
        todo_list_.add();
        task_manager_.push(
          std::bind(std::forward<Function>(f), std::forward<Args>(args)...));
    }
 
    template<class Function, class... Args>
    auto async(Function&& f, Args&&... args)
      -> std::future<decltype(f(args...))>
    {
        auto pack =
          std::bind(std::forward<Function>(f), std::forward<Args>(args)...);
        using pack_t = std::packaged_task<decltype(f(args...))()>;
        auto task_ptr = std::make_shared<pack_t>(std::move(pack));
        this->push([task_ptr] { (*task_ptr)(); });
        return task_ptr->get_future();
    }
 
    void wait() { todo_list_.wait(); }
 
  private:
    void execute_safely(std::function<void()>& task)
    {
        try {
            task();
            todo_list_.cross();
        } catch (...) {
            todo_list_.stop(std::current_exception());
            task_manager_.stop();
        }
    }
 
    detail::TaskManager task_manager_;
    TodoList todo_list_{ 0 };
    std::vector<std::thread> workers_;
};
 
 
template<class Function, class... Args>
void
push(Function&& f, Args&&... args)
{
    ThreadPool::global_instance().push(std::forward<Function>(f),
                                       std::forward<Args>(args)...);
}
 
template<class Function, class... Args>
auto
async(Function&& f, Args&&... args) -> std::future<decltype(f(args...))>
{
    return ThreadPool::global_instance().async(std::forward<Function>(f),
                                               std::forward<Args>(args)...);
}
 
inline void
wait()
{
    ThreadPool::global_instance().wait();
}
 
} // end namespace quickpool