dag_scheduler.h

Go to the documentation of this file.

// BSD 3-Clause License
// Copyright (c) 2024, 🍀☀🌕🌥 🌊
// See the LICENSE file in the project root for full license information.
 
#pragma once
 
#include "dag_job.h"
#include "dag_job_builder.h"
#include "dag_config.h"
 
#include <kcenon/thread/core/thread_pool.h>
#include <kcenon/thread/core/error_handling.h>
 
#include <atomic>
#include <condition_variable>
#include <future>
#include <memory>
#include <mutex>
#include <optional>
#include <shared_mutex>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
 
namespace kcenon::thread
{
    class dag_scheduler
    {
    public:
        explicit dag_scheduler(std::shared_ptr<thread_pool> pool,
                               dag_config config = {});
 
        ~dag_scheduler();
 
        // Non-copyable
        dag_scheduler(const dag_scheduler&) = delete;
        auto operator=(const dag_scheduler&) -> dag_scheduler& = delete;
 
        // Movable
        dag_scheduler(dag_scheduler&&) noexcept;
        auto operator=(dag_scheduler&&) noexcept -> dag_scheduler&;
 
        // ============================================
        // Job Management
        // ============================================
 
        [[nodiscard]] auto add_job(std::unique_ptr<dag_job> j) -> job_id;
 
        [[nodiscard]] auto add_job(dag_job_builder&& builder) -> job_id;
 
        [[nodiscard]] auto add_dependency(job_id dependent, job_id dependency) -> common::VoidResult;
 
        [[nodiscard]] auto remove_job(job_id id) -> common::VoidResult;
 
        // ============================================
        // Execution Control
        // ============================================
 
        [[nodiscard]] auto execute_all() -> std::future<common::VoidResult>;
 
        [[nodiscard]] auto execute(job_id target) -> std::future<common::VoidResult>;
 
        auto cancel_all() -> void;
 
        [[nodiscard]] auto wait() -> common::VoidResult;
 
        [[nodiscard]] auto reset() -> common::VoidResult;
 
        // ============================================
        // Query
        // ============================================
 
        [[nodiscard]] auto get_job_info(job_id id) const -> std::optional<dag_job_info>;
 
        [[nodiscard]] auto get_all_jobs() const -> std::vector<dag_job_info>;
 
        [[nodiscard]] auto get_jobs_in_state(dag_job_state state) const -> std::vector<dag_job_info>;
 
        [[nodiscard]] auto get_ready_jobs() const -> std::vector<job_id>;
 
        [[nodiscard]] auto has_cycles() const -> bool;
 
        [[nodiscard]] auto get_execution_order() const -> std::vector<job_id>;
 
        template<typename T>
        [[nodiscard]] auto get_result(job_id id) const -> const T&
        {
            std::shared_lock lock(mutex_);
            auto it = jobs_.find(id);
            if (it == jobs_.end())
            {
                throw std::runtime_error("Job not found: " + std::to_string(id));
            }
            if (it->second->get_state() != dag_job_state::completed)
            {
                throw std::runtime_error("Job not completed: " + std::to_string(id));
            }
            return it->second->get_result<T>();
        }
 
        // ============================================
        // Visualization
        // ============================================
 
        [[nodiscard]] auto to_dot() const -> std::string;
 
        [[nodiscard]] auto to_json() const -> std::string;
 
        // ============================================
        // Statistics
        // ============================================
 
        [[nodiscard]] auto get_stats() const -> dag_stats;
 
        [[nodiscard]] auto get_config() const -> const dag_config& { return config_; }
 
    private:
        std::shared_ptr<thread_pool> pool_;
 
        dag_config config_;
 
        std::unordered_map<job_id, std::unique_ptr<dag_job>> jobs_;
 
        std::unordered_map<job_id, std::vector<job_id>> dependencies_;
 
        std::unordered_map<job_id, std::vector<job_id>> dependents_;
 
        mutable std::shared_mutex mutex_;
 
        std::condition_variable_any completion_cv_;
 
        std::atomic<bool> executing_{false};
 
        std::atomic<bool> cancelled_{false};
 
        std::atomic<std::size_t> running_count_{0};
 
        std::atomic<std::size_t> active_callbacks_{0};
 
        std::chrono::steady_clock::time_point execution_start_time_;
 
        std::optional<common::error_info> first_error_;
 
        std::unordered_map<job_id, std::size_t> retry_counts_;
 
        // ============================================
        // Internal Methods
        // ============================================
 
        auto on_job_completed(job_id id) -> void;
 
        auto on_job_failed(job_id id, const std::string& error) -> void;
 
        auto schedule_ready_jobs() -> void;
 
        auto execute_job(job_id id) -> void;
 
        [[nodiscard]] auto topological_sort() const -> std::vector<job_id>;
 
        [[nodiscard]] auto detect_cycle() const -> bool;
 
        [[nodiscard]] auto are_dependencies_satisfied(job_id id) const -> bool;
 
        auto skip_dependents(job_id failed_id) -> void;
 
        auto cancel_dependents(job_id failed_id) -> void;
 
        [[nodiscard]] static auto get_state_color(dag_job_state state) -> std::string;
 
        [[nodiscard]] auto is_execution_complete() const -> bool;
    };
 
} // namespace kcenon::thread