core: fix SegmentBudget over-admission and add a wait-list wakeup

Root cause of the M7 RSS gap (core/docs/m7-baseline.md: ~70 MB measured
against a 60 MB target): SegmentBudget::confirm_slot() only checked a task's
own held count against its own target -- never the engine-wide active_ sum.
reallocate_locked()'s two-pass fairness allocation does bound
sum(target) <= max_active_ at the moment it computes a plan, but that bound
says nothing about sum(held): a task can be legitimately holding more than
its own just-lowered target for a while (yield is deferred to a segment
boundary, never mid-segment -- ADR 0011 A1), and another task's target can
correctly rise to claim that capacity before the first task has physically
released it. Both confirm_slot() calls could then succeed against their own,
individually-correct targets while sum(held) exceeded max_active_ --
tools/bench heap-profile caught this directly: budget.active reading 56-86
against a total of 32.

confirm_slot() now also checks active_ < max_active_, unconditionally, as a
backstop that doesn't depend on any task's target bookkeeping being in sync
with what every other task holds. That creates a liveness question the
original design never answered: a task denied only by this new check has a
target that's already correct, so it never changes again and
reallocate_locked()'s plain "fire a callback when a task's target changes"
mechanism never revisits it. Task gained a waiting_for_slot flag, set on
exactly this denial; release_slot()/deregister_task() (the only two places
that free real capacity) now hand a freed slot directly to the
highest-priority waiting task via wake_one_waiter_locked(), if
reallocate_locked()'s own plan didn't already produce a callback for anyone.

download_task.cpp's fill_slots_locked() needed a matching fix: a woken
task's stalled segments (SegState::stalled -- backed off mid-retry, its own
release_slot() already called) have no live worker and never surface
through Segmenter::assign_slot(), which only hands out unassigned or fresh
ranges. fill_slots_locked() now restarts any stalled segment with no live
worker directly (bounded by slot_target, same as its assign_slot() loop)
before looking for new work; a segment it doesn't get to keeps its own
scheduled retry_worker() timer as a second chance.

Also fixes a real TSan-caught data race this work surfaced: SegWorker::
speed_bps was written only by its own segment's curl callback and, before
Progress.speed_bps's polled-path fix, only ever read from that same thread
-- safe without synchronization. snapshot_progress() reading it from
whatever thread calls DownloadHandle::progress() broke that invariant
(workers_mu's shared_lock protects the workers map's structure, not an
individual SegWorker's fields). Now std::atomic<double> with relaxed
ordering -- an informational EMA, nothing synchronizes real state on it --
rather than adding a lock to the write side.

core/tests/segment/budget_test.cpp adds two tests reproducing the actual
gap (budget_active_never_exceeds_max_active_segments_under_concurrent_load,
budget_wait_list_wakes_a_task_whose_target_never_changed) plus a sanity
baseline (budget_release_wakes_a_denied_waiter), and introduces AsyncFakeTask
+ TestTimer for the one existing test that drives the budget from multiple
concurrent threads -- mirroring production's real dispatch (register_task()'s
on_target lambda posts through host.schedule(), download_task.cpp, never a
synchronous call) rather than adding reentrancy-guarding machinery to
SegmentBudget itself to compensate for a synchronous test double being
unlike production. See docs/adr/0017 for the full writeup, including what an
earlier version of this fix got wrong chasing a same-thread reentrancy
hazard that doesn't actually exist in production.

core/docs/m7-baseline.md updated: the RSS number now clears the DoD line
(45.41 MiB via heap-profile), root-caused rather than just re-measured.

Co-Authored-By: Claude Sonnet 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01Q3QrF7rCt21bkAjt9BCDFQ
This commit is contained in:
2026-09-12 10:58:14 +04:00
co-authored by Claude Sonnet 5
parent 7ea04fa79c
commit 322a20efa5
6 changed files with 536 additions and 69 deletions
+50 -6
View File
@@ -37,12 +37,10 @@ SegmentBudget::EngineBudget SegmentBudget::snapshot_locked() const {
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").
// DAEMON's list first, then any registered task not in it (defensive; "a running task
// absent from the list sorts last"). Shared by reallocate_locked() and
// wake_one_waiter_locked(), which need the same priority ordering.
std::vector<TaskId> SegmentBudget::priority_order_locked() const {
std::vector<TaskId> order;
order.reserve(tasks_.size());
for (TaskId id : order_)
@@ -51,6 +49,14 @@ SegmentBudget::Plan SegmentBudget::reallocate_locked() {
for (const auto &[id, _] : tasks_)
if (std::find(order.begin(), order.end(), id) == order.end())
order.push_back(id);
return order;
}
// 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() {
std::vector<TaskId> order = priority_order_locked();
std::unordered_map<TaskId, std::uint32_t> target;
target.reserve(order.size());
@@ -155,6 +161,7 @@ void SegmentBudget::deregister_task(TaskId id) {
active_ -= it->second.held;
tasks_.erase(it);
plan = reallocate_locked();
wake_one_waiter_locked(plan, id); // a departing task frees real slots too
}
run(plan);
}
@@ -182,6 +189,19 @@ bool SegmentBudget::confirm_slot(TaskId id) {
Task &t = it->second;
if (t.held >= t.target)
return false; // target was cut in the race
// reallocate_locked()'s pool math bounds sum(target) <= max_active_ *as computed*, but
// that doesn't bound sum(held): a task can be legitimately over its own just-lowered
// target for a while (yield deferred to a segment boundary, ADR 0011 A1), and another
// task's target can correctly rise to claim that capacity before the first task has
// physically released it. active_ is the one number that's always true regardless of
// any task's target bookkeeping, so it's the backstop. Denials here are remembered
// (waiting_for_slot) rather than left for the caller to somehow ask again at the right
// moment -- see release_slot()'s wake_one_waiter_locked() call.
if (active_ >= max_active_) {
t.waiting_for_slot = true;
return false;
}
t.waiting_for_slot = false;
++t.held;
++active_;
if (snapshot_locked() != last_notified_) {
@@ -191,6 +211,25 @@ bool SegmentBudget::confirm_slot(TaskId id) {
return true;
}
// Appends a retry hint for the highest-priority task with waiting_for_slot set (other
// than `exclude`) to `plan`. Only ever wakes a task confirm_slot() actually turned away --
// not just anyone below its target, which would also fire for tasks that are fairly,
// correctly not entitled to more right now (see reallocate_locked()'s own target math).
void SegmentBudget::wake_one_waiter_locked(Plan &plan, TaskId exclude) const {
for (TaskId id : priority_order_locked()) {
if (id == exclude)
continue;
auto it = tasks_.find(id);
if (it == tasks_.end())
continue;
const Task &t = it->second;
if (t.waiting_for_slot && t.on_target) {
plan.targets.emplace_back(t.on_target, t.target);
return; // exactly one freed slot, exactly one retry hint
}
}
}
void SegmentBudget::release_slot(TaskId id) {
Plan plan;
{
@@ -201,6 +240,11 @@ void SegmentBudget::release_slot(TaskId id) {
--it->second.held;
--active_;
plan = reallocate_locked();
// reallocate_locked() only fires a callback for a task whose *target* changed.
// The task this freed slot is actually owed to (see confirm_slot()) may have a
// target that was already correct and hasn't moved -- nothing else will ever ask
// it to retry, so the budget has to remember and hand this slot to it directly.
wake_one_waiter_locked(plan, id);
}
run(plan);
}
+56 -7
View File
@@ -96,7 +96,17 @@ struct SegWorker {
SteadyTime sample_at{};
std::uint64_t sample_bytes = 0;
double speed_bps = 0;
// Written only by this segment's own curl callback (seg_data, sequential -- no lock
// held across the update, by design: the transfer hot path takes no lock it doesn't
// need). snapshot_progress() reads it from whatever thread calls
// DownloadHandle::progress() (DAEMON's polling, or anyone else's), which workers_mu's
// shared_lock does NOT cover -- that lock only protects the `workers` map's own
// structure, not an individual SegWorker's mutable fields. atomic<double> (relaxed:
// this is an approximate, informational EMA, not something anything synchronizes
// real state on) keeps that read-from-any-thread safe without adding a lock to the
// write side. Found by TSan the first time anything actually read this cross-thread
// (engine_polled_progress_reports_nonzero_speed, added alongside the speed_bps fix).
std::atomic<double> speed_bps{0};
};
struct DownloadTaskState : std::enable_shared_from_this<DownloadTaskState> {
@@ -211,6 +221,11 @@ struct DownloadTaskState : std::enable_shared_from_this<DownloadTaskState> {
void finish_probe_locked();
void apply_slot_target(std::uint32_t n);
void fill_slots_locked();
// Confirms a budget slot and starts a worker for `seg_idx` if nothing already covers
// it. false on either failure (already running, or budget denied) -- the caller's own
// fallback (fill_slots_locked's assign_slot() loop, retry_worker()'s early return)
// stays the same either way. Caller holds mu.
bool try_start_segment_locked(std::uint32_t seg_idx);
void start_worker_locked(std::uint32_t seg_idx);
void restart_probe(bool with_auth);
@@ -368,6 +383,15 @@ void DownloadTaskState::finish_probe_locked() {
host.budget().set_want(id, want_slots());
}
bool DownloadTaskState::try_start_segment_locked(std::uint32_t seg_idx) {
if (workers.count(seg_idx))
return false;
if (!host.budget().confirm_slot(id))
return false;
start_worker_locked(seg_idx);
return true;
}
// 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
@@ -375,6 +399,23 @@ void DownloadTaskState::finish_probe_locked() {
// 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() {
// A segment mid-backoff (SegState::stalled -- seg_finished's retry path: released its
// slot, scheduled a retry_worker() timer, and gave up quietly if confirm_slot() denied
// it then) has no live worker and never surfaces through assign_slot() -- it stays
// assigned to whichever segment iteration created it, just not running. It's not
// "fresh work" the loop below would ever find on its own. Every path that can free
// budget capacity ends up here (apply_slot_target(), driven by SegmentBudget's
// target-changed callback *and* its wait-list wakeup for a task whose target didn't
// move -- see confirm_slot()/release_slot() in budget.cpp), so this is the one place
// that needs to give a stalled segment another try, not every caller of
// retry_worker(). Bounded by slot_target like the assign_slot() loop below; a segment
// this doesn't get to keeps its own scheduled retry_worker() timer as a second chance.
for (auto &v : seg->snapshot()) {
if (workers.size() >= slot_target)
break;
if (v.state == segment::SegState::stalled)
try_start_segment_locked(v.index);
}
while (workers.size() < slot_target) {
auto s = seg->assign_slot();
if (!s) {
@@ -387,7 +428,11 @@ void DownloadTaskState::fill_slots_locked() {
}
start_worker_locked(*s);
}
if (state == EngineState::connecting && !workers.empty())
// Mirrors retry_worker()'s own transition: a stalled-segment restart above can be the
// thing that takes a retry_wait task back to actually transferring, same as connecting
// does for a task starting up.
if ((state == EngineState::connecting || state == EngineState::retry_wait) &&
!workers.empty())
transition(EngineState::downloading, std::nullopt);
}
@@ -547,7 +592,9 @@ net::DataAction DownloadTaskState::seg_data(std::uint32_t seg_idx, ConstByteSpan
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;
double prev = w->speed_bps.load(std::memory_order_relaxed);
w->speed_bps.store(prev == 0 ? inst : 0.7 * prev + 0.3 * inst,
std::memory_order_relaxed);
w->sample_at = now;
w->sample_bytes = w->recv;
}
@@ -1013,10 +1060,11 @@ void DownloadTaskState::emit_progress_if_due() {
std::shared_lock lk(workers_mu);
double agg = 0;
for (auto &[idx, w] : workers) {
agg += w->speed_bps;
double speed = w->speed_bps.load(std::memory_order_relaxed);
agg += speed;
SegmentProgress sp;
sp.index = idx;
sp.speed_bps = static_cast<std::uint64_t>(w->speed_bps);
sp.speed_bps = static_cast<std::uint64_t>(speed);
p.segments.push_back(sp);
}
p.speed_bps = static_cast<std::uint64_t>(agg);
@@ -1230,8 +1278,9 @@ Progress DownloadTaskState::snapshot_progress() {
std::shared_lock wl(workers_mu);
worker_speed.reserve(workers.size());
for (auto &[idx, w] : workers) {
worker_speed.emplace(idx, w->speed_bps);
agg_speed += w->speed_bps;
double speed = w->speed_bps.load(std::memory_order_relaxed);
worker_speed.emplace(idx, speed);
agg_speed += speed;
}
p.effective_segments = static_cast<std::uint32_t>(workers.size());
}