task_group.hpp

#pragma once
 
#include <atomic>
#include <cfenv>
#include <condition_variable>
#include <cstddef>
#include <deque>
#include <exception>
#include <memory>
#include <mutex>
#include <optional>
#include <thread>
#include <type_traits>
#include <utility>
 
#if defined(__unix__) || defined(__APPLE__)
#include <pthread.h>
#define WMTK_WORKER_POOL_RESET_ON_FORK 1
#endif
 
namespace wmtk::threading {
 
namespace detail {
 
struct job
{
    virtual ~job() = default;
    virtual void run() noexcept = 0;
    virtual void finish() noexcept = 0;
};
 
class worker_pool
{
public:
    static worker_pool& instance()
    {
#ifdef WMTK_WORKER_POOL_RESET_ON_FORK
        static const bool registered = (pthread_atfork(nullptr, nullptr, &abandon_in_child), true);
        (void)registered;
#endif
        auto& slot = current();
        worker_pool* p = slot.load(std::memory_order_acquire);
        if (p != nullptr) return *p;
        // Created lazily, so that a forked child can start over. Threads racing to create it
        // agree on one; the losers' pools have no workers yet and are simply deleted.
        auto* fresh = new worker_pool();
        if (slot.compare_exchange_strong(p, fresh, std::memory_order_acq_rel)) return *fresh;
        delete fresh;
        return *p;
    }
 
    void submit(std::unique_ptr<job> j)
    {
        {
            std::unique_lock<std::mutex> lock(m_mutex);
            // Every queued job must have an idle worker of its own: a worker that is busy may be
            // running a task that waits for this one. A worker being started is already owed to
            // the job that started it, so it does not count as free here: another thread taking
            // it in the meantime would leave that job with no worker. Waking an idle worker that
            // then finds the queue already drained by a worker finishing early is harmless; it
            // waits again.
            if (m_jobs.size() + m_spawning < m_idle) {
                m_jobs.push_back(std::move(j));
                lock.unlock();
                m_cv.notify_one();
                return;
            }
            ++m_spawning;
        }
        // No idle worker is free for this job. Start one BEFORE queueing the job, so that a
        // failure to start it (std::system_error) leaves nothing behind for wait() to hang on.
        // If an idle worker turned up in the meantime, one of the two just waits for the next job.
        try {
            std::thread([this] { work(); }).detach();
        } catch (...) {
            std::lock_guard<std::mutex> lock(m_mutex);
            --m_spawning;
            throw;
        }
        {
            std::lock_guard<std::mutex> lock(m_mutex);
            ++m_workers;
            --m_spawning;
            m_jobs.push_back(std::move(j));
        }
        m_cv.notify_one();
    }
 
    bool try_submit(std::unique_ptr<job>& j)
    {
        std::unique_lock<std::mutex> lock(m_mutex);
        if (m_jobs.size() + m_spawning >= m_idle) return false;
        m_jobs.push_back(std::move(j));
        lock.unlock();
        m_cv.notify_one();
        return true;
    }
 
    size_t worker_count()
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        return m_workers;
    }
 
private:
    worker_pool() = default;
 
    static std::atomic<worker_pool*>& current()
    {
        static std::atomic<worker_pool*> pool{nullptr};
        return pool;
    }
 
#ifdef WMTK_WORKER_POOL_RESET_ON_FORK
    static void abandon_in_child() { current().store(nullptr, std::memory_order_relaxed); }
#endif
 
    [[noreturn]] void work()
    {
        std::unique_lock<std::mutex> lock(m_mutex);
        ++m_idle;
        for (;;) {
            m_cv.wait(lock, [this] { return !m_jobs.empty(); });
            --m_idle;
            std::unique_ptr<job> j = std::move(m_jobs.front());
            m_jobs.pop_front();
            lock.unlock();
            // A task expects round-to-nearest, the mode the code that created its thread ran in.
            // A reused worker starts in whatever the previous task left behind, and the exact
            // predicates switch the rounding mode around their interval filters, so an exception
            // in the middle of one would leak an upward mode into every later task on this thread.
            std::fesetround(FE_TONEAREST);
            j->run();
            // Idle again BEFORE the task is reported finished: once it is, wait() may return and
            // its caller start the next group, which must find this worker free rather than start
            // another one. Until it waits again this worker does nothing that can block.
            lock.lock();
            ++m_idle;
            lock.unlock();
            j->finish();
            j.reset();
            lock.lock();
        }
    }
 
    std::mutex m_mutex;
    std::condition_variable m_cv;
    std::deque<std::unique_ptr<job>> m_jobs;
    size_t m_idle = 0; 
    size_t m_spawning = 0; 
    size_t m_workers = 0; 
};
 
} // namespace detail
 
// ---------------------------------------------------------------------------
// task_group: replaces tbb::task_group. run() hands the task to a worker of the persistent
// pool above (one worker per outstanding task, so the semantics are a thread per task);
// wait() blocks until every task run since the last wait has finished, then rethrows the first
// exception any of them threw.
// ---------------------------------------------------------------------------
class task_group
{
    std::mutex m_mutex;
    std::condition_variable m_done;
    size_t m_pending = 0;
    std::exception_ptr m_eptr;
 
public:
    task_group() = default;
    task_group(const task_group&) = delete;
    task_group& operator=(const task_group&) = delete;
 
    template <typename F>
    void run(F&& f)
    {
        // The task is copied or moved into its job here, on the calling thread: a task whose copy
        // or move throws makes run() throw -- as std::thread's constructor did -- before anything
        // is counted or queued.
        auto j = std::make_unique<task_job<std::decay_t<F>>>(*this, std::forward<F>(f));
        {
            std::lock_guard<std::mutex> lock(m_mutex);
            ++m_pending;
        }
        try {
            detail::worker_pool::instance().submit(std::move(j));
        } catch (...) {
            // Nothing was queued (see worker_pool::submit), so nothing will decrement this.
            std::lock_guard<std::mutex> lock(m_mutex);
            --m_pending;
            throw;
        }
    }
 
    template <typename F>
    bool try_run(F&& f)
    {
        std::unique_ptr<detail::job> j =
            std::make_unique<task_job<std::decay_t<F>>>(*this, std::forward<F>(f));
        {
            std::lock_guard<std::mutex> lock(m_mutex);
            ++m_pending;
        }
        if (detail::worker_pool::instance().try_submit(j)) return true;
        // Not queued, so nothing will report it finished.
        std::lock_guard<std::mutex> lock(m_mutex);
        --m_pending;
        return false;
    }
 
    void wait()
    {
        std::unique_lock<std::mutex> lock(m_mutex);
        m_done.wait(lock, [this] { return m_pending == 0; });
        if (m_eptr) {
            std::exception_ptr e = m_eptr;
            m_eptr = nullptr;
            lock.unlock();
            std::rethrow_exception(e);
        }
    }
 
    ~task_group()
    {
        std::unique_lock<std::mutex> lock(m_mutex);
        m_done.wait(lock, [this] { return m_pending == 0; });
    }
 
private:
    template <typename Fn>
    class task_job final : public detail::job
    {
        task_group& m_group;
        std::optional<Fn> m_fn;
 
    public:
        template <typename G>
        task_job(task_group& group, G&& fn)
            : m_group(group)
            , m_fn(std::in_place, std::forward<G>(fn))
        {}
 
        void run() noexcept override
        {
            try {
                (*m_fn)();
            } catch (...) {
                m_group.record(std::current_exception());
            }
            // Destroy the task's own state before it is reported finished: once it is, wait()
            // may return and the caller may tear down whatever the task captured by reference.
            m_fn.reset();
        }
 
        // After this the job touches nothing of the group's.
        void finish() noexcept override { m_group.finish_one(); }
    };
 
    void record(std::exception_ptr e)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        if (!m_eptr) {
            m_eptr = std::move(e);
        }
    }
 
    void finish_one()
    {
        // Notify while holding the lock, so wait() cannot return -- and the group cannot be
        // destroyed -- until this thread has let go of the group's mutex.
        std::lock_guard<std::mutex> lock(m_mutex);
        if (--m_pending == 0) {
            m_done.notify_all();
        }
    }
};
 
} // namespace wmtk::threading