vdm/core/src/segment/segmenter.cpp

// vdm/segment/segmenter.cpp

#include "vdm/segment/segmenter.hpp"

#include <algorithm>
#include <limits>

namespace vdm::segment {
namespace {
constexpr std::uint32_t kNoIndex = std::numeric_limits<std::uint32_t>::max();
constexpr std::uint64_t kU64Max = std::numeric_limits<std::uint64_t>::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<ResumedRange> &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<ResumedRange> 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::uint32_t>(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<std::uint64_t>(t, by_size);
    target_count_ = std::clamp<std::uint32_t>(static_cast<std::uint32_t>(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<std::uint32_t> 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<std::uint32_t> 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);
}

void Segmenter::release_segment(std::uint32_t idx) noexcept {
    std::lock_guard lk(mu_);
    if (idx >= segs_.size())
        return;
    segs_[idx].assigned = false;
    if (segs_[idx].state.load() != SegState::complete)
        segs_[idx].state.store(SegState::idle);
}

// --- 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<std::pair<std::uint64_t, std::uint64_t>> 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<SegmentView> Segmenter::snapshot() const {
    std::lock_guard lk(mu_);
    std::vector<SegmentView> 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