samedi 28 janvier 2023

Thread pool with job queue gets stuck

I want to split jobs among multiple std::thread workers and continue once they are all done. To do so, I implemented a thread pool class mainly based on this SO answer. I noticed, however, that my benchmarks can get stuck, running forever, without any errors thrown.

I wrote a minimal reproducing code, enclosed at the end. Based on terminal output, the issue seems to occur when the jobs are being queued. I checked videos (1, 2), documentation (3) and blog posts (4). I tried replacing the type of the locks, using atomics. I could not find the underlying cause.

Here is the snippet to replicate the issue. The program repeatedly counts the odd elements in the test vector.

#include <atomic>
#include <condition_variable>
#include <functional>
#include <iostream>
#include <mutex>
#include <queue>
#include <thread>
#include <vector>

class Pool {
  public:
    const int worker_count;
    bool to_terminate = false;
    std::atomic<int> unfinished_tasks = 0;
    std::mutex mutex;
    std::condition_variable condition;
    std::vector<std::thread> threads;
    std::queue<std::function<void()>> jobs;

    void thread_loop()
    {
        while (true) {
            std::function<void()> job;
            {
                std::unique_lock<std::mutex> lock(mutex);
                condition.wait(lock, [&] { return (!jobs.empty()) || to_terminate; });

                if (to_terminate)
                    return;

                job = jobs.front();
                jobs.pop();
            }
            job();
            unfinished_tasks -= 1;
        }
    }

  public:
    Pool(int size) : worker_count(size)
    {
        if (size < 0)
            throw std::invalid_argument("Worker count needs to be a positive integer");

        for (int i = 0; i < worker_count; ++i)
            threads.push_back(std::thread(&Pool::thread_loop, this));
    };

    ~Pool()
    {
        {
            std::unique_lock lock(mutex);
            to_terminate = true;
        }
        condition.notify_all();
        for (auto &thread : threads)
            thread.join();
        threads.clear();
    };

    void queue_job(const std::function<void()> &job)
    {
        {
            std::unique_lock<std::mutex> lock(mutex);
            jobs.push(job);
            unfinished_tasks += 1;
            // std::cout << unfinished_tasks;
        }
        condition.notify_one();
    }

    void wait()
    {
        while (unfinished_tasks) {
            ; // spinlock
        };
    }
};

int main()
{
    constexpr int worker_count = 8;
    constexpr int vector_size = 1 << 10;
    Pool pool = Pool(worker_count);

    std::vector<int> test_vector;
    test_vector.reserve(vector_size);
    for (int i = 0; i < vector_size; ++i)
        test_vector.push_back(i);

    std::vector<int> worker_odd_counts(worker_count, 0);

    std::function<void(int)> worker_task = [&](int thread_id) {
        int chunk_size = vector_size / (worker_count) + 1;
        int my_start = thread_id * chunk_size;
        int my_end = std::min(my_start + chunk_size, vector_size);

        int local_odd_count = 0;
        for (int ii = my_start; ii < my_end; ++ii)
            if (test_vector[ii] % 2 != 0)
                ++local_odd_count;

        worker_odd_counts[thread_id] = local_odd_count;
    };

    for (int iteration = 0;; ++iteration) {
        std::cout << "Jobs.." << std::flush;
        for (int i = 0; i < worker_count; ++i)
            pool.queue_job([&worker_task, i] { worker_task(i); });
        std::cout << "..queued. " << std::flush;

        pool.wait();

        int odd_count = 0;
        for (auto elem : worker_odd_counts)
            odd_count += elem;

        std::cout << "Iter:" << iteration << ". Odd:" << odd_count << '\n';
    }
}

Here is the terminal output of one specific run:

[...]
Jobs....queued. Iter:2994. Odd:512
Jobs....queued. Iter:2995. Odd:512
Jobs..

Aucun commentaire:

Enregistrer un commentaire