Thread Pools Work Stealing

Long running Task

  • One task takes a long time to perform

    • It cannot execute the following tasks in the queue until this task completes
    • The following tasks will be delayed
  • Other threads have no work to do

    • Inefficient use of resources

Work Stealing

  • A refinement of work sharing

  • If a thread's queue is empty, the thread "steals" a task

    • The thread takes a task from another thread's queue
  • This ensures that threads are never idle

    • Provided there is enough work for all the threads
  • Can perform better than work sharing

    • If some tasks take much longer than others

Work Stealing Strategy

  • If a thread's queue is empty

    • Do not wait for a task to arrive on the queue
    • Choose another thread's queue at random
    • If there is a task on that queue, pop it and execute it
    • Otherwise, choose a different thread's queue at random
    • Continue until it finds a task to perform
  • If all the queues are empty

    • Pause for a while
    • Then repeat the process

Work Stealing Implementation

  • Add non-blocking functions to the queue

    • Return immediately if they cannot obtain a lock
    • try_pop() returns immediately if the queue is empty
    • try_push() returns immediately if the queue is full
  • The thread pool uses these non-blocking functions

    • worker() tries to find a queue where try_pop() succeeds
    • submit() tries to find a queue where try_push() succeeds

concurrent_queue.h


#pragma once

#include <mutex>
#include <queue>

using namespace std::chrono_literals;

template <class T> class concurrent_queue {
  std::timed_mutex mut;
  std::queue<T> que;
  size_t max{50};

public:
  concurrent_queue() = default;
  concurrent_queue(size_t max) : max(max) {};

  concurrent_queue(concurrent_queue &&) = delete;
  concurrent_queue(const concurrent_queue &) = delete;
  concurrent_queue &operator=(concurrent_queue &&) = delete;
  concurrent_queue &operator=(const concurrent_queue &) = delete;

  bool try_push(T value) {
    // Lock the mutex with a time-out
    std::unique_lock<std::timed_mutex> lck_guard(mut, std::defer_lock);

    // Cannot lock - return immediately
    if (!lck_guard.try_lock_for(1ms) || que.size() > max) {
      return false;
    }

    // Locked - add the element to the queue
    que.push(value);
    return true;
  }

  bool try_pop(T &value) {
    // Lock the mutex with a time-out
    std::unique_lock<std::timed_mutex> lck_guard(mut, std::defer_lock);

    // Cannot lock - return immediately
    if (!lck_guard.try_lock_for(1ms) || que.empty()) {
      return false;
    }

    // Locked - remove front element from the queue
    value = que.front();
    que.pop();
    return true;
  }
};

thr_pool.h


#include "concurrent_queue.h"
#include <cstddef>
#include <functional>
#include <memory>
#include <mutex>
#include <random>
#include <thread>
#include <vector>

// Type alias to simplify the code
// All the task functions will have this type
using Func = std::function<void()>;

// Alias for concurrent_queue type
using Queue = concurrent_queue<Func>;

class thread_pool {
  std::mt19937 mt;

  // Each thread has its own queue of task functions
  std::unique_ptr<Queue[]> work_queues;
  std::vector<std::thread> threads; // Vector of thread objects

  void worker(size_t index); // Their entry point function

  // returns a random number between 0 and thread_count - 1
  unsigned int get_random();

  unsigned int thread_count;

  std::mutex rand_mut;

public:
  thread_pool();
  ~thread_pool();

  // Add a task to the queue
  void submit(Func func);
};

thr_pool.cpp


#include "thr_pool.h"
#include <iostream>
#include <memory>
#include <random>
#include <thread>

thread_pool::thread_pool() {
  thread_count = std::thread::hardware_concurrency() - 1;
  std::cout << "Creating a thread pool with " << thread_count << " threads\n";

  // Create a dynamic array of queues
  work_queues = std::make_unique<Queue[]>(thread_count);

  // Start the threads
  for (unsigned int i = 0; i < thread_count; ++i) {
    threads.push_back(std::thread{&thread_pool::worker, this, i});
  }
}

thread_pool::~thread_pool() {
  // Wait for the threads to finish
  for (auto &thr : threads)
    thr.join();
}

unsigned int thread_pool::get_random() {
  std::lock_guard<std::mutex> lck_guard(rand_mut);
  std::uniform_int_distribution<unsigned int> dist(0, thread_count - 1);
  return dist(mt);
}

// Entry point function for the threads
void thread_pool::worker(size_t index) {
  while (true) {
    // Number of queues we have checked so far
    unsigned int visited = 0;

    // Take a task function off our queue
    size_t i = index;
    Func task;

    while (!work_queues[i].try_pop(task)) {
      // Nothing on this queue. Pick another queue at random
      i = get_random();

      // Hot loop avoidance
      // If we have checked "enough" queues, pause for a while
      // then start again with our own queue
      if (++visited == thread_count) {
        std::this_thread::sleep_for(10ms);
        visited = 0;
        i = index;
      }
    }

    // Invoke the task function
    task();
  }
}

// Choose a thread's queue and add a task to it
void thread_pool::submit(Func func) {
  unsigned int i;

  do {
    i = get_random();
  } while (!work_queues[i].try_push(func));
}

thread_pool_main.cpp


#include "thr_pool.h"
#include <iostream>
#include <thread>

using namespace std::chrono_literals;

// A task function
void task() {
  std::cout << "Thread id: " << std::this_thread::get_id()
            << " starting a task\n";
  std::this_thread::sleep_for(100ms);
  std::cout << "Thread id: " << std::this_thread::get_id()
            << " finishing a task\n";
}

// A long running task function
void task2() {
  std::cout << "Thread id: " << std::this_thread::get_id()
            << "starting a task\n";
  std::this_thread::sleep_for(5s);
  std::cout << "Thread id: " << std::this_thread::get_id()
            << " finishing a task\n";
}

int main() {
  // Create the thread pool
  thread_pool pool;

  // Send some tasks to the thread pool
  pool.submit(task2);
  for (int i = 0; i < 200; ++i)
    pool.submit(task);

  pool.submit([&pool]() {
    std::this_thread::sleep_for(6s);
    std::cout << "All tasks completed \n";
  });
}