manager.hpp

Go to the documentation of this file.

#pragma once
 
 
#include <threadman/config.hpp>
#include <threadman/display.hpp>
#include <threadman/exceptions.hpp>
#include <threadman/future_wait_pool.hpp>
#include <threadman/log.hpp>
#include <threadman/metrics.hpp>
#include <threadman/stats.hpp>
#include <threadman/task.hpp>
#include <threadman/thread.hpp>
#include <threadman/thread_pool.hpp>
 
#include <commons/display_info.hpp>
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <shared_mutex>
#include <string>
#include <string_view>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
 
namespace threadman {
 
class ThreadManager {
public:
    struct Options {
        std::string name = "tm::manager";
        std::chrono::milliseconds housekeeping_interval{THREADMAN_DEFAULT_SCALE_CHECK_INTERVAL_MS};
        std::chrono::milliseconds summary_interval{THREADMAN_DEFAULT_SUMMARY_INTERVAL_MS};
        bool start_housekeeper_eagerly = false;
        std::string default_pool_name = THREADMAN_DEFAULT_POOL_NAME;
        std::size_t default_pool_min_workers = 1;
        std::size_t default_pool_max_workers = std::thread::hardware_concurrency();
    };
 
    using SummaryListener = std::function<void(const ManagerSummary&)>;
    using StuckTaskListener = std::function<void(const StuckTaskEvent&)>;
 
    class SubscriptionToken {
    public:
        SubscriptionToken() noexcept = default;
        SubscriptionToken(const SubscriptionToken&) = delete;
        SubscriptionToken& operator=(const SubscriptionToken&) = delete;
        SubscriptionToken(SubscriptionToken&& other) noexcept
            : mgr_(other.mgr_), id_(other.id_), kind_(other.kind_) {
            other.mgr_ = nullptr;
            other.id_ = 0;
        }
        SubscriptionToken& operator=(SubscriptionToken&& other) noexcept {
            if (this != &other) {
                release();
                mgr_ = other.mgr_;
                id_ = other.id_;
                kind_ = other.kind_;
                other.mgr_ = nullptr;
                other.id_ = 0;
            }
            return *this;
        }
        ~SubscriptionToken() {
            release();
        }
 
        [[nodiscard]] bool valid() const noexcept {
            return mgr_ != nullptr && id_ != 0;
        }
        [[nodiscard]] std::uint64_t id() const noexcept {
            return id_;
        }
 
    private:
        friend class ThreadManager;
        enum class Kind : std::uint8_t { Summary, StuckTask };
        SubscriptionToken(ThreadManager* m, const std::uint64_t i, const Kind k) noexcept
            : mgr_(m), id_(i), kind_(k) {}
 
        void release() noexcept {
            if (mgr_ != nullptr && id_ != 0) {
                mgr_->unsubscribe(id_, kind_);
            }
            mgr_ = nullptr;
            id_ = 0;
        }
 
        ThreadManager* mgr_ = nullptr;
        std::uint64_t id_ = 0;
        Kind kind_ = Kind::Summary;
    };
 
    ThreadManager() : ThreadManager(Options{}) {}
    explicit ThreadManager(Options opts) : opts_(std::move(opts)) {
        if (opts_.start_housekeeper_eagerly) {
            start_housekeeper();
        }
    }
    ~ThreadManager() {
        stop_housekeeper();
        if (default_pool_) {
            try {
                default_pool_->shutdown();
            } catch (...) {
                // dtor must not propagate
            }
            default_pool_.reset();
        }
    }
 
    ThreadManager(const ThreadManager&) = delete;
    ThreadManager& operator=(const ThreadManager&) = delete;
    ThreadManager(ThreadManager&&) = delete;
    ThreadManager& operator=(ThreadManager&&) = delete;
 
    [[nodiscard]] static ThreadManager& instance() noexcept {
        // Construct the metrics subsystem (scope, all metric handles, prom's
        // adapter cell) *before* the singleton, so reverse-order static
        // destruction tears the singleton (and joins its housekeeper) down
        // first, while every metric it publishes is still alive.
        metrics::warm_up();
        static ThreadManager inst;
        return inst;
    }
 
    // --- Registration (called by ManagedThread / ThreadPool) ---------------
    void register_thread(const std::shared_ptr<ManagedThread::ControlBlock>& cb) {
        if (!cb) {
            return;
        }
        {
            std::unique_lock guard(threads_mtx_);
            threads_.insert_or_assign(cb->id, cb);
        }
        total_registered_.fetch_add(1, std::memory_order_relaxed);
        metrics::manager_threads_registered().inc();
    }
    void unregister_thread(const std::uint64_t thread_id) noexcept {
        std::unique_lock guard(threads_mtx_);
        threads_.erase(thread_id);
    }
    void register_pool(ThreadPool& p) {
        {
            std::scoped_lock guard(pools_mtx_);
            pools_.push_back(&p);
        }
        // Lazy housekeeper start on first pool registration.
        start_housekeeper();
    }
    void unregister_pool(const std::uint64_t pool_id) noexcept {
        std::scoped_lock guard(pools_mtx_);
        for (std::size_t i = 0; i < pools_.size(); ++i) {
            if ((pools_[i] != nullptr) && pools_[i]->id() == pool_id) {
                pools_.erase(pools_.begin() + static_cast<std::ptrdiff_t>(i));
                return;
            }
        }
    }
 
    // --- Snapshots ---------------------------------------------------------
    [[nodiscard]] std::vector<ThreadSnapshot> snapshot_threads() const {
        std::vector<ThreadSnapshot> out;
        std::shared_lock guard(threads_mtx_);
        out.reserve(threads_.size());
        for (const auto& weak : threads_ | std::views::values) {
            if (const auto cb = weak.lock()) {
                ThreadSnapshot snap;
                snap.state = cb->state.load(std::memory_order_acquire);
                snap.failed = (snap.state == ThreadState::Failed);
                std::scoped_lock lk(cb->mtx);
                snap.id = cb->id;
                snap.name = cb->name;
                snap.native_id = cb->native_id;
                snap.pool_id = cb->pool_id;
                snap.is_core = cb->is_core;
                if (cb->started_at.time_since_epoch().count() != 0) {
                    const auto end = cb->ended_at.time_since_epoch().count() == 0
                                         ? std::chrono::steady_clock::now()
                                         : cb->ended_at;
                    snap.run_duration = end - cb->started_at;
                }
                out.push_back(std::move(snap));
            }
        }
        return out;
    }
 
    [[nodiscard]] std::optional<ThreadSnapshot> find_by_id(const std::uint64_t id) const {
        std::shared_lock guard(threads_mtx_);
        if (const auto it = threads_.find(id); it != threads_.end()) {
            if (const auto cb = it->second.lock()) {
                ThreadSnapshot snap;
                snap.state = cb->state.load(std::memory_order_acquire);
                snap.failed = (snap.state == ThreadState::Failed);
                std::scoped_lock lk(cb->mtx);
                snap.id = cb->id;
                snap.name = cb->name;
                snap.native_id = cb->native_id;
                snap.pool_id = cb->pool_id;
                snap.is_core = cb->is_core;
                if (cb->started_at.time_since_epoch().count() != 0) {
                    const auto end = cb->ended_at.time_since_epoch().count() == 0
                                         ? std::chrono::steady_clock::now()
                                         : cb->ended_at;
                    snap.run_duration = end - cb->started_at;
                }
                return snap;
            }
        }
        return std::nullopt;
    }
 
    [[nodiscard]] std::vector<ThreadSnapshot> find_by_name(const std::string_view name) const {
        std::vector<ThreadSnapshot> out;
        std::shared_lock guard(threads_mtx_);
        for (const auto& weak : threads_ | std::views::values) {
            const auto cb = weak.lock();
            if (!cb) {
                continue;
            }
            std::scoped_lock lk(cb->mtx);
            if (cb->name == name) {
                ThreadSnapshot snap;
                snap.state = cb->state.load(std::memory_order_acquire);
                snap.failed = (snap.state == ThreadState::Failed);
                snap.id = cb->id;
                snap.name = cb->name;
                snap.native_id = cb->native_id;
                snap.pool_id = cb->pool_id;
                snap.is_core = cb->is_core;
                if (cb->started_at.time_since_epoch().count() != 0) {
                    const auto end = cb->ended_at.time_since_epoch().count() == 0
                                         ? std::chrono::steady_clock::now()
                                         : cb->ended_at;
                    snap.run_duration = end - cb->started_at;
                }
                out.push_back(std::move(snap));
            }
        }
        return out;
    }
 
    [[nodiscard]] std::vector<ThreadPoolStats> snapshot_pools() const {
        std::vector<ThreadPoolStats> out;
        std::scoped_lock guard(pools_mtx_);
        out.reserve(pools_.size());
        for (const auto* p : pools_) {
            if (p != nullptr) {
                out.push_back(p->stats());
            }
        }
        return out;
    }
 
    [[nodiscard]] ManagerSummary build_summary() const {
        ManagerSummary summary;
        summary.wall_clock = std::chrono::system_clock::now();
        summary.per_pool_stats = snapshot_pools();
        summary.total_pools = summary.per_pool_stats.size();
        for (const auto& s : summary.per_pool_stats) {
            summary.total_live_workers += s.workers;
            summary.total_queued += s.queued;
            summary.total_completed += s.completed;
            summary.total_failed += s.failed;
        }
        {
            std::shared_lock g(threads_mtx_);
            for (const auto& w : threads_ | std::views::values) {
                if (w.lock()) {
                    ++summary.live_threads;
                }
            }
        }
        return summary;
    }
 
    // --- Future-wait pool --------------------------------------------------
 
    [[nodiscard]] FutureWaitPool make_future_wait_pool(FutureWaitPoolOptions opts = {}) {
        return FutureWaitPool{std::move(opts), this};
    }
 
    // --- Default pool ------------------------------------------------------
    [[nodiscard]] ThreadPool& default_pool() {
        std::call_once(default_pool_once_, [this] {
            ThreadPoolOptions o;
            o.name = opts_.default_pool_name;
            o.min_workers = opts_.default_pool_min_workers;
            o.max_workers = opts_.default_pool_max_workers;
            o.manager = this;
            default_pool_ = std::make_unique<ThreadPool>(std::move(o));
        });
        return *default_pool_;
    }
 
    // --- Subscriptions ----------------------------------------------------
    [[nodiscard]] SubscriptionToken subscribe_summary(SummaryListener cb) {
        const auto id = next_sub_id_.fetch_add(1, std::memory_order_relaxed);
        {
            std::scoped_lock guard(subs_mtx_);
            summary_subs_.emplace_back(id, std::move(cb));
        }
        return SubscriptionToken{this, id, SubscriptionToken::Kind::Summary};
    }
    [[nodiscard]] SubscriptionToken subscribe_stuck_tasks(StuckTaskListener cb) {
        const auto id = next_sub_id_.fetch_add(1, std::memory_order_relaxed);
        {
            std::scoped_lock guard(subs_mtx_);
            stuck_subs_.emplace_back(id, std::move(cb));
        }
        return SubscriptionToken{this, id, SubscriptionToken::Kind::StuckTask};
    }
 
    // --- Housekeeper control ----------------------------------------------
    void start_housekeeper() {
        if (bool expected = false; !housekeeper_started_.compare_exchange_strong(
                expected, true, std::memory_order_acq_rel)) {
            return;
        }
        ManagedThread::Options topts;
        topts.name = opts_.name + ".housekeeper";
        topts.is_core = true;
        topts.manager = this;
        housekeeper_.emplace(std::move(topts),
                             [this](std::stop_token tok) { housekeeper_loop(std::move(tok)); });
        if (const auto& lg = log::manager()) {
            lg->info("started housekeeper (interval={}ms, summary={}ms)",
                     opts_.housekeeping_interval.count(),
                     opts_.summary_interval.count());
        }
    }
    void stop_housekeeper() noexcept {
        if (!housekeeper_started_.load(std::memory_order_acquire)) {
            return;
        }
        if (housekeeper_.has_value()) {
            housekeeper_->request_stop();
            {
                std::scoped_lock guard(hk_mtx_);
                hk_cv_.notify_all();
            }
            if (housekeeper_->joinable()) {
                housekeeper_->join();
            }
            housekeeper_.reset();
        }
        housekeeper_started_.store(false, std::memory_order_release);
    }
    [[nodiscard]] bool housekeeper_running() const noexcept {
        return housekeeper_started_.load(std::memory_order_acquire);
    }
 
    [[nodiscard]] std::size_t stuck_task_count() const noexcept {
        return known_stuck_size_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] static const comms::DisplayInfo& display_info() {
        static const comms::DisplayInfo info{
            .name = "ThreadManager",
            .description = "Central registry of managed threads and pools with periodic summary "
                           "and stuck-task event dispatch.",
            .icon = comms::Icon::from("mdi:account-supervisor"),
        };
        return info;
    }
 
private:
    void unsubscribe(const std::uint64_t id, const SubscriptionToken::Kind kind) noexcept {
        std::scoped_lock guard(subs_mtx_);
        if (kind == SubscriptionToken::Kind::Summary) {
            for (std::size_t i = 0; i < summary_subs_.size(); ++i) {
                if (summary_subs_[i].first == id) {
                    summary_subs_.erase(summary_subs_.begin() + static_cast<std::ptrdiff_t>(i));
                    return;
                }
            }
        } else {
            for (std::size_t i = 0; i < stuck_subs_.size(); ++i) {
                if (stuck_subs_[i].first == id) {
                    stuck_subs_.erase(stuck_subs_.begin() + static_cast<std::ptrdiff_t>(i));
                    return;
                }
            }
        }
    }
 
    void prune_expired_threads() {
        std::unique_lock guard(threads_mtx_);
        for (auto it = threads_.begin(); it != threads_.end();) {
            if (it->second.expired()) {
                it = threads_.erase(it);
            } else {
                ++it;
            }
        }
    }
 
    void housekeeper_loop(std::stop_token tok) {
        last_summary_at_ = std::chrono::steady_clock::now();
        std::unordered_set<std::uint64_t> known_stuck;
 
        while (!tok.stop_requested()) {
            {
                std::unique_lock lk(hk_mtx_);
                hk_cv_.wait_for(
                    lk, opts_.housekeeping_interval, [&] { return tok.stop_requested(); });
            }
            if (tok.stop_requested()) {
                break;
            }
 
            const auto now = std::chrono::steady_clock::now();
 
            // 1) Prune expired weak refs.
            prune_expired_threads();
 
            // 2) Pool scaling.
            {
                std::scoped_lock guard(pools_mtx_);
                for (auto* p : pools_) {
                    if (p != nullptr) {
                        try {
                            p->scale_tick(now);
                        } catch (const std::exception& e) {
                            if (auto& lg = log::manager()) {
                                lg->warn("scale_tick on pool '{}' threw: {}", p->name(), e.what());
                            }
                        } catch (...) {
                            // ignore
                        }
                    }
                }
            }
 
            // 3) Stuck-task detection.
            std::vector<StuckTaskEvent> all_stuck;
            {
                std::scoped_lock guard(pools_mtx_);
                for (auto* p : pools_) {
                    if (p == nullptr) {
                        continue;
                    }
                    for (auto stuck = p->detect_stuck_tasks(now); auto& s : stuck) {
                        all_stuck.push_back(std::move(s));
                    }
                }
            }
 
            // Dispatch new stuck-task events.
            std::vector<StuckTaskListener> stuck_listeners_copy;
            {
                std::scoped_lock guard(subs_mtx_);
                stuck_listeners_copy.reserve(stuck_subs_.size());
                for (const auto& cb : stuck_subs_ | std::views::values) {
                    stuck_listeners_copy.push_back(cb);
                }
            }
            for (const auto& ev : all_stuck) {
                if (known_stuck.insert(ev.task.id).second) {
                    metrics::tasks_stuck_detected()
                        .labels(prom::Labels{{"pool_id", std::to_string(ev.pool_id)}})
                        .inc();
                    if (stuck_listeners_copy.empty()) {
                        if (auto& lg = log::task()) {
                            lg->warn("task {} on pool {} stuck for {}ms",
                                     ev.task.id,
                                     ev.pool_id,
                                     std::chrono::duration_cast<std::chrono::milliseconds>(
                                         ev.running_for)
                                         .count());
                        }
                    } else {
                        for (const auto& cb : stuck_listeners_copy) {
                            try {
                                cb(ev);
                            } catch (const std::exception& e) {
                                if (auto& lg = log::manager()) {
                                    lg->warn("stuck-task listener threw: {}", e.what());
                                }
                            } catch (...) {
                                // ignore
                            }
                        }
                    }
                }
            }
 
            // Reconcile the dedup set against tasks that are still stuck so it
            // cannot grow without bound over the process lifetime: drop any id
            // that no longer appears in this tick's stuck set. A task's
            // run-duration only increases while it occupies a worker, so a
            // still-stuck task stays in every tick's set and never re-fires;
            // ids are only forgotten once their task leaves the running map.
            {
                std::unordered_set<std::uint64_t> still_stuck;
                still_stuck.reserve(all_stuck.size());
                for (const auto& ev : all_stuck) {
                    still_stuck.insert(ev.task.id);
                }
                std::erase_if(known_stuck,
                              [&](const std::uint64_t id) { return !still_stuck.contains(id); });
                known_stuck_size_.store(known_stuck.size(), std::memory_order_relaxed);
                metrics::tasks_stuck_current().set(known_stuck.size());
            }
 
            // 4) Periodic summary publish.
            if (now - last_summary_at_ >= opts_.summary_interval) {
                last_summary_at_ = now;
                auto summary = build_summary();
 
                // Publish per-pool gauges from the freshly-built summary.
                for (const auto& s : summary.per_pool_stats) {
                    const auto pool_labels = prom::Labels{{"pool", s.name}};
                    metrics::pool_queue_depth().labels(pool_labels).set(s.queued);
                    metrics::pool_workers().labels(pool_labels).set(s.workers);
                    metrics::pool_workers_active().labels(pool_labels).set(s.active);
                    metrics::pool_workers_idle().labels(pool_labels).set(s.idle);
                }
                metrics::manager_pools_live().set(summary.total_pools);
                metrics::manager_summaries_published().inc();
 
                summary.stuck_tasks.reserve(all_stuck.size());
                for (const auto& ev : all_stuck) {
                    summary.stuck_tasks.push_back(ev.task);
                }
                std::vector<SummaryListener> summary_listeners_copy;
                {
                    std::scoped_lock guard(subs_mtx_);
                    summary_listeners_copy.reserve(summary_subs_.size());
                    for (const auto& cb : summary_subs_ | std::views::values) {
                        summary_listeners_copy.push_back(cb);
                    }
                }
                for (const auto& cb : summary_listeners_copy) {
                    try {
                        cb(summary);
                    } catch (const std::exception& e) {
                        if (auto& lg = log::manager()) {
                            lg->warn("summary listener threw: {}", e.what());
                        }
                    } catch (...) {
                        // ignore
                    }
                }
            }
 
            if (auto& lg = log::manager(); lg && lg->should_log(spdlog::level::debug)) {
                std::size_t threads_n = 0;
                std::size_t pools_n = 0;
                {
                    std::shared_lock g(threads_mtx_);
                    threads_n = threads_.size();
                }
                {
                    std::scoped_lock g(pools_mtx_);
                    pools_n = pools_.size();
                }
                lg->debug("housekeeper tick: {} threads, {} pools", threads_n, pools_n);
            }
        }
    }
 
    Options opts_;
 
    mutable std::shared_mutex threads_mtx_;
    std::unordered_map<std::uint64_t, std::weak_ptr<ManagedThread::ControlBlock>> threads_;
    std::atomic<std::uint64_t> total_registered_{0};
 
    mutable std::mutex pools_mtx_;
    std::vector<ThreadPool*> pools_;
 
    mutable std::mutex subs_mtx_;
    std::vector<std::pair<std::uint64_t, SummaryListener>> summary_subs_;
    std::vector<std::pair<std::uint64_t, StuckTaskListener>> stuck_subs_;
    std::atomic<std::uint64_t> next_sub_id_{1};
 
    std::once_flag default_pool_once_;
    std::unique_ptr<ThreadPool> default_pool_;
 
    std::optional<ManagedThread> housekeeper_;
    std::atomic<bool> housekeeper_started_{false};
    std::chrono::steady_clock::time_point last_summary_at_;
    std::atomic<std::size_t> known_stuck_size_{0};
    mutable std::mutex hk_mtx_;
    std::condition_variable hk_cv_;
};
 
// ---------------------------------------------------------------------------
// Deferred manager-coupled definitions on ManagedThread / ThreadPool that
// could not be inlined earlier because ThreadManager was incomplete.
// ---------------------------------------------------------------------------
 
inline void ManagedThread::register_with_manager() {
    auto* m = (opts_manager_ != nullptr) ? opts_manager_ : &ThreadManager::instance();
    m->register_thread(cb_);
}
 
inline void ManagedThread::unregister_from_manager() noexcept {
    if (!cb_) {
        return;
    }
    auto* m = (opts_manager_ != nullptr) ? opts_manager_ : &ThreadManager::instance();
    m->unregister_thread(cb_->id);
}
 
inline void ThreadPool::register_with_manager() {
    manager_ = (opts_.manager != nullptr) ? opts_.manager : &ThreadManager::instance();
    manager_->register_pool(*this);
}
 
inline void ThreadPool::unregister_from_manager() noexcept {
    if (manager_ != nullptr) {
        manager_->unregister_pool(pool_id_);
    }
}
 
}  // namespace threadman