executor.hpp

Go to the documentation of this file.

#pragma once
 
#include "aria/callback_boundary.hpp"
#include "aria/scheduler.hpp"
 
#include <atomic>
#include <condition_variable>
#include <coroutine>
#include <cstdint>
#include <deque>
#include <functional>
#include <future>
#include <iterator>
#include <memory>
#include <mutex>
#include <queue>
#include <stdexcept>
#include <thread>
#include <vector>
 
namespace aria::async {

class IExecutor : public virtual aria::IScheduler {
public:
    ~IExecutor() override = default;

    virtual void post(std::function<void()> fn) = 0;
 
    // ── IScheduler bridge ────────────────────────────────────────────
    [[nodiscard]] aria::SchedulerCaps caps() const noexcept override {
        return aria::SchedulerCaps::Post
             | aria::SchedulerCaps::GraphSafe
             | aria::SchedulerCaps::WorkerSafe;
    }
    void schedule(std::function<void()> fn) override {
        post(std::move(fn));
    }
 
    // ── Legacy capability shims (kept for source compatibility) ──────
    [[nodiscard]] virtual bool is_safe_graph_executor() const noexcept {
        return aria::has_caps(*this, aria::SchedulerCaps::GraphSafe);
    }

    [[nodiscard]] virtual bool is_safe_worker_executor() const noexcept {
        return aria::has_caps(*this, aria::SchedulerCaps::WorkerSafe);
    }
};

class ThreadPoolExecutor : public IExecutor {
public:
    explicit ThreadPoolExecutor(std::size_t threads = std::thread::hardware_concurrency())
        : stop_(false) {
        if (threads == 0) threads = 1;
        try {
            for (std::size_t i = 0; i < threads; ++i) {
                workers_.emplace_back([this]() { worker_loop_(); });
            }
        } catch (...) {
            // A failed thread creation must not destroy joinable std::threads
            // (which terminates the process). Stop and join the started prefix.
            {
                std::scoped_lock lk(mutex_);
                stop_ = true;
            }
            wake_sequence_.fetch_add(1, std::memory_order_release);
            wake_sequence_.notify_all();
            for (auto& worker : workers_) {
                if (worker.joinable()) { worker.join(); }
            }
            throw;
        }
    }
 
    ~ThreadPoolExecutor() override {
        // Drain: wait until the queue is empty AND no worker is running a
        // task.  Detached coroutines that hop off to another executor will
        // re-post back to us when they resume; each post increments
        // `inflight_`, so as long as we wait on `inflight_ == 0` we won't
        // tear down the threadpool out from under a pending resume.
        //
        // CONTRACT: callers MUST ensure no new work is posted to this pool
        // after they begin destroying it.  The standard pattern is:
        //   1. Cancel every CoroutineScope that launches work on this pool.
        //   2. Allow cancelled coroutines to throw/unwind (they may still
        //      post one final resume — `wait_idle` will wait for it).
        //   3. Destroy the pool.
        wait_idle();
        {
            std::lock_guard lk(mutex_);
            stop_ = true;
        }
        wake_sequence_.fetch_add(1, std::memory_order_release);
        wake_sequence_.notify_all();
        for (auto& t : workers_) if (t.joinable()) t.join();
    }
 
    void post(std::function<void()> fn) override {
        {
            std::lock_guard lk(mutex_);
            queue_.push(std::move(fn));
            inflight_.fetch_add(1, std::memory_order_relaxed);
        }
        wake_sequence_.fetch_add(1, std::memory_order_release);
        wake_sequence_.notify_one();
    }

    [[nodiscard]] aria::SchedulerCaps caps() const noexcept override {
        return aria::SchedulerCaps::Post
             | aria::SchedulerCaps::WorkerSafe
             | aria::SchedulerCaps::Autonomous;
    }
 
    [[nodiscard]] std::size_t worker_count() const noexcept { return workers_.size(); }

    void wait_idle() {
        while (true) {
            int pending = 0;
            {
                // Queue publication and its count increment are one locked
                // operation. Sample both before waiting outside the lock.
                std::lock_guard lk(mutex_);
                pending = inflight_.load(std::memory_order_acquire);
                if (queue_.empty() && pending == 0) { return; }
            }
            // Completion between unlock and wait changes the sampled value,
            // so a notification cannot be lost in that handoff window.
            inflight_.wait(pending, std::memory_order_acquire);
        }
    }
 
private:
    void worker_loop_() {
        while (true) {
            std::function<void()> fn;
            {
                std::unique_lock lk(mutex_);
                while (!stop_ && queue_.empty()) {
                    const auto sequence = wake_sequence_.load(std::memory_order_acquire);
                    lk.unlock();
                    wake_sequence_.wait(sequence, std::memory_order_acquire);
                    lk.lock();
                }
                if (stop_ && queue_.empty()) return;
                fn = std::move(queue_.front());
                queue_.pop();
            }
            try {
                fn();
            } catch (...) {
                aria::report_callback_failure(
                    std::string_view{"executor.thread_pool.worker"},
                    std::current_exception());
            }
            // Decrement AFTER the task has finished (not when we dequeued).
            if (inflight_.fetch_sub(1, std::memory_order_acq_rel) == 1) {
                // Atomic waiters need no queue-lock handshake on completion.
                inflight_.notify_all();
            }
        }
    }
 
    std::vector<std::thread> workers_;
    std::queue<std::function<void()>> queue_;
    std::mutex mutex_;
    // Publication epochs avoid the condition-variable re-lock handshake.
    // A worker samples the epoch while holding the queue lock, then waits
    // outside it: a post between unlock and wait changes the sampled value.
    std::atomic<std::uint64_t> wake_sequence_{0};
    // Reserved to retain the existing ThreadPoolExecutor object layout.
    [[maybe_unused]] std::condition_variable idle_cv_;
    std::atomic<int> inflight_{0};
    bool stop_;
};

class InlineExecutor : public IExecutor {
public:
    void post(std::function<void()> fn) override {
        if (!fn) return;
        try {
            fn();
        } catch (...) {
            aria::report_callback_failure(
                std::string_view{"executor.inline.post"},
                std::current_exception());
        }
    }
 
    // Inline runs synchronously on the caller's thread, so it cannot
    // claim main-thread affinity, pumping, or autonomy. The specific
    // "Inline graph + non-Inline worker" race is rejected independently
    // by an explicit dynamic_cast in detail::check_executor_safety_runtime
    // — see async_command.hpp.
    [[nodiscard]] aria::SchedulerCaps caps() const noexcept override {
        return aria::SchedulerCaps::Post
             | aria::SchedulerCaps::GraphSafe
             | aria::SchedulerCaps::WorkerSafe;
    }
};

class MainThreadExecutor : public IExecutor {
public:
    MainThreadExecutor() = default;
    MainThreadExecutor(const MainThreadExecutor&) = delete;
    MainThreadExecutor& operator=(const MainThreadExecutor&) = delete;
    MainThreadExecutor(MainThreadExecutor&&) = delete;
    MainThreadExecutor& operator=(MainThreadExecutor&&) = delete;
 
    ~MainThreadExecutor() override {
        std::deque<std::function<void()>> retired;
        {
            std::scoped_lock lk(m_);
            closed_ = true;
            retired.swap(queue_);
        }
        // Reject work posted by capture destructors while all members are live.
    }
 
    void post(std::function<void()> fn) override {
        if (!fn) { return; }
        {
            std::lock_guard lk(m_);
            if (closed_) { return; }
            queue_.push_back(std::move(fn));
        }
        // An already-bound owner cannot simultaneously be asleep in its own
        // pump. Foreign producers and an unbound executor must still wake it.
        if (owner_.load(std::memory_order_acquire) != std::this_thread::get_id()) {
            cv_.notify_one();
        }
    }

    [[nodiscard]] aria::SchedulerCaps caps() const noexcept override {
        return aria::SchedulerCaps::Post
             | aria::SchedulerCaps::GraphSafe
             | aria::SchedulerCaps::WorkerSafe
             | aria::SchedulerCaps::MainThread
             | aria::SchedulerCaps::Pumpable;
    }
 
    [[nodiscard]] bool is_main_thread() const noexcept override {
        return is_owner_thread();
    }

    std::size_t drain() {
        bind_owner_();
        std::size_t total = 0;
        while (true) {
            std::function<void()> single;
            std::vector<std::function<void()>> local;
            {
                std::lock_guard lk(m_);
                if (queue_.empty()) { break; }
                if (queue_.size() == 1) {
                    single = std::move(queue_.front());
                    queue_.pop_front();
                } else {
                    local.reserve(queue_.size());
                    std::move(queue_.begin(), queue_.end(), std::back_inserter(local));
                    queue_.clear();
                }
            }
            // The common coroutine handoff carries one callable. Keep it on
            // the stack rather than allocating a vector for every resume.
            // Multi-callable batches still detach together before callbacks.
            if (single) {
                invoke_drained(single);
                ++total;
            } else {
                for (auto& fn : local) {
                    invoke_drained(fn);
                    ++total;
                }
            }
        }
        return total;
    }

    template<typename Pred>
    bool pump_until(Pred predicate,
                    std::chrono::milliseconds timeout = std::chrono::seconds{2}) {
        bind_owner_();
        const auto deadline = std::chrono::steady_clock::now() + timeout;
        while (true) {
            drain();
            if (predicate()) return true;
            const auto now = std::chrono::steady_clock::now();
            if (now >= deadline) return false;
            std::unique_lock lk(m_);
            cv_.wait_until(lk, deadline, [this]{ return !queue_.empty(); });
            // Loop: if predicate is already true the next drain is a no-op
            // and we return true. If we woke on timeout the next iteration
            // will see `now >= deadline` and bail.
        }
    }

    void run_one() {
        bind_owner_();
        std::function<void()> fn;
        {
            std::unique_lock lk(m_);
            cv_.wait(lk, [this]{ return !queue_.empty(); });
            fn = std::move(queue_.front());
            queue_.pop_front();
        }
        try {
            fn();
        } catch (...) {
            aria::report_callback_failure(
                std::string_view{"executor.main_thread.run_one"},
                std::current_exception());
        }
    }

    [[nodiscard]] bool is_owner_thread() const noexcept {
        const auto id = owner_.load(std::memory_order_acquire);
        return id == std::thread::id{} || id == std::this_thread::get_id();
    }
 
    [[nodiscard]] std::size_t pending() const noexcept {
        std::lock_guard lk(m_);
        return queue_.size();
    }

    void clear() {
        bind_owner_();
        std::deque<std::function<void()>> retired;
        {
            std::lock_guard lk(m_);
            retired.swap(queue_);
        }
        // Destroy user captures after unlocking: their destructors may post.
 
    }
 
private:
    static void invoke_drained(std::function<void()>& fn) {
        try {
            fn();
        } catch (...) {
            aria::report_callback_failure(
                std::string_view{"executor.main_thread.drain"},
                std::current_exception());
        }
    }
 
    void bind_owner_() {
        const auto self = std::this_thread::get_id();
        if (owner_.load(std::memory_order_acquire) == self) { return; }
        std::thread::id expected{};
        if (owner_.compare_exchange_strong(expected, self,
                                           std::memory_order_acq_rel)) {
            return;  // we just claimed ownership
        }
        // Already bound — must match.
        if (expected != self) {
            throw std::logic_error(
                "MainThreadExecutor must be pumped or cleared on its owner thread");
        }
    }
 
    mutable std::mutex m_;
    std::condition_variable cv_;
    std::deque<std::function<void()>> queue_;
    std::atomic<std::thread::id> owner_{};
    bool closed_ = false;
};

inline auto schedule_on(IExecutor& exec) {
    struct Awaiter {
        IExecutor& exec;
        bool await_ready() const noexcept { return false; }
        void await_suspend(std::coroutine_handle<> h) const {
            exec.post([h]() mutable { h.resume(); });
        }
        void await_resume() const noexcept {}
    };
    return Awaiter{exec};
}
 
}  // namespace aria::async