vdm/core/src/engine.cpp

// vdm/engine.cpp

#include "vdm/engine.hpp"

#include <atomic>
#include <condition_variable>
#include <cstdint>
#include <map>
#include <mutex>
#include <thread>
#include <unordered_map>

#include "task/download_task.hpp"

namespace vdm {

struct Engine::Impl : task::TaskHost {
    explicit Impl(Config c)
        : cfg_(c),
          http_(net::HttpClient::Options{.workers = c.http_workers}),
          prober_(c.probe_pool_size ? c.probe_pool_size : 4),
          budget_(segment::SegmentBudget::Options{.max_active_segments = c.max_active_segments}) {
        timer_ = std::jthread([this](std::stop_token st) { timer_loop(st); });
    }

    ~Impl() override {
        // Quiesce every task first so no callback fires during or after teardown: mark it
        // retired and cancel its transfers. Then stop the timer thread (a fn already
        // running holds a shared_ptr and finishes, but seg_finished early-returns on
        // `retired`). Then drop the registry; http_/prober_/budget_ destruct after.
        {
            std::lock_guard lk(reg_mu_);
            for (auto &[id, t] : tasks_)
                task::quiesce_task(t);
        }
        timer_.request_stop();
        timer_cv_.notify_all();
        if (timer_.joinable())
            timer_.join();
        {
            std::lock_guard lk(reg_mu_);
            tasks_.clear();
        }
    }

    // --- TaskHost -------------------------------------------------------------------
    net::HttpClient &http() override { return http_; }
    segment::SegmentBudget &budget() override { return budget_; }
    rate::RateLimiter &limiter() override { return limiter_; }
    const Config &config() override { return cfg_; }

    task::TimerId schedule(SteadyTime at, std::function<void()> fn) override {
        task::TimerId id;
        {
            std::lock_guard lk(timer_mu_);
            id = ++next_timer_;
            timers_.emplace(at, Entry{id, std::move(fn)});
        }
        timer_cv_.notify_all();
        return id;
    }
    void cancel_timer(task::TimerId id) override {
        std::lock_guard lk(timer_mu_);
        for (auto it = timers_.begin(); it != timers_.end(); ++it)
            if (it->second.id == id) {
                timers_.erase(it);
                return;
            }
    }
    void probe(net::ProbeRequest req, std::function<void(Result<net::ProbeResult>)> done) override {
        prober_.probe(std::move(req), std::move(done));
    }
    void task_retired(TaskId id) override {
        std::lock_guard lk(reg_mu_);
        tasks_.erase(id);
    }

    // --- engine surface ----------------------------------------------------------------
    task::DownloadHandle start(task::DownloadSpec spec, task::DownloadCallbacks cbs) {
        TaskId id{++next_id_};
        auto st = task::create_task(*this, id, std::move(spec), std::move(cbs));
        {
            std::lock_guard lk(reg_mu_);
            tasks_[id] = st;
        }
        return task::DownloadHandle(std::move(st));
    }

    Config cfg_;
    net::HttpClient http_;
    net::Prober prober_;
    segment::SegmentBudget budget_;
    rate::RateLimiter limiter_;

    std::atomic<std::uint64_t> next_id_{0};

    std::mutex reg_mu_;
    std::unordered_map<TaskId, std::shared_ptr<task::DownloadTaskState>> tasks_;

    struct Entry {
        task::TimerId id;
        std::function<void()> fn;
    };
    std::mutex timer_mu_;
    std::condition_variable timer_cv_;
    std::multimap<SteadyTime, Entry> timers_;
    std::atomic<std::uint64_t> next_timer_{0};
    std::jthread timer_;

    void timer_loop(std::stop_token st) {
        std::unique_lock lk(timer_mu_);
        while (!st.stop_requested()) {
            if (timers_.empty()) {
                timer_cv_.wait_for(lk, std::chrono::seconds(1));
                continue;
            }
            auto next_at = timers_.begin()->first;
            if (timer_cv_.wait_until(lk, next_at, [&] {
                    return st.stop_requested() ||
                           (!timers_.empty() && timers_.begin()->first < next_at);
                })) {
                continue;  // stop, or an earlier timer landed — re-evaluate
            }
            // fire everything due
            std::vector<std::function<void()>> due;
            auto now = std::chrono::steady_clock::now();
            for (auto it = timers_.begin(); it != timers_.end() && it->first <= now;)
                due.push_back(std::move(it->second.fn)), it = timers_.erase(it);
            lk.unlock();
            for (auto &fn : due)
                fn();
            lk.lock();
        }
    }
};

// --- Engine ---------------------------------------------------------------------------

Engine::Engine() : Engine(Config{}) {}
Engine::Engine(Config cfg) : impl_(std::make_unique<Impl>(cfg)) {}
Engine::~Engine() = default;

task::DownloadHandle Engine::start(task::DownloadSpec spec, task::DownloadCallbacks cbs) {
    return impl_->start(std::move(spec), std::move(cbs));
}

segment::SegmentBudget &Engine::segment_budget() noexcept {
    return impl_->budget_;
}
rate::RateLimiter &Engine::rate_limiter() noexcept {
    return impl_->limiter_;
}

void Engine::set_default_segments(std::uint32_t n) {
    impl_->cfg_.default_segments = n ? n : 1;
}
void Engine::set_default_buffer_bytes(std::uint64_t b) {
    impl_->cfg_.default_buffer_bytes = b;
}
void Engine::set_max_total_buffer_bytes(std::uint64_t b) {
    impl_->cfg_.max_total_buffer_bytes = b;
}
void Engine::set_probe_pool_size(std::uint32_t) { /* prober pool is fixed at construction in M1 */ }

void Engine::probe(net::ProbeRequest req, std::function<void(Result<net::ProbeResult>)> done) {
    impl_->prober_.probe(std::move(req), std::move(done));
}

}  // namespace vdm