threadpool.h

Go to the documentation of this file.

// Copyright (c) The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or https://www.opensource.org/licenses/mit-license.php.

#ifndef BITCOIN_UTIL_THREADPOOL_H
#define BITCOIN_UTIL_THREADPOOL_H

#include <sync.h>
#include <tinyformat.h>
#include <util/expected.h>
#include <util/check.h>
#include <util/thread.h>

#include <algorithm>
#include <condition_variable>
#include <functional>
#include <future>
#include <queue>
#include <thread>
#include <type_traits>
#include <utility>
#include <vector>

class ThreadPool
{
private:
    std::string m_name;
    Mutex m_mutex;
    std::queue<std::packaged_task<void()>> m_work_queue GUARDED_BY(m_mutex);
    std::condition_variable m_cv;
    // Note: m_interrupt must be guarded by m_mutex, and cannot be replaced by an unguarded atomic bool.
    // This ensures threads blocked on m_cv reliably observe the change and proceed correctly without missing signals.
    // Ref: https://en.cppreference.com/w/cpp/thread/condition_variable
    bool m_interrupt GUARDED_BY(m_mutex){false};
    std::vector<std::thread> m_workers GUARDED_BY(m_mutex);

    void WorkerThread() EXCLUSIVE_LOCKS_REQUIRED(!m_mutex)
    {
        WAIT_LOCK(m_mutex, wait_lock);
        for (;;) {
            std::packaged_task<void()> task;
            {
                // Wait only if needed; avoid sleeping when a new task was submitted while we were processing another one.
                if (!m_interrupt && m_work_queue.empty()) {
                    // Block until the pool is interrupted or a task is available.
                    m_cv.wait(wait_lock, [&]() EXCLUSIVE_LOCKS_REQUIRED(m_mutex) { return m_interrupt || !m_work_queue.empty(); });
                }

                // If stopped and no work left, exit worker
                if (m_interrupt && m_work_queue.empty()) {
                    return;
                }

                task = std::move(m_work_queue.front());
                m_work_queue.pop();
            }

            {
                // Execute the task without the lock
                REVERSE_LOCK(wait_lock, m_mutex);
                task();
            }
        }
    }

public:
    explicit ThreadPool(const std::string& name) : m_name(name) {}

    ~ThreadPool()
    {
        Stop(); // In case it hasn't been stopped.
    }

    void Start(int num_workers) EXCLUSIVE_LOCKS_REQUIRED(!m_mutex)
    {
        assert(num_workers > 0);
        LOCK(m_mutex);
        if (m_interrupt) throw std::runtime_error("Thread pool has been interrupted or is stopping");
        if (!m_workers.empty()) throw std::runtime_error("Thread pool already started");

        // Create workers
        m_workers.reserve(num_workers);
        for (int i = 0; i < num_workers; i++) {
            m_workers.emplace_back(&util::TraceThread, strprintf("%s_pool_%d", m_name, i), [this] { WorkerThread(); });
        }
    }

    void Stop() EXCLUSIVE_LOCKS_REQUIRED(!m_mutex)
    {
        // Notify workers and join them
        std::vector<std::thread> threads_to_join;
        {
            LOCK(m_mutex);
            // Ensure Stop() is not called from a worker thread while workers are still registered,
            // otherwise a self-join deadlock would occur.
            auto id = std::this_thread::get_id();
            for (const auto& worker : m_workers) assert(worker.get_id() != id);
            // Early shutdown to return right away on any concurrent Submit() call
            m_interrupt = true;
            threads_to_join.swap(m_workers);
        }
        m_cv.notify_all();
        // Help draining queue
        while (ProcessTask()) {}
        // Free resources
        for (auto& worker : threads_to_join) worker.join();

        // Since we currently wait for tasks completion, sanity-check empty queue
        LOCK(m_mutex);
        Assume(m_work_queue.empty());
        // Re-allow Start() now that all workers have exited
        m_interrupt = false;
    }

    enum class SubmitError {
        Inactive,
        Interrupted,
    };

    template <class F>
    [[nodiscard]] util::Expected<std::future<std::invoke_result_t<F>>, SubmitError> Submit(F&& fn) noexcept EXCLUSIVE_LOCKS_REQUIRED(!m_mutex)
    {
        std::packaged_task<std::invoke_result_t<F>()> task{std::forward<F>(fn)};
        auto future{task.get_future()};
        {
            LOCK(m_mutex);
            if (m_workers.empty()) return util::Unexpected{SubmitError::Inactive};
            if (m_interrupt) return util::Unexpected{SubmitError::Interrupted};

            m_work_queue.emplace(std::move(task));
        }
        m_cv.notify_one();
        return {std::move(future)};
    }

    bool ProcessTask() EXCLUSIVE_LOCKS_REQUIRED(!m_mutex)
    {
        std::packaged_task<void()> task;
        {
            LOCK(m_mutex);
            if (m_work_queue.empty()) return false;

            // Pop the task
            task = std::move(m_work_queue.front());
            m_work_queue.pop();
        }
        task();
        return true;
    }

    void Interrupt() EXCLUSIVE_LOCKS_REQUIRED(!m_mutex)
    {
        WITH_LOCK(m_mutex, m_interrupt = true);
        m_cv.notify_all();
    }

    size_t WorkQueueSize() EXCLUSIVE_LOCKS_REQUIRED(!m_mutex)
    {
        return WITH_LOCK(m_mutex, return m_work_queue.size());
    }

    size_t WorkersCount() EXCLUSIVE_LOCKS_REQUIRED(!m_mutex)
    {
        return WITH_LOCK(m_mutex, return m_workers.size());
    }
};

constexpr std::string_view SubmitErrorString(const ThreadPool::SubmitError err) noexcept {
    switch (err) {
        case ThreadPool::SubmitError::Inactive:
            return "No active workers";
        case ThreadPool::SubmitError::Interrupted:
            return "Interrupted";
    }
    Assume(false); // Unreachable
    return "Unknown error";
}

#endif // BITCOIN_UTIL_THREADPOOL_H