This is a follow up to my previous post which also follows up on my first post regarding my thread pool implementation. I have since made some further changes and attempted to improve performance with work stealing.
The most up to date version of the code is available on my Github.
Previously, I removed the use of multiple queues and went with a single queue in conjunction with a std::condition_variable but have now moved back to using multiple queues combined with a std::binary_semaphore as the signal for the thread to do work. I've also attempted to implement a basic work stealing strategy and it seems to have improved the base performance of the thread pool compared to a using a single queue without work stealing.
I think the next logical improvement to make would be to use a lock free queue but I still need to do some testing to see if the performance gain would be worth it.
Something I wasn't able to solve is allowing users to pass a refernece to a lambda function that is given to the thread pool. For example, the following code would not work. x would never be modified.
int x = 2;
pool.enqueue_detach([](int &value) { value = 4; }, std::ref(x));
thread_pool.h
#pragma once
#include <atomic>
#include <concepts>
#include <deque>
#include <functional>
#include <future>
#include <memory>
#include <semaphore>
#include <thread>
#include <type_traits>
#include "thread_pool/thread_safe_queue.h"
namespace dp {
template <typename FunctionType = std::function<void()>>
requires std::invocable<FunctionType> &&
std::is_same_v<void, std::invoke_result_t<FunctionType>>
class thread_pool {
public:
explicit thread_pool(
const unsigned int &number_of_threads = std::thread::hardware_concurrency())
: tasks_(number_of_threads) {
for (std::size_t i = 0; i < number_of_threads; ++i) {
threads_.emplace_back([&, id = i](const std::stop_token &stop_tok) {
do {
// wait until signaled
tasks_[id].signal.acquire();
do {
// invoke the task
while (auto task = tasks_[id].tasks.pop()) {
try {
pending_tasks_.fetch_sub(1, std::memory_order_release);
std::invoke(std::move(task.value()));
} catch (...) {
}
}
// try to steal a task
for (std::size_t j = 1; j < tasks_.size(); ++j) {
const std::size_t index = (id + j) % tasks_.size();
if (auto task = tasks_[index].tasks.steal()) {
pending_tasks_.fetch_sub(1, std::memory_order_release);
std::invoke(std::move(task.value()));
}
}
} while (pending_tasks_.load(std::memory_order_acquire) > 0);
} while (!stop_tok.stop_requested());
});
}
}
~thread_pool() {
// stop all threads
for (std::size_t i = 0; i < threads_.size(); ++i) {
threads_[i].request_stop();
tasks_[i].signal.release();
threads_[i].join();
}
}
/// thread pool is non-copyable
thread_pool(const thread_pool &) = delete;
thread_pool &operator=(const thread_pool &) = delete;
template <typename Function, typename... Args,
typename ReturnType = std::invoke_result_t<Function &&, Args &&...>>
requires std::invocable<Function, Args...>
[[nodiscard]] std::future<ReturnType> enqueue(Function f, Args... args) {
auto shared_promise = std::make_shared<std::promise<ReturnType>>();
auto task = [func = std::move(f), ... largs = std::move(args),
promise = shared_promise]() {
try {
promise->set_value(func(largs...));
} catch (...) {
promise->set_exception(std::current_exception());
}
};
// get the future before enqueuing the task
auto future = shared_promise->get_future();
// enqueue the task
enqueue_task(std::move(task));
return future;
}
template <typename Function, typename... Args>
requires std::invocable<Function, Args...> &&
std::is_same_v<void, std::invoke_result_t<Function &&, Args &&...>>
void enqueue_detach(Function &&func, Args &&...args) {
enqueue_task(
std::move([f = std::forward<Function>(func),
... largs = std::forward<Args>(args)]() mutable -> decltype(auto) {
// suppress exceptions
try {
std::invoke(f, largs...);
} catch (...) {
}
}));
}
private:
template <typename Function>
void enqueue_task(Function &&f) {
const std::size_t i = count_++ % tasks_.size();
pending_tasks_.fetch_add(1, std::memory_order_relaxed);
tasks_[i].tasks.push(std::forward<Function>(f));
tasks_[i].signal.release();
}
struct task_item {
dp::thread_safe_queue<FunctionType> tasks{};
std::binary_semaphore signal{0};
};
std::vector<std::jthread> threads_;
std::deque<task_item> tasks_;
std::size_t count_{};
std::atomic_int_fast64_t pending_tasks_{};
};
} // namespace dp
And for clarity, my thread safe queue implementation:
thread_safe_queue.h
#pragma once
#include <deque>
#include <mutex>
#include <optional>
namespace dp {
template <typename T>
class thread_safe_queue {
public:
using value_type = T;
using size_type = typename std::deque<T>::size_type;
thread_safe_queue() = default;
void push(T&& value) {
std::lock_guard lock(mutex_);
data_.push_back(std::forward<T>(value));
}
[[nodiscard]] bool empty() const {
std::lock_guard lock(mutex_);
return data_.empty();
}
[[nodiscard]] std::optional<T> pop() {
std::lock_guard lock(mutex_);
if (data_.empty()) return std::nullopt;
auto front = data_.front();
data_.pop_front();
return front;
}
[[nodiscard]] std::optional<T> steal() {
std::lock_guard lock(mutex_);
if (data_.empty()) return std::nullopt;
auto back = data_.back();
data_.pop_back();
return back;
}
private:
using mutex_type = std::mutex;
std::deque<T> data_{};
mutable mutex_type mutex_{};
};
} // namespace dp
Any and all feedback is much appreciated. Would you use this in one of your projects? Why or why not? My goal is to have something that is performant but also reasonably robust for the majority of use cases. Working on this project has been a great learning experience in the context of concurrent programing. It is a very interesting, albeit deceptively difficult, field.
