I implemented work stealing queue inspired by Sean Parent's talk on code::dive 2016. Full implementation is here. I am looking to get feedback on improvements to make code more effective and common best practices.
TaskQueue
is thread safe task queue. It has pair of blocking Pop/Push
methods and non-blocking ones. Non-blocking methods are used in work stealing queue thread pool. Push
returns std::future
object. When m_done
flag is true tasks cannot be popped from the queue.
TaskQueue.h
#pragma once
#include <deque>
#include <optional>
#include <functional>
#include <condition_variable>
#include <future>
class TaskQueue
{
public:
using TaskType = std::function<void()>;
TaskQueue() = default;
~TaskQueue() = default;
TaskQueue(TaskQueue&&) = default;
TaskQueue& operator=(TaskQueue&&) = default;
void SetDone(bool done)
{
{
LockType lock{ m_mutex };
m_done = done;
}
if (done)
m_ready.notify_all();
}
auto IsDone() const
{
LockType lock{ m_mutex };
return m_done;
}
auto WaitAndPop(TaskType& task)
{
LockType lock{ m_mutex };
while (m_queue.empty() && !m_done)
m_ready.wait(lock);
if (!m_queue.empty() && !m_done)
{
task = std::move(m_queue.front());
m_queue.pop_front();
return true;
}
return false;
}
template<typename TTask>
auto Push(TTask&& task) // -> std::future<decltype(task())>
{
using TaskReturnType = decltype(task());
// std::packaged_task<> is move only type.
// We need to wrap it in a shared_ptr:
auto packagedTask = std::make_shared<std::packaged_task<TaskReturnType()>>(std::forward<TTask>(task));
auto future = packagedTask->get_future();
{
LockType lock{ m_mutex };
m_queue.emplace_back([packagedTask] { (*packagedTask)(); });
}
m_ready.notify_one();
return future;
}
auto TryPop(TaskType& task)
{
LockType lock{ m_mutex, std::try_to_lock };
if (!lock || m_queue.empty() || m_done)
return false;
task = move(m_queue.front());
m_queue.pop_front();
return true;
}
template<typename TTask>
auto TryPush(TTask&& task) -> std::optional<std::future<decltype(task())>>
{
using TaskReturnType = decltype(task());
// std::packaged_task<void()> is not movable
// We need to wrap it in a shared_ptr:
auto packagedTask = std::make_shared<std::packaged_task<TaskReturnType()>>(std::forward<TTask>(task));
auto future = packagedTask->get_future();
{
LockType lock{ m_mutex, std::try_to_lock };
if (!lock)
return {};
m_queue.emplace_back([packagedTask]() { (*packagedTask)(); });
}
m_ready.notify_one();
return future;
}
private:
using LockType = std::unique_lock<std::mutex>;
std::deque<TaskType> m_queue;
bool m_done{ false };
mutable std::mutex m_mutex;
std::condition_variable m_ready;
TaskQueue(const TaskQueue&) = delete;
TaskQueue& operator=(const TaskQueue&) = delete;
};
Work-stealing queue thread pool. Each thread has its own task queue. The main idea of work stealing is to Pop/Push
tasks in queue that is not locked by another thread.
WorkStealingQueueThreadPool.h
#pragma once
#include "TaskQueue.h"
#include <algorithm>
#include <thread>
class WorkStealingQueueThreadPool
{
public:
explicit WorkStealingQueueThreadPool(size_t threadCount = std::max(2u, std::thread::hardware_concurrency()));
~WorkStealingQueueThreadPool();
template<typename TaskT>
auto ExecuteAsync(TaskT&& task)
{
const auto index = m_queueIndex++;
for (size_t n = 0; n != m_queueCount*m_tryoutCount; ++n)
{
auto result = m_queues[(index + n) % m_queueCount].TryPush(std::forward<TaskT>(task));
if (result.has_value())
return std::move(*result);
}
return m_queues[index % m_queueCount].Push(std::forward<TaskT>(task));
}
private:
void Run(size_t queueIndex);
std::vector<TaskQueue> m_queues;
size_t m_queueIndex{ 0 };
const size_t m_queueCount;
const size_t m_tryoutCount{ 1 };
std::vector<std::thread> m_threads;
};
WorkStealingQueueThreadPool::WorkStealingQueueThreadPool(size_t threadCount)
: m_queues{ threadCount }
, m_queueCount{ threadCount }
{
for (size_t index = 0; index != threadCount; ++index)
m_threads.emplace_back([this, index] { Run(index); });
}
WorkStealingQueueThreadPool::~WorkStealingQueueThreadPool()
{
for (auto& queue : m_queues)
queue.SetDone(true);
for (auto& thread : m_threads)
thread.join();
}
void WorkStealingQueueThreadPool::Run(size_t queueIndex)
{
while (!m_queues[queueIndex].IsDone())
{
TaskQueue::TaskType task;
for (size_t n = 0; n != m_queueCount*m_tryoutCount; ++n)
{
if (m_queues[(queueIndex + n) % m_queueCount].TryPop(task))
break;
}
if (!task && !m_queues[queueIndex].WaitAndPop(task))
return;
task();
}
}
pop
by any thread in pool. \$\endgroup\$