vdm/core/include/vdm/segment/budget.hpp

// vdm/segment/budget.hpp — the global segment allocator (ADR 0011).
//
// One instance per engine. It owns exactly one ceiling — `maxActiveSegments`, in segment
// units — and the min-1-before-seconds fairness rule (ADR 0011 §3). DAEMON's scheduler
// counts tasks and never touches this except through the read-outs and setters below;
// CORE's download_task (stage 8) drives the task-facing half.
//
// Fairness (two-pass, recomputed on every edge): a guarantee pass gives every task that
// wants a slot and holds zero exactly one, in DAEMON's priority order; then a growth pass
// round-robins the remainder up to each task's effective cap. A task's target can drop
// below what it holds (a lower-priority task shedding for a higher-priority arrival, or a
// live `set_max_active_segments` cut) — the task then *yields*: it releases a slot at its
// next segment boundary, never mid-segment (ADR 0011 A1). A finishing worker whose target
// still covers it *steals* instead (slot-neutral). That steal-vs-yield choice lives in
// stage 8, driven by comparing this budget's target to the task's live worker count.
//
// This header compiles standalone.

#ifndef VDM_SEGMENT_BUDGET_HPP
#define VDM_SEGMENT_BUDGET_HPP

#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <functional>
#include <mutex>
#include <optional>
#include <span>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

#include "vdm/ids.hpp"

namespace vdm::segment {

class SegmentBudget {
  public:
    struct EngineBudget {
        std::uint32_t total = 0;          // == maxActiveSegments
        std::uint32_t active = 0;         // slots held across all tasks
        std::uint32_t tasks_starved = 0;  // running tasks holding zero slots (ADR 0011 §3.6)
        bool operator==(const EngineBudget &) const = default;
    };

    struct Options {
        std::uint32_t max_active_segments = 32;
        std::chrono::milliseconds notify_period{250};  // <=4 Hz, per event.task.progress
    };
    SegmentBudget();  // default Options
    explicit SegmentBudget(Options opts);
    ~SegmentBudget();
    SegmentBudget(const SegmentBudget &) = delete;
    SegmentBudget &operator=(const SegmentBudget &) = delete;

    // ---- task-facing (download_task, stage 8) -------------------------------------------
    struct TaskParams {
        std::string host;                // key for the per-host segment cap
        std::uint32_t per_task_cap = 1;  // min(spec.segments ?? maxSegmentsPerDownload, 32)
        bool resumable = false;          // false => effective cap forced to 1
    };
    // Called with the new absolute slot target for the task. Runs on a budget thread (or
    // the caller's, for the edge case) — must not block and must not re-enter the budget
    // beyond confirm_slot()/release_slot(). The task starts or yields workers to match.
    using SlotTargetFn = std::function<void(std::uint32_t target)>;

    void register_task(TaskId id, const TaskParams &params, SlotTargetFn on_target);
    void deregister_task(TaskId id);  // pause / complete / cancel

    // How many slots the task could use right now (0 .. per_task_cap): incomplete
    // segments it has range for. 0 while retry_wait / assembling / verifying / paused —
    // which is exactly why those states are never counted as starvation.
    void set_want(TaskId id, std::uint32_t want);

    // A worker actually started on a granted slot. false => the target was cut in the
    // race and the worker must not start.
    [[nodiscard]] bool confirm_slot(TaskId id);
    // A held slot is free: segment complete with no steal, failed, paused, or yielded.
    void release_slot(TaskId id);

    // ---- DAEMON-facing (sched/) -------------------------------------------------------
    void set_max_active_segments(std::uint32_t n);                   // drain-not-kill (ADR 0011 §2)
    void set_host_segment_cap(std::string host, std::uint32_t cap);  // 0 clears
    void set_task_order(std::span<const TaskId> priority_order);     // pushed on change

    [[nodiscard]] EngineBudget budget() const;
    [[nodiscard]] std::uint32_t segments_active(TaskId id) const;
    [[nodiscard]] std::vector<TaskId> starved_tasks() const;
    [[nodiscard]] std::optional<SteadyTime> starved_since(TaskId id) const;
    void on_budget_changed(std::function<void(EngineBudget)> cb);

  private:
    struct Task {
        std::string host;
        std::uint32_t per_task_cap = 1;
        bool resumable = false;
        std::uint32_t want = 0;
        std::uint32_t held = 0;
        std::uint32_t target = 0;  // last published
        SlotTargetFn on_target;
        std::optional<SteadyTime> starved_since;
    };

    // A unit of deferred work: callbacks are copied out here so the public entry points
    // can invoke them AFTER dropping mu_.
    struct Plan {
        std::vector<std::pair<SlotTargetFn, std::uint32_t>> targets;
        std::optional<std::pair<std::function<void(EngineBudget)>, EngineBudget>> notify_now;
    };

    Plan reallocate_locked();
    static void run(Plan &p);
    [[nodiscard]] std::uint32_t effective_cap_locked(const Task &t) const;
    [[nodiscard]] EngineBudget snapshot_locked() const;
    void notifier_loop(std::stop_token st);

    mutable std::mutex mu_;
    std::unordered_map<TaskId, Task> tasks_;
    std::vector<TaskId> order_;
    std::unordered_map<std::string, std::uint32_t> host_caps_;
    std::uint32_t max_active_;
    std::uint32_t active_ = 0;

    std::function<void(EngineBudget)> on_changed_;
    EngineBudget last_notified_;
    std::uint32_t last_starved_ = 0;
    bool dirty_ = false;

    std::chrono::milliseconds notify_period_;
    std::condition_variable notify_cv_;
    std::jthread notifier_;
};

}  // namespace vdm::segment

#endif  // VDM_SEGMENT_BUDGET_HPP