diff --git a/core/CMakeLists.txt b/core/CMakeLists.txt index 1140360..fca1a84 100644 --- a/core/CMakeLists.txt +++ b/core/CMakeLists.txt @@ -20,6 +20,8 @@ add_library(veloxcore STATIC src/io/sparse_file.cpp src/io/write_buffer.cpp src/meta/veloxpart.cpp + src/segment/segmenter.cpp + src/segment/budget.cpp ) add_library(velox::core ALIAS veloxcore) diff --git a/core/include/vdm/ids.hpp b/core/include/vdm/ids.hpp new file mode 100644 index 0000000..914b4a4 --- /dev/null +++ b/core/include/vdm/ids.hpp @@ -0,0 +1,38 @@ +// vdm/ids.hpp — opaque engine-internal identifiers. +// +// TaskId is CORE's handle for a download. DAEMON owns the wire UUID and keeps a +// UUID <-> TaskId map; CORE never sees the UUID (layering rule — no wire types in the +// engine). Assigned by CORE when DAEMON registers a task. +// +// This header compiles standalone. + +#ifndef VDM_IDS_HPP +#define VDM_IDS_HPP + +#include +#include +#include +#include +#include + +namespace vdm { + +struct TaskId { + std::uint64_t value = 0; + + [[nodiscard]] constexpr bool valid() const noexcept { return value != 0; } + friend constexpr auto operator<=>(const TaskId &, const TaskId &) = default; +}; + +using SteadyTime = std::chrono::steady_clock::time_point; + +} // namespace vdm + +template <> +struct std::hash { + std::size_t operator()(vdm::TaskId id) const noexcept { + return std::hash{}(id.value); + } +}; + +#endif // VDM_IDS_HPP diff --git a/core/include/vdm/segment/budget.hpp b/core/include/vdm/segment/budget.hpp new file mode 100644 index 0000000..cf758a7 --- /dev/null +++ b/core/include/vdm/segment/budget.hpp @@ -0,0 +1,137 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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 register_task(TaskId id, const TaskParams ¶ms, 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 priority_order); // pushed on change + + [[nodiscard]] EngineBudget budget() const; + [[nodiscard]] std::uint32_t segments_active(TaskId id) const; + [[nodiscard]] std::vector starved_tasks() const; + [[nodiscard]] std::optional starved_since(TaskId id) const; + void on_budget_changed(std::function 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 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> targets; + std::optional, 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 tasks_; + std::vector order_; + std::unordered_map host_caps_; + std::uint32_t max_active_; + std::uint32_t active_ = 0; + + std::function 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 diff --git a/core/include/vdm/segment/segmenter.hpp b/core/include/vdm/segment/segmenter.hpp new file mode 100644 index 0000000..1957ec1 --- /dev/null +++ b/core/include/vdm/segment/segmenter.hpp @@ -0,0 +1,174 @@ +// vdm/segment/segmenter.hpp — per-download range management and dynamic segment stealing. +// +// docs/04 §3. Owns the split of [0, total_size) for one download: the initial layout, a +// split when a slot is granted (`assign_slot`), a *steal* when a worker finishes and may +// keep its slot (`on_complete` — take the second half of the largest remaining range), +// and a re-split of an orphaned range when a segment fails 3x on the same host with a +// mirror available (`on_failed` -> requeue). +// +// Thread model: one mutex — the segmenter's — *is* "the task lock" (docs/04 §3: the steal +// is "atomic under the task lock"). Every method takes it. `advance()` and the per-worker +// accessors are called once per buffer flush (a few Hz per segment), not from the curl +// write callback, so a lock there is free; the no-lock/no-alloc rule is about that +// callback and its ring buffer. Segment fields stay std::atomic so the store type is +// trivially relocatable and reads never tear. The store is a std::deque so a steal's +// push_back never moves an existing record. +// +// This header compiles standalone. + +#ifndef VDM_SEGMENT_SEGMENTER_HPP +#define VDM_SEGMENT_SEGMENTER_HPP + +#include +#include +#include +#include +#include +#include + +namespace vdm::segment { + +inline constexpr std::uint64_t kDefaultMinSegmentBytes = 1u << 20; // 1 MiB (docs/04 §3) +inline constexpr std::uint32_t kDefaultSegments = 8; +inline constexpr std::uint32_t kMaxSegments = 32; + +enum class SegState : std::uint8_t { + idle, // range assigned, no worker connected yet + connecting, + downloading, + stalled, // low-speed; still holds its slot + complete, + failed, // gave up (orphaned; range requeued or lost to a steal) +}; + +// A flat, copyable view of one segment record. Ranges are absolute byte offsets, +// **inclusive** on both ends (contract Segment.endByte / ADR 0010). +struct SegmentView { + std::uint32_t index = 0; + std::uint64_t start = 0; + std::uint64_t end = 0; + std::uint64_t completed = 0; + SegState state = SegState::idle; + std::uint32_t consecutive_failures = 0; + + [[nodiscard]] std::uint64_t length() const noexcept { + return end >= start ? end - start + 1 : 0; + } + [[nodiscard]] std::uint64_t remaining() const noexcept { + return length() - (completed < length() ? completed : length()); + } + bool operator==(const SegmentView &) const = default; +}; + +// One resumed range, as read back from .veloxpart.meta (kept independent of meta/ so this +// header stands alone). +struct ResumedRange { + std::uint64_t start = 0; + std::uint64_t end = 0; + std::uint64_t completed = 0; +}; + +enum class FailAction { + retry, // same range, backoff (owned by the task/state machine) + requeue, // 3x connection failure + a mirror exists: orphan the remaining range and + // re-split it; the segment's slot is released +}; + +class Segmenter { + public: + // total_size 0 => unknown (chunked): forces a single segment. resumable == false also + // forces a single segment (docs/04 §3: "Non-resumable servers -> exactly 1 segment"). + Segmenter(std::uint64_t total_size, std::uint32_t requested_segments, bool resumable, + std::uint64_t min_segment_bytes = kDefaultMinSegmentBytes); + + // Resume: rebuild from a persisted segment table. Ranges must tile [0, total_size) + // with no gap or overlap; a malformed table falls back to a single segment. + Segmenter(std::uint64_t total_size, std::uint32_t requested_segments, + const std::vector &resumed, bool resumable, + std::uint64_t min_segment_bytes = kDefaultMinSegmentBytes); + + Segmenter(const Segmenter &) = delete; + Segmenter &operator=(const Segmenter &) = delete; + + // The count this download would use with an unlimited budget. 1 when non-resumable / + // unknown size; otherwise min(requested, floor(total / min_segment_bytes), 32). + [[nodiscard]] std::uint32_t target_segment_count() const noexcept { return target_count_; } + [[nodiscard]] std::uint64_t total_size() const noexcept { return total_size_; } + + // Give a worker something to download. Called when the budget grants a slot. Prefers + // an orphaned range from a requeue; otherwise splits the largest remaining range in + // half and hands back the tail. Returns nullopt when the target count is already met + // or nothing splits to >= min_segment_bytes. `state` of the returned segment is + // `connecting`. + [[nodiscard]] std::optional assign_slot(); + + // A worker finished its range. If `may_steal` is false the slot is being yielded — + // returns nullopt and the caller releases the slot to the budget. If true and a + // remaining range splits to >= min_segment_bytes, steals its second half: returns a + // NEW segment index for the same worker to continue on (slot-neutral). Otherwise + // nullopt (nothing to steal -> release). + [[nodiscard]] std::optional on_complete(std::uint32_t idx, bool may_steal); + + // A worker's segment errored. `connection_error` distinguishes a transport failure + // (reset/timeout/refused) from an HTTP/content one. Returns requeue only on the 3rd + // consecutive connection error when `has_mirror`; on requeue the remaining range is + // orphaned for assign_slot() to re-split and the segment is marked failed. + FailAction on_failed(std::uint32_t idx, bool connection_error, bool has_mirror); + + // A successful (re)connection resets the consecutive-failure counter. + void note_connected(std::uint32_t idx); + + // Progress from the write path. Lock-free. `bytes` is the absolute completed count + // within the segment; clamped to the segment length. + void advance(std::uint32_t idx, std::uint64_t bytes) noexcept; + + // Per-segment fields for a worker. A worker reads `segment_end` before every write so + // a concurrent steal that shrank its range stops it cleanly. `segment_start` never + // changes for a given index. + [[nodiscard]] std::uint64_t segment_start(std::uint32_t idx) const noexcept; + [[nodiscard]] std::uint64_t segment_end(std::uint32_t idx) const noexcept; + [[nodiscard]] std::uint64_t segment_completed(std::uint32_t idx) const noexcept; + [[nodiscard]] SegState segment_state(std::uint32_t idx) const noexcept; + void set_segment_state(std::uint32_t idx, SegState s) noexcept; + + // Sum of bytes done across every segment (active + already complete). Locks. + [[nodiscard]] std::uint64_t downloaded() const; + [[nodiscard]] bool all_complete() const; + + // Every segment record, for TaskDetail projection and the .veloxpart.meta writer. + [[nodiscard]] std::vector snapshot() const; + + private: + struct Seg { + std::uint32_t index; + std::uint64_t start; + std::atomic end; + std::atomic completed{0}; + std::atomic state{SegState::idle}; + std::uint32_t consecutive_failures = 0; + bool assigned = false; // a worker holds this segment right now + + Seg(std::uint32_t i, std::uint64_t s, std::uint64_t e) : index(i), start(s), end(e) {} + }; + + void compute_target(std::uint32_t requested); + std::uint32_t add_seg_locked(std::uint64_t start, std::uint64_t end, std::uint64_t completed, + SegState state, bool assigned); + std::uint32_t split_largest_remaining_locked(); // returns new index, or UINT32_MAX + [[nodiscard]] std::uint64_t remaining_of_locked(const Seg &s) const noexcept; + [[nodiscard]] std::uint32_t assigned_count_locked() const noexcept; + + mutable std::mutex mu_; + std::deque segs_; + std::vector orphans_; // requeued ranges awaiting re-split + + std::uint64_t total_size_; + std::uint64_t min_seg_; + bool resumable_; + std::uint32_t target_count_ = 1; + std::uint32_t next_index_ = 0; +}; + +} // namespace vdm::segment + +#endif // VDM_SEGMENT_SEGMENTER_HPP diff --git a/core/src/segment/budget.cpp b/core/src/segment/budget.cpp new file mode 100644 index 0000000..10a7f61 --- /dev/null +++ b/core/src/segment/budget.cpp @@ -0,0 +1,299 @@ +// vdm/segment/budget.cpp + +#include "vdm/segment/budget.hpp" + +#include + +namespace vdm::segment { + +SegmentBudget::SegmentBudget() : SegmentBudget(Options{}) {} + +SegmentBudget::SegmentBudget(Options opts) + : max_active_(opts.max_active_segments ? opts.max_active_segments : 1), + notify_period_(opts.notify_period) { + notifier_ = std::jthread([this](std::stop_token st) { notifier_loop(st); }); +} + +SegmentBudget::~SegmentBudget() { + notifier_.request_stop(); + notify_cv_.notify_all(); +} + +// --- allocation ----------------------------------------------------------------------- + +std::uint32_t SegmentBudget::effective_cap_locked(const Task &t) const { + std::uint32_t base = t.resumable ? t.per_task_cap : 1; + base = std::clamp(base, 1, 32); + if (auto it = host_caps_.find(t.host); it != host_caps_.end() && it->second > 0) + base = std::min(base, it->second); + return base; +} + +SegmentBudget::EngineBudget SegmentBudget::snapshot_locked() const { + std::uint32_t starved = 0; + for (const auto &[id, t] : tasks_) + if (t.want >= 1 && t.held == 0) + ++starved; + return EngineBudget{max_active_, active_, starved}; +} + +// The two-pass fairness allocation. Recomputes every task's target from scratch (so a +// live cap cut naturally produces target < held -> yield), diffs against the last +// published target, and collects the callbacks to fire once mu_ is released. +SegmentBudget::Plan SegmentBudget::reallocate_locked() { + // Priority order: DAEMON's list first, then any registered task not in it (defensive; + // "a running task absent from the list sorts last"). + std::vector order; + order.reserve(tasks_.size()); + for (TaskId id : order_) + if (tasks_.count(id)) + order.push_back(id); + for (const auto &[id, _] : tasks_) + if (std::find(order.begin(), order.end(), id) == order.end()) + order.push_back(id); + + std::unordered_map target; + target.reserve(order.size()); + std::uint32_t pool = max_active_; + + auto capped_want = [&](TaskId id) { + const Task &t = tasks_.at(id); + return std::min(t.want, effective_cap_locked(t)); + }; + + // Guarantee pass: one slot each, in priority order, to anyone who wants one. + for (TaskId id : order) { + if (pool == 0) + break; + if (capped_want(id) >= 1) { + target[id] = 1; + --pool; + } + } + // Growth pass: round-robin the remainder, up to each task's effective cap. + while (pool > 0) { + bool granted = false; + for (TaskId id : order) { + if (pool == 0) + break; + std::uint32_t &tv = target[id]; + if (tv < capped_want(id)) { + ++tv; + --pool; + granted = true; + } + } + if (!granted) + break; + } + + const SteadyTime now = std::chrono::steady_clock::now(); + Plan plan; + for (auto &[id, t] : tasks_) { + std::uint32_t nt = target.count(id) ? target[id] : 0; + if (nt != t.target) { + t.target = nt; + if (t.on_target) + plan.targets.emplace_back(t.on_target, nt); + } + // starvation timestamp bookkeeping + bool starved_now = t.want >= 1 && t.held == 0; + if (starved_now && !t.starved_since) + t.starved_since = now; + if (!starved_now) + t.starved_since.reset(); + } + + EngineBudget eb = snapshot_locked(); + bool starved_edge = (eb.tasks_starved == 0) != (last_starved_ == 0); + if (eb != last_notified_) + dirty_ = true; + last_starved_ = eb.tasks_starved; + + if (starved_edge && on_changed_) { + plan.notify_now = std::make_pair(on_changed_, eb); + last_notified_ = eb; + dirty_ = false; + } + if (dirty_) + notify_cv_.notify_one(); + return plan; +} + +void SegmentBudget::run(Plan &p) { + for (auto &[fn, n] : p.targets) + if (fn) + fn(n); + if (p.notify_now && p.notify_now->first) + p.notify_now->first(p.notify_now->second); +} + +// --- task-facing -------------------------------------------------------------------- + +void SegmentBudget::register_task(TaskId id, const TaskParams ¶ms, SlotTargetFn on_target) { + Plan plan; + { + std::lock_guard lk(mu_); + Task t; + t.host = params.host; + t.per_task_cap = params.per_task_cap ? params.per_task_cap : 1; + t.resumable = params.resumable; + t.on_target = std::move(on_target); + tasks_[id] = std::move(t); + plan = reallocate_locked(); + } + run(plan); +} + +void SegmentBudget::deregister_task(TaskId id) { + Plan plan; + { + std::lock_guard lk(mu_); + auto it = tasks_.find(id); + if (it == tasks_.end()) + return; + active_ -= it->second.held; + tasks_.erase(it); + plan = reallocate_locked(); + } + run(plan); +} + +void SegmentBudget::set_want(TaskId id, std::uint32_t want) { + Plan plan; + { + std::lock_guard lk(mu_); + auto it = tasks_.find(id); + if (it == tasks_.end()) + return; + if (it->second.want == want) + return; + it->second.want = want; + plan = reallocate_locked(); + } + run(plan); +} + +bool SegmentBudget::confirm_slot(TaskId id) { + std::lock_guard lk(mu_); + auto it = tasks_.find(id); + if (it == tasks_.end()) + return false; + Task &t = it->second; + if (t.held >= t.target) + return false; // target was cut in the race + ++t.held; + ++active_; + if (snapshot_locked() != last_notified_) { + dirty_ = true; + notify_cv_.notify_one(); + } + return true; +} + +void SegmentBudget::release_slot(TaskId id) { + Plan plan; + { + std::lock_guard lk(mu_); + auto it = tasks_.find(id); + if (it == tasks_.end() || it->second.held == 0) + return; + --it->second.held; + --active_; + plan = reallocate_locked(); + } + run(plan); +} + +// --- DAEMON-facing --------------------------------------------------------------------- + +void SegmentBudget::set_max_active_segments(std::uint32_t n) { + Plan plan; + { + std::lock_guard lk(mu_); + n = n ? n : 1; + if (n == max_active_) + return; + max_active_ = n; + plan = reallocate_locked(); + } + run(plan); +} + +void SegmentBudget::set_host_segment_cap(std::string host, std::uint32_t cap) { + Plan plan; + { + std::lock_guard lk(mu_); + if (cap == 0) + host_caps_.erase(host); + else + host_caps_[std::move(host)] = cap; + plan = reallocate_locked(); + } + run(plan); +} + +void SegmentBudget::set_task_order(std::span priority_order) { + Plan plan; + { + std::lock_guard lk(mu_); + order_.assign(priority_order.begin(), priority_order.end()); + plan = reallocate_locked(); + } + run(plan); +} + +SegmentBudget::EngineBudget SegmentBudget::budget() const { + std::lock_guard lk(mu_); + return snapshot_locked(); +} + +std::uint32_t SegmentBudget::segments_active(TaskId id) const { + std::lock_guard lk(mu_); + auto it = tasks_.find(id); + return it == tasks_.end() ? 0 : it->second.held; +} + +std::vector SegmentBudget::starved_tasks() const { + std::lock_guard lk(mu_); + std::vector out; + for (const auto &[id, t] : tasks_) + if (t.want >= 1 && t.held == 0) + out.push_back(id); + return out; +} + +std::optional SegmentBudget::starved_since(TaskId id) const { + std::lock_guard lk(mu_); + auto it = tasks_.find(id); + return it == tasks_.end() ? std::nullopt : it->second.starved_since; +} + +void SegmentBudget::on_budget_changed(std::function cb) { + std::lock_guard lk(mu_); + on_changed_ = std::move(cb); +} + +// --- notifier thread: coalesced <=4 Hz ----------------------------------------------- + +void SegmentBudget::notifier_loop(std::stop_token st) { + std::unique_lock lk(mu_); + while (!st.stop_requested()) { + notify_cv_.wait_for(lk, notify_period_, [&] { return dirty_ || st.stop_requested(); }); + if (st.stop_requested()) + break; + if (!dirty_) + continue; + EngineBudget eb = snapshot_locked(); + auto cb = on_changed_; + last_notified_ = eb; + last_starved_ = eb.tasks_starved; + dirty_ = false; + lk.unlock(); + if (cb) + cb(eb); + lk.lock(); + } +} + +} // namespace vdm::segment diff --git a/core/src/segment/segmenter.cpp b/core/src/segment/segmenter.cpp new file mode 100644 index 0000000..2587024 --- /dev/null +++ b/core/src/segment/segmenter.cpp @@ -0,0 +1,312 @@ +// vdm/segment/segmenter.cpp + +#include "vdm/segment/segmenter.hpp" + +#include +#include + +namespace vdm::segment { +namespace { +constexpr std::uint32_t kNoIndex = std::numeric_limits::max(); +constexpr std::uint64_t kU64Max = std::numeric_limits::max(); + +bool is_live(SegState s) noexcept { + return s == SegState::idle || s == SegState::connecting || s == SegState::downloading || + s == SegState::stalled; +} +} // namespace + +// Seg holds std::atomics, so it is neither copyable nor movable — every insertion is an +// emplace_back that constructs it in place, followed by stores. This helper centralises +// that. Caller holds mu_. +std::uint32_t Segmenter::add_seg_locked(std::uint64_t start, std::uint64_t end, + std::uint64_t completed, SegState state, bool assigned) { + segs_.emplace_back(next_index_++, start, end); + Seg &s = segs_.back(); + s.completed.store(completed); + s.state.store(state); + s.assigned = assigned; + return s.index; +} + +// --- construction --------------------------------------------------------------------- + +Segmenter::Segmenter(std::uint64_t total_size, std::uint32_t requested_segments, bool resumable, + std::uint64_t min_segment_bytes) + : total_size_(total_size), + min_seg_(min_segment_bytes ? min_segment_bytes : 1), + resumable_(resumable) { + compute_target(requested_segments); +} + +Segmenter::Segmenter(std::uint64_t total_size, std::uint32_t requested_segments, + const std::vector &resumed, bool resumable, + std::uint64_t min_segment_bytes) + : total_size_(total_size), + min_seg_(min_segment_bytes ? min_segment_bytes : 1), + resumable_(resumable) { + compute_target(requested_segments); + + // Validate the resumed table tiles [0, total_size) exactly. + bool ok = resumable_ && total_size_ > 0 && !resumed.empty(); + if (ok) { + std::vector sorted = resumed; + std::sort(sorted.begin(), sorted.end(), + [](const auto &a, const auto &b) { return a.start < b.start; }); + std::uint64_t cursor = 0; + for (const auto &r : sorted) { + if (r.start != cursor || r.end < r.start || r.completed > r.end - r.start + 1) { + ok = false; + break; + } + cursor = r.end + 1; + } + if (ok && cursor != total_size_) + ok = false; + if (ok) { + for (const auto &r : sorted) { + bool done = r.completed == r.end - r.start + 1; + add_seg_locked(r.start, r.end, r.completed, + done ? SegState::complete : SegState::idle, false); + } + std::uint32_t incomplete = 0; + for (const auto &s : segs_) + if (s.state.load() != SegState::complete) + ++incomplete; + target_count_ = std::clamp(std::max(incomplete, 1u), 1, kMaxSegments); + return; + } + } + // Fall back to a fresh single/target layout (segs created lazily by assign_slot()). + segs_.clear(); +} + +void Segmenter::compute_target(std::uint32_t requested) { + if (!resumable_ || total_size_ == 0) { + target_count_ = 1; + return; + } + std::uint64_t by_size = total_size_ / min_seg_; + if (by_size == 0) + by_size = 1; + std::uint64_t t = requested == 0 ? kDefaultSegments : requested; + t = std::min(t, by_size); + target_count_ = std::clamp(static_cast(t), 1, kMaxSegments); +} + +// --- helpers (mu_ held) -------------------------------------------------------------- + +std::uint64_t Segmenter::remaining_of_locked(const Seg &s) const noexcept { + std::uint64_t end = s.end.load(); + std::uint64_t done = s.start + s.completed.load(); + return done > end ? 0 : end - done + 1; +} + +std::uint32_t Segmenter::assigned_count_locked() const noexcept { + std::uint32_t n = 0; + for (const auto &s : segs_) + if (s.assigned) + ++n; + return n; +} + +// Split the largest remaining range; hand back its second half as a new segment. +std::uint32_t Segmenter::split_largest_remaining_locked() { + Seg *victim = nullptr; + std::uint64_t best = 0; + for (auto &s : segs_) { + if (!s.assigned || !is_live(s.state.load())) + continue; + std::uint64_t rem = remaining_of_locked(s); + if (rem > best) { + best = rem; + victim = &s; + } + } + if (!victim || best < 2 * min_seg_) + return kNoIndex; + + const std::uint64_t v_end = victim->end.load(); + const std::uint64_t half = best / 2; // >= min_seg_ since best >= 2*min + const std::uint64_t mid = v_end - half; // victim keeps [start, mid] + const std::uint64_t cur = victim->start + victim->completed.load(); + if (mid < cur || mid - cur + 1 < min_seg_) + return kNoIndex; // victim would be too small + + victim->end.store(mid); // the victim's worker reads end before each write and stops here + return add_seg_locked(mid + 1, v_end, 0, SegState::idle, false); +} + +// --- public: structural (take the lock) -------------------------------------------- + +std::optional Segmenter::assign_slot() { + std::lock_guard lk(mu_); + + if (!orphans_.empty()) { + ResumedRange o = orphans_.front(); + orphans_.erase(orphans_.begin()); + return add_seg_locked(o.start, o.end, o.completed, SegState::connecting, true); + } + + if (assigned_count_locked() >= target_count_) + return std::nullopt; + + if (segs_.empty()) { + std::uint64_t end = total_size_ > 0 ? total_size_ - 1 : kU64Max - 1; + return add_seg_locked(0, end, 0, SegState::connecting, true); + } + + // Some resumed segments may be unassigned idle ranges — hand one out before splitting. + for (auto &s : segs_) { + if (!s.assigned && s.state.load() == SegState::idle) { + s.assigned = true; + s.state.store(SegState::connecting); + return s.index; + } + } + + std::uint32_t idx = split_largest_remaining_locked(); + if (idx == kNoIndex) + return std::nullopt; + segs_[idx].assigned = true; + segs_[idx].state.store(SegState::connecting); + return idx; +} + +std::optional Segmenter::on_complete(std::uint32_t idx, bool may_steal) { + std::lock_guard lk(mu_); + if (idx >= segs_.size()) + return std::nullopt; + Seg &seg = segs_[idx]; + seg.completed.store(seg.end.load() - seg.start + 1); + seg.state.store(SegState::complete); + seg.assigned = false; + + if (!may_steal) + return std::nullopt; // yielding the slot + + if (!orphans_.empty()) { + ResumedRange o = orphans_.front(); + orphans_.erase(orphans_.begin()); + return add_seg_locked(o.start, o.end, o.completed, SegState::connecting, true); + } + + std::uint32_t new_idx = split_largest_remaining_locked(); + if (new_idx == kNoIndex) + return std::nullopt; // nothing to steal -> release the slot + segs_[new_idx].assigned = true; + segs_[new_idx].state.store(SegState::connecting); + return new_idx; +} + +FailAction Segmenter::on_failed(std::uint32_t idx, bool connection_error, bool has_mirror) { + std::lock_guard lk(mu_); + if (idx >= segs_.size()) + return FailAction::retry; + Seg &seg = segs_[idx]; + ++seg.consecutive_failures; + + if (connection_error && seg.consecutive_failures >= 3 && has_mirror) { + std::uint64_t cur = seg.start + seg.completed.load(); + std::uint64_t end = seg.end.load(); + if (cur <= end) + orphans_.push_back({cur, end, 0}); + seg.state.store(SegState::failed); + seg.assigned = false; + return FailAction::requeue; + } + return FailAction::retry; +} + +void Segmenter::note_connected(std::uint32_t idx) { + std::lock_guard lk(mu_); + if (idx < segs_.size()) + segs_[idx].consecutive_failures = 0; +} + +// --- public: per-worker accessors ---------------------------------------------- +// +// These take mu_. They are called from the write path once per buffer flush (a few per +// second per segment), not from the curl write callback — the no-lock/no-alloc rule is +// about that callback and its ring buffer, not about progress bookkeeping. The segment +// fields are still std::atomic so a reader that already holds a stable reference sees a +// torn-free value, and so the deque element type is safe to relocate-free. + +void Segmenter::advance(std::uint32_t idx, std::uint64_t bytes) noexcept { + std::lock_guard lk(mu_); + if (idx >= segs_.size()) + return; + Seg &s = segs_[idx]; + std::uint64_t len = s.end.load() - s.start + 1; + s.completed.store(bytes < len ? bytes : len); +} + +std::uint64_t Segmenter::segment_start(std::uint32_t idx) const noexcept { + std::lock_guard lk(mu_); + return idx < segs_.size() ? segs_[idx].start : 0; +} +std::uint64_t Segmenter::segment_end(std::uint32_t idx) const noexcept { + std::lock_guard lk(mu_); + return idx < segs_.size() ? segs_[idx].end.load() : 0; +} +std::uint64_t Segmenter::segment_completed(std::uint32_t idx) const noexcept { + std::lock_guard lk(mu_); + return idx < segs_.size() ? segs_[idx].completed.load() : 0; +} +SegState Segmenter::segment_state(std::uint32_t idx) const noexcept { + std::lock_guard lk(mu_); + return idx < segs_.size() ? segs_[idx].state.load() : SegState::failed; +} +void Segmenter::set_segment_state(std::uint32_t idx, SegState st) noexcept { + std::lock_guard lk(mu_); + if (idx < segs_.size()) + segs_[idx].state.store(st); +} + +// --- public: queries (take the lock) ---------------------------------------------- + +std::uint64_t Segmenter::downloaded() const { + std::lock_guard lk(mu_); + std::uint64_t sum = 0; + for (const auto &s : segs_) + sum += s.completed.load(); + return sum; +} + +bool Segmenter::all_complete() const { + std::lock_guard lk(mu_); + if (segs_.empty()) + return false; + if (total_size_ == 0) + return segs_.front().state.load() == SegState::complete; + if (!orphans_.empty()) + return false; + + std::vector> done; // [start, start+completed) + for (const auto &s : segs_) { + std::uint64_t c = s.completed.load(); + if (c > 0) + done.emplace_back(s.start, s.start + c); + } + std::sort(done.begin(), done.end()); + std::uint64_t cursor = 0; + for (auto [a, b] : done) { + if (a > cursor) + return false; // gap + if (b > cursor) + cursor = b; + } + return cursor >= total_size_; +} + +std::vector Segmenter::snapshot() const { + std::lock_guard lk(mu_); + std::vector out; + out.reserve(segs_.size()); + for (const auto &s : segs_) + out.push_back(SegmentView{s.index, s.start, s.end.load(), s.completed.load(), + s.state.load(), s.consecutive_failures}); + return out; +} + +} // namespace vdm::segment diff --git a/core/tests/CMakeLists.txt b/core/tests/CMakeLists.txt index 8ea7ed8..28db7d1 100644 --- a/core/tests/CMakeLists.txt +++ b/core/tests/CMakeLists.txt @@ -28,6 +28,8 @@ vdm_add_test(veloxcore_url_test net/url_test.cpp) vdm_add_test(veloxcore_sparse_file_test io/sparse_file_test.cpp) vdm_add_test(veloxcore_write_buffer_test io/write_buffer_test.cpp) vdm_add_test(veloxcore_veloxpart_test meta/veloxpart_test.cpp) +vdm_add_test(veloxcore_segmenter_test segment/segmenter_test.cpp) +vdm_add_test(veloxcore_budget_test segment/budget_test.cpp) set(_testserver ${CMAKE_SOURCE_DIR}/tools/testserver/testserver.py) foreach(net_it http_client probe) diff --git a/core/tests/segment/budget_test.cpp b/core/tests/segment/budget_test.cpp new file mode 100644 index 0000000..9d81428 --- /dev/null +++ b/core/tests/segment/budget_test.cpp @@ -0,0 +1,264 @@ +#include "vdm/segment/budget.hpp" + +#include +#include +#include +#include +#include + +#include "vtest.hpp" + +using namespace vdm; +using namespace vdm::segment; +using EB = SegmentBudget::EngineBudget; + +namespace { + +TaskId tid(std::uint64_t v) { + return TaskId{v}; +} + +// A test task that reacts to slot targets the way stage 8's download_task will: start +// workers up to the target, release them when the target drops. Purely bookkeeping. +struct FakeTask { + SegmentBudget *budget = nullptr; + TaskId id{}; + std::mutex mu; + std::uint32_t workers = 0; + std::uint32_t target = 0; + + FakeTask() = default; + FakeTask(SegmentBudget *b, TaskId i) : budget(b), id(i) {} + + void on_target(std::uint32_t t) { + std::lock_guard lk(mu); + target = t; + while (workers < target) { + if (!budget->confirm_slot(id)) + break; + ++workers; + } + // over target -> yield the excess immediately (a real task waits for a boundary) + while (workers > target) { + budget->release_slot(id); + --workers; + } + } + std::uint32_t held() { + std::lock_guard lk(mu); + return workers; + } +}; + +} // namespace + +VT_TEST(budget_single_task_grows_to_cap) { + SegmentBudget b({.max_active_segments = 32}); + FakeTask t{&b, tid(1)}; + b.register_task(tid(1), {.host = "h", .per_task_cap = 8, .resumable = true}, + [&](std::uint32_t n) { t.on_target(n); }); + b.set_want(tid(1), 8); + VT_CHECK_EQ(t.held(), 8u); + VT_CHECK_EQ(b.segments_active(tid(1)), 8u); + VT_CHECK_EQ(b.budget().active, 8u); + VT_CHECK_EQ(b.budget().tasks_starved, 0u); +} + +VT_TEST(budget_min_one_before_seconds) { + // Budget of 3, two tasks each wanting 8. min-1 first: each gets 1, then the higher- + // priority one grows to 2. + SegmentBudget b({.max_active_segments = 3}); + FakeTask a{&b, tid(1)}, c{&b, tid(2)}; + b.register_task(tid(1), {.host = "h1", .per_task_cap = 8, .resumable = true}, + [&](std::uint32_t n) { a.on_target(n); }); + b.register_task(tid(2), {.host = "h2", .per_task_cap = 8, .resumable = true}, + [&](std::uint32_t n) { c.on_target(n); }); + std::vector order = {tid(1), tid(2)}; + b.set_task_order(order); + b.set_want(tid(1), 8); + b.set_want(tid(2), 8); + + VT_CHECK(a.held() >= 1); // guarantee + VT_CHECK(c.held() >= 1); // guarantee — the load-bearing property + VT_CHECK_EQ(a.held() + c.held(), 3u); + VT_CHECK_EQ(a.held(), 2u); // higher priority took the growth slot +} + +VT_TEST(budget_new_high_priority_task_gets_min_one_via_yield) { + SegmentBudget b({.max_active_segments = 4}); + FakeTask a{&b, tid(1)}; + b.register_task(tid(1), {.host = "h", .per_task_cap = 8, .resumable = true}, + [&](std::uint32_t n) { a.on_target(n); }); + b.set_task_order(std::vector{tid(1)}); + b.set_want(tid(1), 8); + VT_CHECK_EQ(a.held(), 4u); // hogging the whole budget + + // a second, higher-priority task arrives + FakeTask c{&b, tid(2)}; + b.register_task(tid(2), {.host = "h2", .per_task_cap = 8, .resumable = true}, + [&](std::uint32_t n) { c.on_target(n); }); + b.set_task_order(std::vector{tid(2), tid(1)}); + b.set_want(tid(2), 8); + + // a yields so c gets at least its guaranteed slot; the surplus is shared round-robin. + VT_CHECK(c.held() >= 1); // min-1 — the load-bearing guarantee + VT_CHECK(a.held() >= 1); // a keeps its own min-1 + VT_CHECK(a.held() < 4); // a really did yield at least one + VT_CHECK_EQ(a.held() + c.held(), 4u); + VT_CHECK_EQ(b.budget().active, 4u); + VT_CHECK_EQ(b.budget().tasks_starved, 0u); +} + +VT_TEST(budget_host_cap_clamps_effective_target) { + SegmentBudget b({.max_active_segments = 32}); + FakeTask t{&b, tid(1)}; + b.set_host_segment_cap("slowcdn", 4); + b.register_task(tid(1), {.host = "slowcdn", .per_task_cap = 16, .resumable = true}, + [&](std::uint32_t n) { t.on_target(n); }); + b.set_want(tid(1), 16); + VT_CHECK_EQ(t.held(), 4u); // clamped by the host cap, not per_task_cap + + b.set_host_segment_cap("slowcdn", 0); // clear + VT_CHECK_EQ(t.held(), 16u); +} + +VT_TEST(budget_non_resumable_task_capped_at_one) { + SegmentBudget b({.max_active_segments = 32}); + FakeTask t{&b, tid(1)}; + b.register_task(tid(1), {.host = "h", .per_task_cap = 8, .resumable = false}, + [&](std::uint32_t n) { t.on_target(n); }); + b.set_want(tid(1), 8); + VT_CHECK_EQ(t.held(), 1u); +} + +VT_TEST(budget_live_lower_sheds_via_yield_lowest_priority_first) { + SegmentBudget b({.max_active_segments = 24}); + FakeTask a{&b, tid(1)}, c{&b, tid(2)}, d{&b, tid(3)}; + for (auto *ft : {&a, &c, &d}) + b.register_task(ft->id, {.host = "h", .per_task_cap = 8, .resumable = true}, + [ft](std::uint32_t n) { ft->on_target(n); }); + b.set_task_order(std::vector{tid(1), tid(2), tid(3)}); + for (auto id : {tid(1), tid(2), tid(3)}) + b.set_want(id, 8); + VT_CHECK_EQ(a.held() + c.held() + d.held(), 24u); // 8 + 8 + 8 + + b.set_max_active_segments(10); // live cut + VT_CHECK_EQ(a.held() + c.held() + d.held(), 10u); + VT_CHECK(a.held() >= c.held() && c.held() >= d.held()); // priority order preserved + VT_CHECK(a.held() >= 1 && c.held() >= 1 && d.held() >= 1); // min-1 still honoured +} + +VT_TEST(budget_live_lower_below_task_count_starves_the_tail) { + SegmentBudget b({.max_active_segments = 6}); + std::vector ts(4); + for (std::uint32_t i = 0; i < 4; ++i) { + ts[i].budget = &b; + ts[i].id = tid(i + 1); + } + for (auto &ft : ts) + b.register_task(ft.id, {.host = "h", .per_task_cap = 4, .resumable = true}, + [&ft](std::uint32_t n) { ft.on_target(n); }); + b.set_task_order(std::vector{tid(1), tid(2), tid(3), tid(4)}); + for (auto &ft : ts) + b.set_want(ft.id, 4); + VT_CHECK_EQ(b.budget().tasks_starved, 0u); + + b.set_max_active_segments(3); // below the running-task count + VT_CHECK_EQ(ts[0].held(), 1u); + VT_CHECK_EQ(ts[3].held(), 0u); // lowest priority shed to zero + VT_CHECK_EQ(b.budget().tasks_starved, 1u); + VT_REQUIRE(b.starved_tasks().size() == 1); + VT_CHECK_EQ(b.starved_tasks()[0], tid(4)); + VT_CHECK(b.starved_since(tid(4)).has_value()); + VT_CHECK(!b.starved_since(tid(1)).has_value()); +} + +VT_TEST(budget_deregister_frees_slots_to_starved) { + SegmentBudget b({.max_active_segments = 4}); + FakeTask a{&b, tid(1)}, c{&b, tid(2)}; + b.register_task(tid(1), {.host = "h", .per_task_cap = 8, .resumable = true}, + [&](std::uint32_t n) { a.on_target(n); }); + b.set_task_order(std::vector{tid(1)}); + b.set_want(tid(1), 8); + VT_CHECK_EQ(a.held(), 4u); + + b.register_task(tid(2), {.host = "h", .per_task_cap = 8, .resumable = true}, + [&](std::uint32_t n) { c.on_target(n); }); + b.set_task_order(std::vector{tid(1), tid(2)}); + b.set_want(tid(2), 8); + VT_CHECK(c.held() >= 1); // min-1 from a's yield + + b.deregister_task(tid(1)); + VT_CHECK_EQ(c.held(), 4u); // c grows into the whole freed budget + VT_CHECK_EQ(b.budget().active, 4u); +} + +VT_TEST(budget_on_changed_fires_on_starved_edge) { + SegmentBudget b({.max_active_segments = 1, .notify_period = std::chrono::milliseconds{40}}); + std::mutex m; + std::vector seen; + b.on_budget_changed([&](EB e) { + std::lock_guard lk(m); + seen.push_back(e); + }); + + FakeTask a{&b, tid(1)}, c{&b, tid(2)}; + b.register_task(tid(1), {.host = "h", .per_task_cap = 4, .resumable = true}, + [&](std::uint32_t n) { a.on_target(n); }); + b.register_task(tid(2), {.host = "h", .per_task_cap = 4, .resumable = true}, + [&](std::uint32_t n) { c.on_target(n); }); + b.set_task_order(std::vector{tid(1), tid(2)}); + b.set_want(tid(1), 4); + b.set_want(tid(2), 4); // budget is 1 -> tid(2) is starved: 0 -> nonzero edge + + // the edge fire is synchronous on the triggering call + bool saw_starved = false; + { + std::lock_guard lk(m); + for (auto &e : seen) + if (e.tasks_starved > 0) + saw_starved = true; + } + VT_CHECK(saw_starved); + + b.deregister_task(tid(1)); // frees the slot -> tid(2) no longer starved: edge back + std::this_thread::sleep_for(std::chrono::milliseconds(120)); + bool saw_unstarved_after = false; + { + std::lock_guard lk(m); + VT_CHECK(!seen.empty()); + saw_unstarved_after = seen.back().tasks_starved == 0; + } + VT_CHECK(saw_unstarved_after); +} + +VT_TEST(budget_concurrent_confirm_release_stays_consistent) { + SegmentBudget b({.max_active_segments = 16}); + constexpr int kTasks = 6; + std::vector> ts; + for (int i = 0; i < kTasks; ++i) { + ts.push_back(std::make_unique()); + ts.back()->budget = &b; + ts.back()->id = tid(i + 1); + FakeTask *ft = ts.back().get(); + b.register_task(ft->id, {.host = "h", .per_task_cap = 6, .resumable = true}, + [ft](std::uint32_t n) { ft->on_target(n); }); + } + std::vector drivers; + for (int i = 0; i < kTasks; ++i) { + drivers.emplace_back([&, id = tid(i + 1)] { + for (int r = 0; r < 4000; ++r) + b.set_want(id, (r % 7)); + }); + } + drivers.clear(); // join + for (auto &ft : ts) + b.set_want(ft->id, 0); + + // With everyone wanting nothing, the budget must be fully released. + VT_CHECK_EQ(b.budget().active, 0u); + std::uint32_t sum = 0; + for (auto &ft : ts) + sum += b.segments_active(ft->id); + VT_CHECK_EQ(sum, 0u); +} diff --git a/core/tests/segment/segmenter_test.cpp b/core/tests/segment/segmenter_test.cpp new file mode 100644 index 0000000..5633277 --- /dev/null +++ b/core/tests/segment/segmenter_test.cpp @@ -0,0 +1,257 @@ +#include "vdm/segment/segmenter.hpp" + +#include +#include +#include +#include + +#include "vtest.hpp" + +using namespace vdm::segment; + +namespace { +constexpr std::uint64_t MiB = 1u << 20; + +// Assign `n` slots (bounded by what the segmenter hands out) and return the indices. +std::vector fill(Segmenter &s, int n) { + std::vector idx; + for (int i = 0; i < n; ++i) { + auto a = s.assign_slot(); + if (!a) + break; + idx.push_back(*a); + } + return idx; +} + +// Do the segments (by their [start,end] at this instant) tile [0,total) with no overlap? +bool tiles_exactly(const Segmenter &s) { + auto snap = s.snapshot(); + std::vector> r; + for (auto &v : snap) + if (v.state != SegState::failed) + r.emplace_back(v.start, v.end); + std::sort(r.begin(), r.end()); + std::uint64_t cursor = 0; + for (auto [a, b] : r) { + if (a != cursor) + return false; + cursor = b + 1; + } + return cursor == s.total_size(); +} +} // namespace + +VT_TEST(seg_target_count_clamps) { + VT_CHECK_EQ(Segmenter(100 * MiB, 8, true).target_segment_count(), 8u); + VT_CHECK_EQ(Segmenter(100 * MiB, 64, true).target_segment_count(), 32u); // max 32 + VT_CHECK_EQ(Segmenter(100 * MiB, 0, true).target_segment_count(), 8u); // default + VT_CHECK_EQ(Segmenter(3 * MiB + 1, 8, true).target_segment_count(), 3u); // total/min + VT_CHECK_EQ(Segmenter(100 * MiB, 8, false).target_segment_count(), 1u); // non-resumable + VT_CHECK_EQ(Segmenter(0, 8, true).target_segment_count(), 1u); // chunked +} + +VT_TEST(seg_non_resumable_is_one_segment) { + Segmenter s(50 * MiB, 8, false); + auto idx = fill(s, 8); + VT_REQUIRE(idx.size() == 1); + auto snap = s.snapshot(); + VT_REQUIRE(snap.size() == 1); + VT_CHECK_EQ(snap[0].start, 0u); + VT_CHECK_EQ(snap[0].end, 50u * MiB - 1); +} + +VT_TEST(seg_initial_split_covers_range) { + Segmenter s(80 * MiB, 8, true); + auto idx = fill(s, 8); + VT_CHECK_EQ(idx.size(), 8u); + VT_CHECK(tiles_exactly(s)); + // no segment below the 1 MiB floor + for (auto &v : s.snapshot()) + VT_CHECK(v.length() >= MiB); +} + +VT_TEST(seg_split_stops_at_min_floor) { + // 5 MiB, floor 1 MiB, ask for 8: only ~5 splits possible (each half >= 1 MiB needs + // the parent >= 2 MiB), so we get fewer than 8. + Segmenter s(5 * MiB, 8, true); + auto idx = fill(s, 8); + VT_CHECK(idx.size() >= 1 && idx.size() <= 5); + VT_CHECK(tiles_exactly(s)); +} + +VT_TEST(seg_steal_takes_second_half_of_largest_remaining) { + Segmenter s(80 * MiB, 4, true); + auto idx = fill(s, 4); + VT_REQUIRE(idx.size() == 4); + + // Spread progress unevenly. Index != file position after splits, so identify the + // largest-remaining segment by scanning the snapshot, not by index. + s.advance(idx[1], 3 * MiB); + s.advance(idx[2], 7 * MiB); + s.advance(idx[3], 12 * MiB); + + SegmentView pre_victim{}; + std::uint64_t worst = 0; + for (auto &v : s.snapshot()) + if (v.index != idx[0] && v.remaining() > worst) { + worst = v.remaining(); + pre_victim = v; + } + + auto cont = s.on_complete(idx[0], /*may_steal=*/true); + VT_REQUIRE(cont.has_value()); + + SegmentView victim{}, fresh{}; + for (auto &v : s.snapshot()) { + if (v.index == pre_victim.index) + victim = v; + if (v.index == *cont) + fresh = v; + } + VT_CHECK_EQ(fresh.end, pre_victim.end); // fresh takes the tail of the victim's range + VT_CHECK_EQ(victim.end + 1, fresh.start); // contiguous, no gap / no overlap + VT_CHECK(victim.end < pre_victim.end); // the victim really did shrink + VT_CHECK(fresh.length() >= MiB); + VT_CHECK(victim.remaining() >= MiB); + // fresh got roughly the back half of what was remaining + VT_CHECK(fresh.length() >= worst / 2 - MiB && fresh.length() <= worst / 2 + MiB); + VT_CHECK(tiles_exactly(s)); +} + +VT_TEST(seg_complete_without_steal_releases) { + Segmenter s(4 * MiB, 2, true); + auto idx = fill(s, 2); // 2 x 2 MiB + VT_REQUIRE(idx.size() == 2); + s.advance(idx[1], 2 * MiB); + s.set_segment_state(idx[1], SegState::complete); + // idx[0] done, nothing left worth >= 1 MiB to steal -> release + s.advance(idx[0], 2 * MiB); + auto cont = s.on_complete(idx[0], true); + VT_CHECK(!cont.has_value()); +} + +VT_TEST(seg_yield_returns_nullopt) { + Segmenter s(80 * MiB, 4, true); + auto idx = fill(s, 4); + s.advance(idx[0], 20 * MiB); + auto cont = s.on_complete(idx[0], /*may_steal=*/false); // yielding + VT_CHECK(!cont.has_value()); +} + +VT_TEST(seg_third_connection_failure_with_mirror_requeues) { + Segmenter s(40 * MiB, 2, true); + auto idx = fill(s, 2); + s.advance(idx[0], 4 * MiB); + + VT_CHECK(s.on_failed(idx[0], /*conn=*/true, /*mirror=*/true) == FailAction::retry); + VT_CHECK(s.on_failed(idx[0], true, true) == FailAction::retry); + VT_CHECK(s.on_failed(idx[0], true, true) == FailAction::requeue); // 3rd + + VT_CHECK_EQ(s.segment_state(idx[0]), SegState::failed); + // the orphaned tail is now assignable again + auto again = s.assign_slot(); + VT_REQUIRE(again.has_value()); + auto snap = s.snapshot(); + SegmentView reborn{}; + for (auto &v : snap) + if (v.index == *again) + reborn = v; + VT_CHECK_EQ(reborn.start, 4u * MiB); // resumes where the failed one stopped + VT_CHECK_EQ(reborn.end, 20u * MiB - 1); // its half of the file +} + +VT_TEST(seg_failure_without_mirror_always_retries) { + Segmenter s(40 * MiB, 2, true); + auto idx = fill(s, 2); + for (int i = 0; i < 6; ++i) + VT_CHECK(s.on_failed(idx[0], true, /*mirror=*/false) == FailAction::retry); + // a non-connection error also retries regardless of count + VT_CHECK(s.on_failed(idx[1], /*conn=*/false, /*mirror=*/true) == FailAction::retry); +} + +VT_TEST(seg_note_connected_resets_failure_count) { + Segmenter s(40 * MiB, 2, true); + auto idx = fill(s, 2); + s.on_failed(idx[0], true, true); + s.on_failed(idx[0], true, true); + s.note_connected(idx[0]); + VT_CHECK(s.on_failed(idx[0], true, true) == FailAction::retry); // count restarted +} + +VT_TEST(seg_resume_from_meta_table) { + std::vector table = { + {0, 9 * MiB - 1, 9 * MiB}, // fully done + {9 * MiB, 19 * MiB - 1, 3 * MiB}, // partial + {19 * MiB, 40 * MiB - 1, 0}, // untouched + }; + Segmenter s(40 * MiB, 8, table, true); + auto snap = s.snapshot(); + VT_REQUIRE(snap.size() == 3); + VT_CHECK_EQ(snap[0].state, SegState::complete); + VT_CHECK_EQ(snap[1].completed, 3u * MiB); + VT_CHECK_EQ(s.downloaded(), 12u * MiB); + VT_CHECK(tiles_exactly(s)); + + // assign hands out the two incomplete ranges before splitting + auto a = s.assign_slot(); + auto b = s.assign_slot(); + VT_REQUIRE(a && b); +} + +VT_TEST(seg_resume_from_bad_table_falls_back) { + std::vector gappy = {{0, 4 * MiB - 1, 0}, {8 * MiB, 40 * MiB - 1, 0}}; + Segmenter s(40 * MiB, 8, gappy, true); + VT_CHECK(s.snapshot().empty()); // lazy fresh layout + auto idx = fill(s, 8); + VT_CHECK(idx.size() >= 1); + VT_CHECK(tiles_exactly(s)); +} + +VT_TEST(seg_all_complete_and_downloaded) { + Segmenter s(8 * MiB, 4, true); + auto idx = fill(s, 4); + VT_CHECK(!s.all_complete()); + for (auto i : idx) { + std::uint64_t len = s.segment_end(i) - s.segment_start(i) + 1; + s.advance(i, len); + s.set_segment_state(i, SegState::complete); + } + VT_CHECK(s.all_complete()); + VT_CHECK_EQ(s.downloaded(), 8u * MiB); +} + +// --- the steal path under the sanitizers ----------------------------------------------- + +VT_TEST(seg_concurrent_steal_and_advance) { + constexpr std::uint64_t total = 64 * MiB; + Segmenter s(total, 8, true); + auto idx = fill(s, 8); + VT_REQUIRE(idx.size() == 8); + + std::vector workers; + for (std::uint32_t w = 0; w < 8; ++w) { + workers.emplace_back([&s, seg = idx[w]]() mutable { + std::uint32_t cur = seg; + for (int guard = 0; guard < 200000; ++guard) { + const std::uint64_t start = s.segment_start(cur); + const std::uint64_t end = s.segment_end(cur); // may shrink under a steal + const std::uint64_t len = end - start + 1; + const std::uint64_t done = s.segment_completed(cur); + if (done >= len) { + auto nxt = s.on_complete(cur, /*may_steal=*/true); + if (!nxt) + return; // nothing left to steal — this worker is finished + cur = *nxt; + continue; + } + s.advance(cur, std::min(done + 128 * 1024, len)); + } + }); + } + workers.clear(); // join + + VT_CHECK(s.all_complete()); + VT_CHECK_EQ(s.downloaded(), total); + VT_CHECK(tiles_exactly(s)); +}