sami/vdm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
c710481018050e0941b293d9954ead3c566c5772
vdm/core/src/task/download_task.cpp
T
samiandClaude Sonnet 5 ef816c21fb core: fix quiesce() racing an in-flight write callback (ASan use-after-free) 
DownloadTaskState::quiesce() (Engine shutdown / ~Engine, via quiesce_task())
cancelled every live worker's transfer, then immediately cleared `workers` on
the calling thread. transfer.cancel() only *requests* the HttpClient worker
thread stop the transfer -- it does not wait for that to happen. If that
thread was mid write-callback (seg_data -> WriteBuffer::append ->
SparseFile::write_at), clearing the map destroyed the SegWorker (and its
ring buffer) it was still writing through: a heap-use-after-free, caught by
ASan via tools/bench alloc-check, which by design drops its Engine while a
download is still active mid-sample.

Every other exit path (verify/fail/auto_pause/demote, via begin_drain_locked)
already gets this right: cancel, then let each worker remove and flush
itself through seg_finished once HttpClient actually confirms the transfer
stopped, on the correct thread. quiesce() now does the same instead of
tearing the map down itself -- wait on a condition variable, notified from
seg_finished right after it erases, until `workers` is empty.

Co-Authored-By: Claude Sonnet 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01Q3QrF7rCt21bkAjt9BCDFQ
2026-09-11 13:12:57 +04:00

1292 lines
47 KiB
C++
Raw Blame History

// vdm/task/download_task.cpp — the download state machine behind DownloadHandle.
//
// One std::mutex `mu` is the task lock: every state transition, worker start/stop, and
// terminal path runs under it. The curl write-path (seg_data / seg_head) touches only
// per-segment state, guarded by a shared_mutex over the worker map. Segment completions
// and timers are posted to the engine timer thread and run under `mu`. Callbacks to
// DAEMON are collected under `mu` (defer()) and fired only after it is released
// (flush_deferred(), which guards its own queue with a separate mutex).
#include "task/download_task.hpp"
#include <unistd.h>
#include <algorithm>
#include <atomic>
#include <cctype>
#include <cerrno>
#include <condition_variable>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <shared_mutex>
#include <unordered_map>
#include "task/digest.hpp"
#include "vdm/io/sparse_file.hpp"
#include "vdm/io/write_buffer.hpp"
#include "vdm/meta/veloxpart.hpp"
#include "vdm/net/url.hpp"
#include "vdm/segment/segmenter.hpp"
namespace vdm::task {
namespace {
constexpr std::uint64_t kBufFloor = 64u * 1024;
constexpr std::uint64_t kBufCeil = 16u * 1024 * 1024;
constexpr std::uint64_t kNoEnd = ~std::uint64_t{0};
bool is_connection_error(Error e) noexcept {
switch (e) {
case Error::connection_reset:
case Error::timeout:
case Error::connect_failed:
case Error::resolve_failed:
case Error::tls_failed:
return true;
default:
return false;
}
}
std::chrono::milliseconds backoff_for(int attempt) {
// 1st retry ~1s, then 2, 4, 8, ... capped at 60s, +/-20% jitter (docs/04 §7).
int shift = std::clamp(attempt - 1, 0, 6);
std::uint64_t base = 1000ull << shift;
base = std::min<std::uint64_t>(base, 60'000);
std::uint64_t jitter = base / 5;
std::uint64_t r = static_cast<std::uint64_t>(std::rand()) % (2 * jitter + 1);
return std::chrono::milliseconds{base - jitter + r};
}
std::string lower(std::string s) {
std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return std::tolower(c); });
return s;
}
} // namespace
// What to do once every worker has drained (see DownloadTaskState::begin_drain_locked).
// A worker that hits one of these outcomes must not tear the task down itself: siblings may
// still be mid-transfer, and cutting them off synchronously — cancel every worker, clear the
// map, proceed — drops their buffered-but-unflushed bytes while segment_completed() already
// counts those bytes as done. That's exactly the class of bug that makes a later resume skip
// real data (seg_data() runs append() without workers_mu, so a sibling's buffer can't be
// flushed safely from here anyway). Instead: cancel the siblings, remember what to do, and
// let each one's own seg_finished (which already flushes on every exit path) run it once the
// worker map is actually empty.
enum class PendingAction { none, verify, fail, auto_pause, demote };
struct SegWorker {
std::uint32_t seg_index = 0;
net::Transfer transfer;
std::unique_ptr<io::WriteBuffer> buf;
std::uint64_t base_completed = 0;
std::uint64_t recv = 0;
long http_status = 0;
bool needs_auth = false;
bool wrong_status = false;
bool range_bad = false;
bool auth_handshake = false; // saw a 401/407 and let libcurl resend with credentials
std::string resp_etag, resp_last_modified; // captured on a wrong_status 200, for demote
std::optional<ErrorInfo> flush_error;
int retries = 0;
SteadyTime sample_at{};
std::uint64_t sample_bytes = 0;
double speed_bps = 0;
};
struct DownloadTaskState : std::enable_shared_from_this<DownloadTaskState> {
TaskHost &host;
TaskId id;
DownloadSpec spec;
DownloadCallbacks cbs;
std::mutex mu;
std::shared_mutex workers_mu;
std::mutex deferred_mu;
// Notified (holding mu) whenever seg_finished removes an entry from `workers`; quiesce()
// waits on it instead of clearing the map itself — see quiesce()'s comment.
std::condition_variable workers_drained_cv;
EngineState state = EngineState::probing;
std::optional<ErrorInfo> last_error;
std::atomic<bool> retired{false};
std::string part_path, meta_path;
net::ProbeResult probe;
bool have_probe = false;
bool resumable = false;
bool probe_needed_auth = false;
std::optional<std::uint64_t> total_size;
std::string origin_host;
std::unique_ptr<segment::Segmenter> seg;
std::unique_ptr<io::SparseFile> file;
std::unordered_map<std::uint32_t, std::unique_ptr<SegWorker>> workers;
std::unordered_map<std::uint32_t, int> retry_counts;
std::uint32_t slot_target = 0;
std::uint32_t effective_buffer = 0;
std::uint32_t requested_segments = 8;
bool registered = false;
bool pause_requested = false;
bool cancel_requested = false;
bool discard_on_cancel = false;
bool awaiting_auth = false;
bool awaiting_decision = false;
int max_retries = 10;
// Set by begin_drain_locked while waiting for sibling workers to drain; see
// PendingAction above.
PendingAction pending_action = PendingAction::none;
std::optional<ErrorInfo> pending_error;
bool pending_auth = false;
bool pending_decision = false;
std::atomic<std::int64_t> last_progress_ns{0};
SteadyTime started_at{};
std::vector<std::function<void()>> deferred;
DownloadTaskState(TaskHost &h, TaskId i, DownloadSpec s, DownloadCallbacks c)
: host(h), id(i), spec(std::move(s)), cbs(std::move(c)) {}
// --- deferred callbacks -------------------------------------------------------------
void defer(std::function<void()> fn) {
std::lock_guard lk(deferred_mu);
deferred.push_back(std::move(fn));
}
void flush_deferred() { // MUST be called with `mu` NOT held
std::vector<std::function<void()>> run;
{
std::lock_guard lk(deferred_mu);
run.swap(deferred);
}
for (auto &fn : run)
fn();
}
// caller holds mu
void transition(EngineState to, std::optional<ErrorInfo> err) {
if (state == to && !err)
return;
EngineState from = state;
state = to;
last_error = err;
auto cb = cbs.on_state;
defer([cb, from, to, err] {
if (cb)
cb(from, to, err);
});
}
[[nodiscard]] std::string current_url() const {
return have_probe && !probe.effective_url.empty() ? probe.effective_url : spec.url;
}
[[nodiscard]] bool has_mirror() const { return !spec.mirrors.empty(); }
[[nodiscard]] std::uint32_t want_slots() const {
if (!seg)
return requested_segments ? requested_segments : 1;
std::uint32_t target = seg->target_segment_count();
std::uint32_t incomplete = 0;
bool any = false;
for (auto &v : seg->snapshot()) {
any = true;
if (v.state != segment::SegState::complete && v.state != segment::SegState::failed)
++incomplete;
}
if (!any)
return target; // fresh: nothing created yet
if (incomplete == 0)
return target; // orphans pending re-split
return std::clamp<std::uint32_t>(incomplete, 1, target);
}
void begin();
void on_probe_result(Result<net::ProbeResult> r);
void finish_probe_locked();
void apply_slot_target(std::uint32_t n);
void fill_slots_locked();
void start_worker_locked(std::uint32_t seg_idx);
void restart_probe(bool with_auth);
net::DataAction seg_head(std::uint32_t seg_idx, const net::ResponseHead &h);
net::DataAction seg_data(std::uint32_t seg_idx, ConstByteSpan span);
void on_transfer_done(std::uint32_t seg_idx, Result<net::TransferStats> r);
void seg_finished(std::uint32_t seg_idx, Result<net::TransferStats> r);
void retry_worker(std::uint32_t seg_idx);
void begin_verify_locked();
void fail_locked(ErrorInfo e);
void auto_pause_locked(ErrorInfo e, bool auth, bool decision);
void demote_to_single_segment_locked();
void begin_drain_locked(PendingAction action, ErrorInfo e, bool auth, bool decision);
void finalize_cancel_locked();
void write_sidecar_locked();
void emit_progress_if_due();
void do_pause();
void do_resume();
void do_cancel(bool discard);
void do_provide_auth(std::string u, std::string p, bool remember);
void do_decide(Decision d);
void do_refresh_url(std::string url, std::vector<net::HeaderField> headers);
EngineState snapshot_state();
Progress snapshot_progress();
void quiesce(); // engine shutdown: stop everything, fire nothing
};
// ==================================================================================
void DownloadTaskState::begin() {
bool probe_hint;
{
std::unique_lock lk(mu);
part_path = spec.save_path + ".veloxpart";
meta_path = part_path + ".meta";
started_at = std::chrono::steady_clock::now();
max_retries =
static_cast<int>(spec.max_retries.value_or(host.config().default_max_retries));
requested_segments = spec.segments.value_or(host.config().default_segments);
probe_hint = spec.probe_hint.has_value();
if (probe_hint) {
probe = *spec.probe_hint;
have_probe = true;
finish_probe_locked();
} else {
transition(EngineState::probing, std::nullopt);
}
}
flush_deferred();
if (!probe_hint)
restart_probe(false);
}
void DownloadTaskState::restart_probe(bool with_auth) {
net::ProbeRequest pr;
pr.url = spec.url;
pr.headers = spec.headers;
pr.cookies = spec.cookies;
pr.referrer = spec.referrer;
pr.user_agent = spec.user_agent;
pr.proxy = spec.proxy;
if (with_auth)
pr.auth = spec.auth;
auto wp = weak_from_this();
host.probe(std::move(pr), [wp](Result<net::ProbeResult> r) {
if (auto s = wp.lock())
s->on_probe_result(std::move(r));
});
}
void DownloadTaskState::on_probe_result(Result<net::ProbeResult> r) {
{
std::unique_lock lk(mu);
if (retired.load() || is_terminal(state))
return;
if (!r.has_value()) {
fail_locked(std::move(r).error());
} else {
probe = std::move(r).value();
have_probe = true;
if (probe.requires_auth) {
probe_needed_auth = true;
auto_pause_locked(
ErrorInfo(Error::auth_required, "probe 401/407", probe.http_status), true,
false);
} else {
finish_probe_locked();
}
}
}
flush_deferred();
}
void DownloadTaskState::finish_probe_locked() {
resumable = probe.resumable;
total_size = probe.total_size;
origin_host = net::split_url(current_url()).host;
const std::uint64_t min_seg = host.config().min_segment_bytes;
const std::uint64_t total = total_size.value_or(0);
std::vector<segment::ResumedRange> resumed;
bool do_resume = spec.allow_resume;
if (do_resume) {
auto m = meta::read_veloxpart_file(meta_path);
if (m.has_value() && total > 0 && m.value().total_size == total &&
(m.value().etag.empty() || m.value().etag == probe.etag) &&
(m.value().last_modified.empty() || m.value().last_modified == probe.last_modified)) {
for (auto &s : m.value().segments)
resumed.push_back({s.start, s.end, s.completed});
} else {
do_resume = false;
}
}
file = std::make_unique<io::SparseFile>();
io::SparseFile::OpenOptions oo;
oo.total_size = total;
oo.preallocate = total > 0;
oo.truncate_existing = !do_resume;
if (auto o = file->open(part_path, oo); !o.has_value()) {
fail_locked(std::move(o).error());
return;
}
if (do_resume && !resumed.empty())
seg = std::make_unique<segment::Segmenter>(total, requested_segments, resumed, resumable,
min_seg);
else
seg = std::make_unique<segment::Segmenter>(total, requested_segments, resumable, min_seg);
const std::uint32_t cap = resumable ? std::min<std::uint32_t>(requested_segments, 32) : 1;
std::uint64_t want_buf = spec.buffer_bytes.value_or(host.config().default_buffer_bytes);
want_buf = std::clamp<std::uint64_t>(want_buf, kBufFloor, kBufCeil);
std::uint64_t max_per = std::max<std::uint64_t>(
kBufFloor, host.config().max_total_buffer_bytes / std::max(1u, cap));
effective_buffer = static_cast<std::uint32_t>(std::min(want_buf, max_per));
auto wp = weak_from_this();
host.budget().register_task(id, {origin_host, cap, resumable}, [wp](std::uint32_t n) {
auto s = wp.lock();
if (!s)
return;
s->host.schedule(std::chrono::steady_clock::now(), [wp, n] {
if (auto s2 = wp.lock())
s2->apply_slot_target(n);
});
});
registered = true;
host.limiter().attach_task(id, std::nullopt);
transition(EngineState::connecting, std::nullopt);
host.budget().set_want(id, want_slots());
}
// Start workers up to `slot_target`, given whatever the budget currently confirms. Shared
// by apply_slot_target() (the budget's async callback, whenever the computed target
// actually changes) and demote_to_single_segment_locked() -- the demoted target can
// legitimately equal what it was before the rebuild (e.g. a download already down to its
// last segment), in which case SegmentBudget::set_want() no-ops and the async callback
// never fires, so nothing else would ever start the replacement worker.
void DownloadTaskState::fill_slots_locked() {
while (workers.size() < slot_target) {
auto s = seg->assign_slot();
if (!s) {
host.budget().set_want(id, static_cast<std::uint32_t>(workers.size()));
break;
}
if (!host.budget().confirm_slot(id)) {
seg->set_segment_state(*s, segment::SegState::idle);
break;
}
start_worker_locked(*s);
}
if (state == EngineState::connecting && !workers.empty())
transition(EngineState::downloading, std::nullopt);
}
void DownloadTaskState::apply_slot_target(std::uint32_t n) {
{
std::unique_lock lk(mu);
if (retired.load() || is_terminal(state) || pause_requested || cancel_requested ||
awaiting_auth || awaiting_decision || pending_action != PendingAction::none || !seg)
return;
slot_target = n;
fill_slots_locked();
}
flush_deferred();
}
void DownloadTaskState::start_worker_locked(std::uint32_t seg_idx) {
auto w = std::make_unique<SegWorker>();
w->seg_index = seg_idx;
const std::uint64_t sstart = seg->segment_start(seg_idx);
const std::uint64_t completed = seg->segment_completed(seg_idx);
const std::uint64_t send = seg->segment_end(seg_idx);
w->base_completed = completed;
w->sample_at = std::chrono::steady_clock::now();
if (auto it = retry_counts.find(seg_idx); it != retry_counts.end())
w->retries = it->second;
auto wp = weak_from_this();
w->buf = std::make_unique<io::WriteBuffer>(
sstart + completed, effective_buffer,
[wp](std::uint64_t off, ConstByteSpan sp) -> Result<void> {
if (auto s = wp.lock())
return s->file->write_at(off, sp);
return ErrorInfo(Error::canceled, "task gone");
});
net::Request req;
req.url = current_url();
req.headers = spec.headers;
req.cookies = spec.cookies;
req.referrer = spec.referrer;
req.user_agent = spec.user_agent;
req.proxy = spec.proxy;
req.auth = spec.auth;
req.follow_redirects = true;
req.accept_encoding = false;
req.low_speed_bytes_per_sec = 1024;
req.low_speed_secs = 30;
// Only range-request a resumable source; a non-resumable single segment is a plain GET.
if (resumable && total_size && send != kNoEnd)
req.range = net::ByteRange{sstart + completed, send};
if (completed > 0 && (!probe.etag.empty() || !probe.last_modified.empty()))
req.headers.push_back({"If-Range", !probe.etag.empty() ? probe.etag : probe.last_modified});
net::TransferCallbacks tc;
tc.on_head = [wp, seg_idx](const net::ResponseHead &h) {
if (auto s = wp.lock())
return s->seg_head(seg_idx, h);
return net::DataAction::abort;
};
tc.on_data = [wp, seg_idx](ConstByteSpan sp) {
if (auto s = wp.lock())
return s->seg_data(seg_idx, sp);
return net::DataAction::abort;
};
tc.on_finished = [wp, seg_idx](Result<net::TransferStats> r) {
if (auto s = wp.lock())
s->on_transfer_done(seg_idx, std::move(r));
};
auto *raw = w.get();
{
std::unique_lock wl(workers_mu);
workers[seg_idx] = std::move(w);
}
raw->transfer = host.http().start(std::move(req), std::move(tc));
seg->set_segment_state(seg_idx, segment::SegState::connecting);
}
// --- curl write path -------------------------------------------------------------------
net::DataAction DownloadTaskState::seg_head(std::uint32_t seg_idx, const net::ResponseHead &h) {
SegWorker *w = nullptr;
{
std::shared_lock lk(workers_mu);
auto it = workers.find(seg_idx);
if (it == workers.end())
return net::DataAction::abort;
w = it->second.get();
}
w->http_status = h.status;
if (h.status == 401 || h.status == 407) {
// libcurl with CURLAUTH_ANY does a challenge round-trip: the first response is a
// 401/407, then it resends with credentials. Don't kill that handshake — let it
// proceed once and wait for the real status on the next header block. Only when we
// have nothing to try, or we already tried, is this a genuine "auth required".
const bool have_creds =
spec.auth.scheme != net::AuthScheme::none && !spec.auth.username.empty();
if (have_creds && !w->auth_handshake) {
w->auth_handshake = true;
return net::DataAction::proceed;
}
w->needs_auth = true;
return net::DataAction::abort;
}
// A 200 where we sent a Range is only "the file changed under us" when we are running
// a resumable, multi-segment transfer. A non-resumable single-segment GET legitimately
// gets a 200 (the source has no Range support).
if (resumable && total_size && *total_size > 0 && h.status == 200) {
w->wrong_status = true;
if (auto v = h.headers.get("ETag"))
w->resp_etag.assign(*v);
if (auto v = h.headers.get("Last-Modified"))
w->resp_last_modified.assign(*v);
return net::DataAction::abort;
}
if (h.status == 416) {
w->range_bad = true;
return net::DataAction::abort;
}
if (h.status >= 400)
return net::DataAction::abort;
seg->set_segment_state(seg_idx, segment::SegState::downloading);
return net::DataAction::proceed;
}
net::DataAction DownloadTaskState::seg_data(std::uint32_t seg_idx, ConstByteSpan span) {
SegWorker *w = nullptr;
{
std::shared_lock lk(workers_mu);
auto it = workers.find(seg_idx);
if (it == workers.end())
return net::DataAction::abort;
w = it->second.get();
}
if (span.empty())
return net::DataAction::proceed;
// Body of libcurl's pre-auth 401/407 response — discard it; the real body follows the
// resend under a 2xx header.
if (w->http_status == 401 || w->http_status == 407)
return net::DataAction::proceed;
auto wait = host.limiter().acquire(id, span.size());
if (wait.count() > 0) {
net::Transfer t = w->transfer;
host.schedule(std::chrono::steady_clock::now() + wait, [t]() mutable { t.resume(); });
return net::DataAction::pause;
}
if (auto r = w->buf->append(span); !r.has_value()) {
w->flush_error = std::move(r).error();
return net::DataAction::abort;
}
w->recv += span.size();
seg->advance(seg_idx, w->base_completed + w->recv);
auto now = std::chrono::steady_clock::now();
auto dt = std::chrono::duration<double>(now - w->sample_at).count();
if (dt >= 0.5) {
double inst = static_cast<double>(w->recv - w->sample_bytes) / dt;
w->speed_bps = w->speed_bps == 0 ? inst : 0.7 * w->speed_bps + 0.3 * inst;
w->sample_at = now;
w->sample_bytes = w->recv;
}
emit_progress_if_due();
return net::DataAction::proceed;
}
void DownloadTaskState::on_transfer_done(std::uint32_t seg_idx, Result<net::TransferStats> r) {
auto wp = weak_from_this();
host.schedule(std::chrono::steady_clock::now(), [wp, seg_idx, r = std::move(r)]() mutable {
if (auto s = wp.lock())
s->seg_finished(seg_idx, std::move(r));
});
}
void DownloadTaskState::seg_finished(std::uint32_t seg_idx, Result<net::TransferStats> r) {
std::unique_lock lk(mu);
std::unique_ptr<SegWorker> w;
{
std::unique_lock wl(workers_mu);
auto it = workers.find(seg_idx);
if (it == workers.end()) {
lk.unlock();
flush_deferred();
return;
}
w = std::move(it->second);
workers.erase(it);
}
workers_drained_cv.notify_all(); // quiesce() may be waiting for `workers` to empty out
if (retired.load()) { // engine shutting down / already terminal — no more callbacks
if (w->buf)
(void)w->buf->flush();
lk.unlock();
return;
}
// The bytes may all have arrived even though the connection closed dirty
// (content-length-mismatch's honest-length lie, flaky-reset's tail, a proxy RST after
// the last byte). If the segment is fully covered, that's a success.
if (seg && !cancel_requested && !pause_requested && !w->needs_auth && !w->wrong_status &&
!w->flush_error && !w->range_bad) {
const std::uint64_t len = seg->segment_end(seg_idx) - seg->segment_start(seg_idx) + 1;
if (len != 0 && seg->segment_completed(seg_idx) >= len) {
r = Result<net::TransferStats>(net::TransferStats{});
}
}
auto release_slot = [&] {
if (registered)
host.budget().release_slot(id);
};
auto done = [&] {
lk.unlock();
flush_deferred();
};
if (cancel_requested) {
if (w->buf)
(void)w->buf->flush();
release_slot();
if (workers.empty())
finalize_cancel_locked();
return done();
}
if (pause_requested) {
if (w->buf)
(void)w->buf->flush();
seg->advance(seg_idx, w->base_completed + w->recv);
// Hand the range back unassigned so resume's assign_slot() picks it up instead of
// treating it as still-held and splitting it (which would orphan its front half).
seg->release_segment(seg_idx);
release_slot();
if (workers.empty()) {
(void)file->sync();
write_sidecar_locked();
transition(EngineState::paused, std::nullopt);
}
return done();
}
if (pending_action != PendingAction::none) {
// A sibling already decided the task is finishing (verify / fail / auto-pause /
// demote); this worker's own outcome no longer matters. Drain it like every other
// exit path: flush its buffer so segment_completed() stays true to disk, then hand
// off to whichever worker finds the map empty.
if (w->buf)
(void)w->buf->flush(); // best-effort: we're already tearing down for another
// reason, and the pending action doesn't depend on
// this segment reaching any particular state.
seg->advance(seg_idx, w->base_completed + w->recv);
release_slot();
if (workers.empty()) {
PendingAction action = std::exchange(pending_action, PendingAction::none);
ErrorInfo e = pending_error.value_or(ErrorInfo(Error::internal, ""));
bool auth = pending_auth, decision = pending_decision;
switch (action) {
case PendingAction::verify:
begin_verify_locked();
break;
case PendingAction::fail:
fail_locked(e);
break;
case PendingAction::auto_pause:
auto_pause_locked(e, auth, decision);
break;
case PendingAction::demote:
demote_to_single_segment_locked();
break;
case PendingAction::none:
break;
}
}
return done();
}
if (w->needs_auth) {
release_slot();
auto_pause_locked(ErrorInfo(Error::auth_required, "401/407", w->http_status), true, false);
return done();
}
if (w->wrong_status) {
release_slot();
// A 200 where 206 was expected is ambiguous: the file really changed (ask, don't
// corrupt — docs/04 §5), or the server just stopped honouring Range for this
// connection while it's still the same file (docs/04 §7: demote to 1 segment and
// continue). ETag/Last-Modified from the 200 itself, compared against what the
// probe recorded, is the only signal that tells them apart. Prefer ETag strictly
// when both sides have one — same rule If-Range itself uses — and only fall back to
// Last-Modified when there's no ETag to compare; a coarse (often second-resolution)
// Last-Modified that happens to match is weak evidence next to a mismatching ETag.
bool same_file;
if (!probe.etag.empty() && !w->resp_etag.empty())
same_file = probe.etag == w->resp_etag;
else if (!probe.last_modified.empty() && !w->resp_last_modified.empty())
same_file = probe.last_modified == w->resp_last_modified;
else
same_file = false; // no validator to compare -> can't prove it, ask
if (same_file) {
demote_to_single_segment_locked();
} else {
auto_pause_locked(ErrorInfo(Error::server_file_changed, "200 where 206 expected"),
false, true);
}
return done();
}
if (w->flush_error) {
ErrorInfo e = *w->flush_error;
release_slot();
if (e.code == Error::disk_full) {
auto_pause_locked(e, false, false);
} else {
fail_locked(e);
}
return done();
}
if (w->range_bad) {
// 416 mid-download means our range metadata is stale (docs/04 §7): re-probe and
// re-split rather than retrying the same now-invalid range until it exhausts.
release_slot();
auto_pause_locked(ErrorInfo(Error::range_not_satisfiable, "416"), false, true);
return done();
}
if (!r.has_value()) {
ErrorInfo e = std::move(r).error();
// Persist what this attempt received before the connection dropped. seg_data() has
// already advanced `completed` past these bytes, but they are still only in the
// WriteBuffer (nothing forces a flush until the buffer fills or the segment ends).
// Without this, a retry resumes from `completed` and skips an unwritten hole — for
// a small file that never fills the buffer, every reset loses its whole payload.
if (w->buf) {
if (auto fl = w->buf->flush(); !fl.has_value()) {
release_slot();
fail_locked(std::move(fl).error());
return done();
}
}
const bool conn = is_connection_error(e.code);
segment::FailAction fa = seg->on_failed(seg_idx, conn, has_mirror());
if (w->recv > 0)
w->retries = 0; // progress this attempt: not a stuck segment
if (fa == segment::FailAction::requeue) {
retry_counts.erase(seg_idx);
release_slot();
host.budget().set_want(id, want_slots());
} else if (++w->retries > max_retries) {
ErrorInfo x(Error::max_retries_exhausted, e.to_string());
x.cause = e.code;
release_slot();
fail_locked(x);
} else {
retry_counts[seg_idx] = w->retries;
seg->set_segment_state(seg_idx, segment::SegState::stalled);
release_slot(); // give the slot back while we back off
int attempt = w->retries;
auto wp = weak_from_this();
host.schedule(std::chrono::steady_clock::now() + backoff_for(attempt), [wp, seg_idx] {
if (auto s = wp.lock())
s->retry_worker(seg_idx);
});
if (workers.empty())
transition(EngineState::retry_wait, std::nullopt);
}
return done();
}
// success
if (w->buf) {
if (auto f = w->buf->flush(); !f.has_value()) {
release_slot();
fail_locked(std::move(f).error());
return done();
}
}
retry_counts.erase(seg_idx);
const bool may_steal = (workers.size() + 1) <= slot_target && !pause_requested;
auto cont = seg->on_complete(seg_idx, may_steal);
(void)file->sync();
write_sidecar_locked();
if (cont)
start_worker_locked(*cont); // slot-neutral steal
else
release_slot();
if (seg->all_complete())
begin_verify_locked(); // drain-aware: defers if other workers are still live
return done();
}
void DownloadTaskState::retry_worker(std::uint32_t seg_idx) {
{
std::unique_lock lk(mu);
if (retired.load() || is_terminal(state) || pause_requested || cancel_requested ||
pending_action != PendingAction::none || !seg)
return;
if (workers.count(seg_idx))
return;
if (!host.budget().confirm_slot(id))
return; // budget is full; try again later
if (state == EngineState::retry_wait)
transition(EngineState::connecting, std::nullopt);
start_worker_locked(seg_idx);
if (!workers.empty() && state == EngineState::connecting)
transition(EngineState::downloading, std::nullopt);
}
flush_deferred();
}
void DownloadTaskState::begin_verify_locked() {
if (!workers.empty()) {
begin_drain_locked(PendingAction::verify, ErrorInfo(Error::internal, ""), false, false);
return;
}
transition(EngineState::assembling, std::nullopt);
transition(EngineState::verifying, std::nullopt);
(void)file->sync();
(void)file->close();
std::optional<std::string> got;
if (spec.checksum) {
auto h = hash_file(part_path, spec.checksum->algo);
if (!h.has_value()) {
fail_locked(std::move(h).error());
return;
}
got = h.value();
if (*got != lower(spec.checksum->hex)) {
fail_locked(ErrorInfo(Error::checksum_mismatch,
"want " + lower(spec.checksum->hex) + " got " + *got));
return;
}
}
if (::rename(part_path.c_str(), spec.save_path.c_str()) != 0) {
fail_locked(ErrorInfo(Error::io_error, std::string("rename: ") + std::strerror(errno)));
return;
}
::unlink(meta_path.c_str());
if (registered)
host.budget().deregister_task(id);
host.limiter().detach_task(id);
registered = false;
transition(EngineState::complete, std::nullopt);
DownloadOutcome o;
o.final_path = spec.save_path;
o.bytes = seg->downloaded();
o.sha256_hex =
(spec.checksum && spec.checksum->algo == Checksum::Algo::sha256) ? got : std::nullopt;
o.elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
std::chrono::steady_clock::now() - started_at);
auto cb = cbs.on_finished;
defer([cb, o] {
if (cb)
cb(Result<DownloadOutcome>(o));
});
retired.store(true);
TaskHost *h = &host;
TaskId tid = id;
defer([h, tid] { h->task_retired(tid); });
}
void DownloadTaskState::fail_locked(ErrorInfo e) {
if (!workers.empty()) {
begin_drain_locked(PendingAction::fail, std::move(e), false, false);
return;
}
if (file)
(void)file->close();
if (seg)
write_sidecar_locked();
if (registered)
host.budget().deregister_task(id);
host.limiter().detach_task(id);
registered = false;
transition(EngineState::failed, e);
auto cb = cbs.on_finished;
defer([cb, e] {
if (cb)
cb(Result<DownloadOutcome>(ErrorInfo(e)));
});
retired.store(true);
TaskHost *h = &host;
TaskId tid = id;
defer([h, tid] { h->task_retired(tid); });
}
void DownloadTaskState::auto_pause_locked(ErrorInfo e, bool auth, bool decision) {
if (!workers.empty()) {
begin_drain_locked(PendingAction::auto_pause, std::move(e), auth, decision);
return;
}
awaiting_auth = auth;
awaiting_decision = decision;
if (file)
(void)file->sync();
if (seg)
write_sidecar_locked();
if (registered)
host.budget().set_want(id, 0);
transition(EngineState::paused, e);
if (auth) {
AuthChallenge ac;
ac.host = origin_host;
auto cb = cbs.on_auth_required;
defer([cb, ac] {
if (cb)
cb(ac);
});
}
if (decision) {
DecisionRequest dr;
dr.kind = e.code == Error::server_file_changed
? DecisionRequest::Kind::server_file_changed
: DecisionRequest::Kind::range_metadata_stale;
dr.detail = e.context;
auto cb = cbs.on_decision_needed;
defer([cb, dr] {
if (cb)
cb(dr);
});
}
}
// Cancel every live worker and remember what to do once they've all drained through
// seg_finished's PendingAction branch (see the enum's comment). Never clears `workers`
// itself — each worker removes itself, flushed, when its own transfer actually completes.
void DownloadTaskState::begin_drain_locked(PendingAction action, ErrorInfo e, bool auth,
bool decision) {
pending_action = action;
pending_error = std::move(e);
pending_auth = auth;
pending_decision = decision;
std::shared_lock wl(workers_mu);
for (auto &[idx, w] : workers)
w->transfer.cancel();
}
// The 200-where-206-expected we just saw carried the same ETag/Last-Modified the probe
// recorded: same file, the server (or this connection) just doesn't honour Range. Rebuild
// as a single non-resumable segment covering the whole file and keep going with a plain
// GET. It re-transfers bytes we may already have — there's no way to ask a Range-blind
// server for a suffix — but it never truncates or discards what's on disk, and a source
// that hasn't changed serves identical bytes, so the result is still byte-correct.
void DownloadTaskState::demote_to_single_segment_locked() {
if (!workers.empty()) {
begin_drain_locked(PendingAction::demote, ErrorInfo(Error::internal, ""), false, false);
return;
}
resumable = false;
seg = std::make_unique<segment::Segmenter>(total_size.value_or(0), 1, /*resumable=*/false,
host.config().min_segment_bytes);
transition(EngineState::connecting, std::nullopt);
if (registered) {
// set_want() alone is not enough: if the demoted target happens to equal what it
// was before the rebuild (e.g. this was already the last live segment), it's a
// no-op and the async budget callback never fires. Drive slot assignment directly.
slot_target = want_slots();
host.budget().set_want(id, slot_target);
fill_slots_locked();
}
}
void DownloadTaskState::finalize_cancel_locked() {
if (file)
(void)file->close();
if (registered)
host.budget().deregister_task(id);
host.limiter().detach_task(id);
registered = false;
if (discard_on_cancel) {
::unlink(part_path.c_str());
::unlink(meta_path.c_str());
} else if (seg) {
write_sidecar_locked();
}
transition(EngineState::cancelled, std::nullopt);
auto cb = cbs.on_finished;
defer([cb] {
if (cb)
cb(Result<DownloadOutcome>(ErrorInfo(Error::canceled)));
});
retired.store(true);
TaskHost *h = &host;
TaskId tid = id;
defer([h, tid] { h->task_retired(tid); });
}
void DownloadTaskState::write_sidecar_locked() {
if (!seg || !total_size)
return;
meta::VeloxPart vp;
vp.total_size = *total_size;
vp.urls = {spec.url};
if (!current_url().empty() && current_url() != spec.url)
vp.urls.push_back(current_url());
for (auto &m : spec.mirrors)
vp.urls.push_back(m);
vp.etag = probe.etag;
vp.last_modified = probe.last_modified;
vp.content_type = probe.mime;
std::uint64_t dl = 0;
for (auto &v : seg->snapshot()) {
if (v.state == segment::SegState::failed)
continue;
vp.segments.push_back({v.start, v.end, v.completed});
dl += v.completed;
}
vp.downloaded = dl;
(void)meta::write_veloxpart_file(meta_path, vp, true);
}
void DownloadTaskState::emit_progress_if_due() {
const auto now = std::chrono::steady_clock::now();
const std::int64_t now_ns = now.time_since_epoch().count();
std::int64_t prev = last_progress_ns.load(std::memory_order_relaxed);
if (now_ns - prev < 250'000'000)
return;
if (!last_progress_ns.compare_exchange_strong(prev, now_ns))
return;
Progress p;
{
std::shared_lock lk(workers_mu);
double agg = 0;
for (auto &[idx, w] : workers) {
agg += w->speed_bps;
SegmentProgress sp;
sp.index = idx;
sp.speed_bps = static_cast<std::uint64_t>(w->speed_bps);
p.segments.push_back(sp);
}
p.speed_bps = static_cast<std::uint64_t>(agg);
p.effective_segments = static_cast<std::uint32_t>(workers.size());
}
if (seg) {
p.downloaded = seg->downloaded();
for (auto &v : seg->snapshot())
for (auto &sp : p.segments)
if (sp.index == v.index) {
sp.start = v.start;
sp.end = v.end;
sp.completed = v.completed;
sp.state = v.state;
}
}
p.total = total_size;
p.effective_buffer_bytes = effective_buffer;
if (p.speed_bps > 0 && total_size && *total_size > p.downloaded)
p.eta_seconds = static_cast<std::uint32_t>((*total_size - p.downloaded) / p.speed_bps);
auto cb = cbs.on_progress;
if (cb)
cb(p);
}
// --- handle-facing --------------------------------------------------------------------
void DownloadTaskState::do_pause() {
{
std::unique_lock lk(mu);
if (state == EngineState::paused || is_terminal(state))
return;
pause_requested = true;
if (registered)
host.budget().set_want(id, 0);
if (workers.empty()) {
if (file)
(void)file->sync();
write_sidecar_locked();
transition(EngineState::paused, std::nullopt);
} else {
std::shared_lock wl(workers_mu);
for (auto &[idx, w] : workers)
w->transfer.cancel();
}
}
flush_deferred();
}
void DownloadTaskState::do_resume() {
{
std::unique_lock lk(mu);
if (state != EngineState::paused || awaiting_auth || awaiting_decision)
return;
pause_requested = false;
transition(EngineState::connecting, std::nullopt);
if (registered)
host.budget().set_want(id, want_slots());
}
flush_deferred();
}
void DownloadTaskState::do_cancel(bool discard) {
{
std::unique_lock lk(mu);
if (is_terminal(state))
return;
cancel_requested = true;
discard_on_cancel = discard;
if (workers.empty()) {
finalize_cancel_locked();
} else {
std::shared_lock wl(workers_mu);
for (auto &[idx, w] : workers)
w->transfer.cancel();
}
}
flush_deferred();
}
void DownloadTaskState::do_provide_auth(std::string u, std::string p, bool) {
bool reprobe = false;
{
std::unique_lock lk(mu);
if (!awaiting_auth)
return;
spec.auth = net::AuthConfig{net::AuthScheme::any, std::move(u), std::move(p)};
awaiting_auth = false;
reprobe = probe_needed_auth;
probe_needed_auth = false;
transition(EngineState::connecting, std::nullopt);
if (!reprobe && registered)
host.budget().set_want(id, want_slots());
}
flush_deferred();
if (reprobe)
restart_probe(true);
}
void DownloadTaskState::do_decide(Decision d) {
{
std::unique_lock lk(mu);
if (!awaiting_decision)
return;
awaiting_decision = false;
if (d == Decision::abort) {
// Surface the reason the decision was actually asked for (range_metadata_stale
// sets last_error to range_not_satisfiable, server_file_changed to itself), not
// a hardcoded label that would misreport a 416 as a changed file.
fail_locked(last_error.value_or(ErrorInfo(Error::server_file_changed, "user aborted")));
} else {
if (d == Decision::restart) {
::unlink(part_path.c_str());
::unlink(meta_path.c_str());
spec.allow_resume = false;
io::SparseFile::OpenOptions oo;
oo.total_size = total_size.value_or(0);
oo.preallocate = oo.total_size > 0;
oo.truncate_existing = true;
file = std::make_unique<io::SparseFile>();
(void)file->open(part_path, oo);
seg = std::make_unique<segment::Segmenter>(total_size.value_or(0),
requested_segments, resumable,
host.config().min_segment_bytes);
}
transition(EngineState::connecting, std::nullopt);
if (registered)
host.budget().set_want(id, want_slots());
}
}
flush_deferred();
}
void DownloadTaskState::do_refresh_url(std::string url, std::vector<net::HeaderField> headers) {
{
std::unique_lock lk(mu);
if (is_terminal(state))
return;
spec.url = std::move(url);
if (!headers.empty())
spec.headers = std::move(headers);
std::shared_lock wl(workers_mu);
for (auto &[idx, w] : workers)
w->transfer.cancel();
}
flush_deferred();
auto wp = weak_from_this();
net::ProbeRequest pr;
pr.url = spec.url;
pr.headers = spec.headers;
pr.auth = spec.auth;
pr.proxy = spec.proxy;
host.probe(std::move(pr), [wp](Result<net::ProbeResult> r) {
auto s = wp.lock();
if (!s)
return;
std::unique_lock lk(s->mu);
if (s->retired.load() || is_terminal(s->state))
return;
if (r.has_value()) {
s->probe.effective_url = r.value().effective_url;
s->probe.etag = r.value().etag;
s->probe.last_modified = r.value().last_modified;
}
if (s->registered)
s->host.budget().set_want(s->id, s->want_slots());
lk.unlock();
s->flush_deferred();
});
}
void DownloadTaskState::quiesce() {
std::unique_lock lk(mu);
retired.store(true);
{
std::shared_lock wl(workers_mu);
for (auto &[idx, w] : workers)
w->transfer.cancel();
}
// Do not clear `workers` here: transfer.cancel() only requests the HttpClient worker
// thread stop the transfer, asynchronously -- it does not wait for that to happen. A
// worker's SegWorker (and its WriteBuffer) may still be in active use by a curl write
// callback running on that other thread right now. Clearing the map out from under it
// was a real, ASan-caught heap-use-after-free (ring buffer freed here while
// SparseFile::write_at() on the HttpClient worker thread was still writing through it).
// Every worker removes and flushes itself, safely, via seg_finished once HttpClient
// actually confirms the transfer has stopped (same path every other exit uses; see the
// `retired` branch there) -- just wait for that to happen for all of them. Bounded by
// however long a cancelled curl transfer takes to unwind, not user-controllable.
workers_drained_cv.wait(lk, [this] { return workers.empty(); });
}
EngineState DownloadTaskState::snapshot_state() {
std::lock_guard lk(mu);
return state;
}
Progress DownloadTaskState::snapshot_progress() {
Progress p;
std::lock_guard lk(mu);
if (seg) {
p.downloaded = seg->downloaded();
for (auto &v : seg->snapshot()) {
SegmentProgress sp;
sp.index = v.index;
sp.start = v.start;
sp.end = v.end;
sp.completed = v.completed;
sp.state = v.state;
p.segments.push_back(sp);
}
}
p.total = total_size;
{
std::shared_lock wl(workers_mu);
p.effective_segments = static_cast<std::uint32_t>(workers.size());
}
p.effective_buffer_bytes = effective_buffer;
return p;
}
// ==================================================================================
std::shared_ptr<DownloadTaskState> create_task(TaskHost &host, TaskId id, DownloadSpec spec,
DownloadCallbacks cbs) {
auto s = std::make_shared<DownloadTaskState>(host, id, std::move(spec), std::move(cbs));
s->begin();
return s;
}
void quiesce_task(const std::shared_ptr<DownloadTaskState> &s) {
if (s)
s->quiesce();
}
// --- DownloadHandle bodies ----------------------------------------------------------
TaskId DownloadHandle::id() const noexcept {
return state_ ? state_->id : TaskId{};
}
void DownloadHandle::pause() {
if (state_)
state_->do_pause();
}
void DownloadHandle::resume() {
if (state_)
state_->do_resume();
}
void DownloadHandle::cancel(bool discard_partial) {
if (state_)
state_->do_cancel(discard_partial);
}
void DownloadHandle::provide_auth(std::string u, std::string p, bool remember) {
if (state_)
state_->do_provide_auth(std::move(u), std::move(p), remember);
}
void DownloadHandle::decide(Decision d) {
if (state_)
state_->do_decide(d);
}
void DownloadHandle::refresh_url(std::string url, std::vector<net::HeaderField> headers) {
if (state_)
state_->do_refresh_url(std::move(url), std::move(headers));
}
EngineState DownloadHandle::state() const {
return state_ ? state_->snapshot_state() : EngineState::failed;
}
Progress DownloadHandle::progress() const {
return state_ ? state_->snapshot_progress() : Progress{};
}
} // namespace vdm::task
Reference in New Issue View Git Blame Copy Permalink