executor.hpp

Go to the documentation of this file.

#pragma once
 
#include "aria/callback_boundary.hpp"
#include "aria/scheduler.hpp"
 
#include <atomic>
#include <cassert>
#include <condition_variable>
#include <coroutine>
#include <deque>
#include <functional>
#include <future>
#include <iterator>
#include <memory>
#include <mutex>
#include <queue>
#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;
        for (std::size_t i = 0; i < threads; ++i) {
            workers_.emplace_back([this]() { worker_loop_(); });
        }
    }
 
    ~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;
        }
        cv_.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);
        }
        cv_.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() {
        std::unique_lock lk(mutex_);
        idle_cv_.wait(lk, [this] {
            return queue_.empty() && inflight_.load() == 0;
        });
    }
 
private:
    void worker_loop_() {
        while (true) {
            std::function<void()> fn;
            {
                std::unique_lock lk(mutex_);
                cv_.wait(lk, [this] { return stop_ || !queue_.empty(); });
                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) {
                // Last one — notify a possibly-waiting destructor.
                std::lock_guard lk(mutex_);
                idle_cv_.notify_all();
            }
        }
    }
 
    std::vector<std::thread> workers_;
    std::queue<std::function<void()>> queue_;
    std::mutex mutex_;
    std::condition_variable cv_;       // wakes workers
    std::condition_variable idle_cv_;  // wakes wait_idle()
    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:
    void post(std::function<void()> fn) override {
        {
            std::lock_guard lk(m_);
            queue_.push_back(std::move(fn));
        }
        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::vector<std::function<void()>> local;
            {
                std::lock_guard lk(m_);
                if (queue_.empty()) break;
                local.reserve(queue_.size());
                std::move(queue_.begin(), queue_.end(), std::back_inserter(local));
                queue_.clear();
            }
            for (auto& fn : local) {
                try {
                    fn();
                } catch (...) {
                    aria::report_callback_failure(
                        std::string_view{"executor.main_thread.drain"},
                        std::current_exception());
                }
                ++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() noexcept {
        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:
    void bind_owner_() noexcept {
        std::thread::id expected{};
        const auto self = std::this_thread::get_id();
        if (owner_.compare_exchange_strong(expected, self,
                                           std::memory_order_acq_rel)) {
            return;  // we just claimed ownership
        }
        // Already bound — must match.
        assert(expected == self
               && "MainThreadExecutor pumped from a non-owner thread. "
                  "post() is fine from any thread, but drain/pump/run_one "
                  "must run on the thread that originally bound the executor.");
    }
 
    mutable std::mutex m_;
    std::condition_variable cv_;
    std::deque<std::function<void()>> queue_;
    std::atomic<std::thread::id> owner_{};
};

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