// tools/bench/vdm_bench.cpp — the M1/M7 performance gates from docs/04-engine-design.md §8,
// and the sanitizer-clean concurrent-load regression that falls out of the same harness.
//
// Three subcommands, one binary, one way of driving vdm::Engine and reading back
// wall-clock/CPU/RSS:
//
//   vdm_bench throughput [--size 5G] [--segments N] [--require-mbps 125] [--max-cpu-pct 8]
//       A single download. Reports achieved throughput and process CPU as a percentage of
//       one core. The M1 DoD line is "5 GB saturates a 1 Gbit link at <=8% of one core" --
//       pass --size 5G --require-mbps 940 --max-cpu-pct 8 against a real link for the
//       actual sign-off. Against the bundled local server (see support/local_server.hpp)
//       the numbers are still meaningful for regression tracking; they just aren't a
//       real-network measurement, which is why the ctest-registered run below doesn't gate
//       on them.
//
//   vdm_bench load [--tasks 20] [--task-size 4M] [--segments N] [--require-rss-kb N]
//       docs/04 §8's "<=60 MB RSS with 20 active downloads at default buffers, given
//       max_active_segments=32" scenario: one Engine, default Config, N concurrent tasks,
//       peak RSS read back via getrusage(). This is the same binary the ctest below runs
//       under ASan/UBSan/TSan as the M1 DoD's "20-task load test" -- there the job is
//       purely "no sanitizer error, every task completes correctly"; --require-rss-kb is
//       for a --preset release run, where the number means something (a sanitizer roughly
//       doubles-to-quadruples RSS via redzones/shadow memory).
//
//   vdm_bench heap-profile [--tasks 20] [--task-size 4M] [--segments N] [--top N]
//       Answers "where does the RSS actually go", for the M7 RSS target (docs/04 §8) --
//       no massif/heaptrack in this environment (no root to install them), so this is a
//       small in-process equivalent: the same operator new/delete override tracks, per
//       allocation, which call site made it (one return address, via
//       __builtin_return_address(0) right in operator new -- cheap enough to run for the
//       whole scenario, unlike a full backtrace) and keeps a live (not-yet-freed) byte
//       count per site. Runs the same 20-task concurrent scenario as `load`, waits for
//       peak RSS to stop growing, then prints the top sites by live bytes -- the
//       breakdown ADR 0012's "45-50 MB" estimate needs to be checked against, and
//       core/docs/m7-baseline.md's ~70 MB measurement needs explained.
//
//   vdm_bench alloc-check [--size 128M] [--window-s 2]
//       docs/agents/AGENT-CORE.md: "no allocation in the curl write callback... checked in
//       review and by a bench assertion." operator new/delete are overridden process-wide
//       below; this samples the count across a steady mid-transfer window (no segment
//       start/stop, so no probe/segmenter/sidecar activity) and requires it stay within a
//       small time-proportional budget -- not a strict zero, because emit_progress_if_due()
//       legitimately builds a Progress::segments vector up to 4x/sec regardless of
//       throughput. What must NOT happen is that count scaling with bytes transferred; the
//       budget is sized so it can't, while tolerating that fixed, small, rate-independent
//       bookkeeping cost. The transfer is paced by tools/testserver's `throttled` mode
//       (support/testserver_client.hpp), not the engine's own rate limiter -- the
//       limiter's pause/resume path allocates on every throttle event, which would measure
//       the limiter instead of the write path.
//
// Every subcommand is report-only (exit 0 once the download(s) complete correctly) unless
// its --require-* flag is passed, so the ctest registrations in CMakeLists.txt are stable
// under CI's shared, sanitizer-slowed, virtualized hardware.

#include "vdm/engine.hpp"

#include <cxxabi.h>
#include <dlfcn.h>
#include <sys/resource.h>
#include <unistd.h>

#include <algorithm>
#include <atomic>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <future>
#include <mutex>
#include <new>
#include <string>
#include <unordered_map>
#include <vector>

#include "support/local_server.hpp"
#include "support/testserver_client.hpp"

using namespace vdm;
using namespace vdm::task;
using namespace std::chrono_literals;

// --- allocation tracking: a plain counter for alloc-check, and (only when heap-profile
// turns it on) a per-call-site live-byte tracker standing in for massif/heaptrack, neither
// of which is installable here (no root). Both share one operator new/delete override
// since a process gets exactly one. ---------------------------------------------------

namespace {
std::atomic<std::uint64_t> g_alloc_count{0};

std::atomic<bool> g_site_tracking{false};
// Re-entrancy guard: our own bookkeeping below calls into std::unordered_map, which
// allocates. Without this, that nested operator new call would recurse into the same
// tracking logic -- not infinitely (the nested call's own bookkeeping call would itself be
// guarded the same way, so it terminates), but it would count the tracker's own node
// allocations as if they were the program's, polluting the profile. When set, an
// allocation is just satisfied via malloc/free with no bookkeeping.
thread_local bool g_in_tracker = false;

struct LiveAlloc {
    std::size_t size;
    void *site;  // __builtin_return_address(0) at the call to operator new
};
struct SiteStats {
    std::uint64_t live_bytes = 0;
    std::uint64_t alloc_count = 0;  // lifetime, not just live -- how *often* a site fires
};

std::mutex g_track_mu;
std::unordered_map<void *, LiveAlloc> g_live;         // pointer -> what/where
std::unordered_map<void *, SiteStats> g_site_stats;    // call site -> aggregate

void track_alloc(void *p, std::size_t n) {
    if (!g_site_tracking.load(std::memory_order_relaxed) || g_in_tracker)
        return;
    g_in_tracker = true;
    void *site = __builtin_return_address(0);
    {
        std::lock_guard lk(g_track_mu);
        g_live.emplace(p, LiveAlloc{n, site});
        auto &st = g_site_stats[site];
        st.live_bytes += n;
        st.alloc_count += 1;
    }
    g_in_tracker = false;
}

void track_free(void *p) {
    if (!p || !g_site_tracking.load(std::memory_order_relaxed) || g_in_tracker)
        return;
    g_in_tracker = true;
    {
        std::lock_guard lk(g_track_mu);
        if (auto it = g_live.find(p); it != g_live.end()) {
            g_site_stats[it->second.site].live_bytes -= it->second.size;
            g_live.erase(it);
        }
    }
    g_in_tracker = false;
}

std::string describe_site(void *site);  // defined below; used as a per-site fallback here

// Resolves a batch of call-site addresses to "function (file:line)" via one `addr2line`
// invocation (reads .symtab + DWARF, unlike dladdr below -- most of this codebase's
// call sites are internal-linkage or file-static and never make it into the *dynamic*
// symbol table dladdr is limited to). `sites` and the returned vector are in the same
// order. dladdr's `dli_fbase` (the object's runtime load base -- works even when it can't
// name the specific symbol, which is exactly the "site not in .dynsym" case this exists
// for) converts each ASLR'd runtime address into the file-relative offset addr2line wants.
// Falls back to a raw "0x..." string per site if addr2line isn't on PATH or the batch call
// otherwise fails -- describe_site() below still has dladdr as a second-line fallback for
// an individual site in that case.
std::vector<std::string> resolve_sites_via_addr2line(const std::vector<void *> &sites) {
    std::vector<std::string> out(sites.size());
    for (std::size_t i = 0; i < sites.size(); ++i) {
        char buf[32];
        std::snprintf(buf, sizeof buf, "%p", sites[i]);
        out[i] = buf;
    }
    if (sites.empty())
        return out;

    // `-e /proc/self/exe` would resolve inside the *addr2line* process this spawns, i.e.
    // to addr2line's own binary, not vdm_bench's -- readlink it here, in-process, first.
    char exe_path[4096];
    ssize_t exe_len = ::readlink("/proc/self/exe", exe_path, sizeof(exe_path) - 1);
    if (exe_len <= 0)
        return out;
    exe_path[exe_len] = '\0';

    std::string cmd = std::string("addr2line -f -C -e '") + exe_path + "'";
    for (void *site : sites) {
        Dl_info info{};
        std::uintptr_t off = reinterpret_cast<std::uintptr_t>(site);
        if (::dladdr(site, &info) && info.dli_fbase)
            off -= reinterpret_cast<std::uintptr_t>(info.dli_fbase);
        char buf[24];
        std::snprintf(buf, sizeof buf, " %zx", off);
        cmd += buf;
    }
    cmd += " 2>/dev/null";

    FILE *pipe = ::popen(cmd.c_str(), "r");
    if (!pipe)
        return out;
    std::vector<std::string> lines;
    char line[1024];
    while (std::fgets(line, sizeof line, pipe)) {
        std::string s = line;
        while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back();
        lines.push_back(std::move(s));
    }
    ::pclose(pipe);

    // addr2line -f prints two lines per address: function name, then file:line.
    for (std::size_t i = 0; i * 2 + 1 < lines.size() && i < sites.size(); ++i) {
        const std::string &fn = lines[i * 2];
        const std::string &loc = lines[i * 2 + 1];
        if (fn == "??" && loc == "??:0")
            out[i] = describe_site(sites[i]);  // addr2line drew a blank -- try dladdr
        else
            out[i] = fn + " (" + loc + ")";
    }
    return out;
}

// Single-site fallback (dladdr only) for when addr2line can't place an address (it read
// only .symtab -- e.g. no debug/symbol data at all for that address) or isn't available.
std::string describe_site(void *site) {
    Dl_info info{};
    if (!::dladdr(site, &info) || !info.dli_sname) {
        char buf[32];
        std::snprintf(buf, sizeof buf, "%p", site);
        return buf;
    }
    std::string name = info.dli_sname;
    int status = 0;
    if (char *dem = abi::__cxa_demangle(name.c_str(), nullptr, nullptr, &status);
        dem && status == 0) {
        name = dem;
        std::free(dem);
    }
    auto off = static_cast<std::uintptr_t>(reinterpret_cast<char *>(site) -
                                           reinterpret_cast<char *>(info.dli_saddr));
    return name + " +0x" + [&] {
        char buf[24];
        std::snprintf(buf, sizeof buf, "%zx", off);
        return std::string(buf);
    }();
}

}  // namespace

void *operator new(std::size_t n) {
    g_alloc_count.fetch_add(1, std::memory_order_relaxed);
    void *p = std::malloc(n ? n : 1);
    if (!p)
        throw std::bad_alloc();
    track_alloc(p, n);
    return p;
}
// The nothrow overload is a distinct allocation function from the plain one above, not an
// alternate spelling of it: leaving it un-overridden means it falls through to ASan's own
// default (nothrow) operator new, while a deallocation of that same pointer still reaches
// the plain `operator delete` override below -- a real alloc-dealloc-mismatch ASan catches
// (hit via std::stable_sort's std::get_temporary_buffer, which allocates this way; see
// core/src/segment/budget.cpp's SegmentBudget::run() for why that path no longer uses
// stable_sort at all, but this override is fixed too since anything else pulling in a
// nothrow allocation -- directly or via another std:: temp-buffer user -- would hit it
// again otherwise).
void *operator new(std::size_t n, const std::nothrow_t &) noexcept {
    g_alloc_count.fetch_add(1, std::memory_order_relaxed);
    void *p = std::malloc(n ? n : 1);
    if (p)
        track_alloc(p, n);
    return p;
}
void operator delete(void *p) noexcept {
    track_free(p);
    std::free(p);
}
void operator delete(void *p, std::size_t) noexcept {
    track_free(p);
    std::free(p);
}
void operator delete(void *p, const std::nothrow_t &) noexcept {
    track_free(p);
    std::free(p);
}

namespace {

// --- small helpers ---------------------------------------------------------------------

std::uint64_t parse_size(std::string_view s) {
    if (s.empty())
        return 0;
    char suffix = s.back();
    std::uint64_t mult = 1;
    std::string_view digits = s;
    if (suffix == 'k' || suffix == 'K') {
        mult = 1024;
        digits.remove_suffix(1);
    } else if (suffix == 'm' || suffix == 'M') {
        mult = 1024ull * 1024;
        digits.remove_suffix(1);
    } else if (suffix == 'g' || suffix == 'G') {
        mult = 1024ull * 1024 * 1024;
        digits.remove_suffix(1);
    }
    return std::strtoull(std::string(digits).c_str(), nullptr, 10) * mult;
}

// Trivial `--flag value` / `--flag` (bool) parser: no library dependency worth adding for
// a handful of options across three subcommands.
class Args {
  public:
    Args(int argc, char **argv, int start) {
        for (int i = start; i < argc; ++i) raw_.emplace_back(argv[i]);
    }
    [[nodiscard]] std::string get(std::string_view flag, std::string def) const {
        for (std::size_t i = 0; i < raw_.size(); ++i)
            if (raw_[i] == flag && i + 1 < raw_.size())
                return raw_[i + 1];
        return def;
    }
    [[nodiscard]] std::uint64_t get_size(std::string_view flag, std::string def) const {
        return parse_size(get(flag, std::move(def)));
    }
    [[nodiscard]] double get_double(std::string_view flag, double def) const {
        auto s = get(flag, "");
        return s.empty() ? def : std::strtod(s.c_str(), nullptr);
    }
    [[nodiscard]] long get_long(std::string_view flag, long def) const {
        auto s = get(flag, "");
        return s.empty() ? def : std::strtol(s.c_str(), nullptr, 10);
    }

  private:
    std::vector<std::string> raw_;
};

struct Rusage {
    double cpu_s;
    long peak_rss_kb;
};

Rusage sample_rusage() {
    struct ::rusage ru {};
    ::getrusage(RUSAGE_SELF, &ru);
    double cpu = (double)ru.ru_utime.tv_sec + ru.ru_utime.tv_usec / 1e6 +
                (double)ru.ru_stime.tv_sec + ru.ru_stime.tv_usec / 1e6;
    return {cpu, ru.ru_maxrss};  // ru_maxrss is KiB on Linux, and is a lifetime peak, not
                                 // a snapshot -- fine for us, we only ever want the peak.
}

// Synchronous single-download driver: start it, block for on_finished. Used by throughput
// and alloc-check, which only ever run one transfer at a time.
Result<DownloadOutcome> run_one(Engine &eng, DownloadSpec spec, std::chrono::seconds timeout) {
    std::promise<Result<DownloadOutcome>> p;
    auto f = p.get_future();
    std::atomic<bool> fired{false};
    DownloadCallbacks cbs;
    cbs.on_finished = [&](Result<DownloadOutcome> r) {
        if (!fired.exchange(true))
            p.set_value(std::move(r));
    };
    auto h = eng.start(std::move(spec), std::move(cbs));
    if (f.wait_for(timeout) != std::future_status::ready)
        return Err{Error::timeout, "vdm_bench: download did not finish in time"};
    return f.get();
}

std::string tmp_workdir() {
    std::string p = "/tmp/vdm_bench_XXXXXX";
    return ::mkdtemp(p.data()) ? p : "/tmp";
}

void report(const char *label, double v, const char *unit) {
    std::fprintf(stderr, "  %-22s %10.2f %s\n", label, v, unit);
}

// --- subcommands -------------------------------------------------------------------

int cmd_throughput(const Args &a) {
    const std::uint64_t size = a.get_size("--size", "256M");
    const long require_mbps = a.get_long("--require-mbps", 0);
    const long max_cpu_pct = a.get_long("--max-cpu-pct", 0);
    const auto segments = static_cast<std::uint32_t>(a.get_long("--segments", 0));

    std::string dir = tmp_workdir();
    if (!vdm::bench::make_sparse_file(dir + "/payload.bin", size)) {
        std::fprintf(stderr, "vdm_bench: could not create %llu-byte payload in %s\n",
                     (unsigned long long)size, dir.c_str());
        return 2;
    }
    vdm::bench::LocalServer srv(dir);
    if (!srv.available()) {
        std::fprintf(stderr,
                     "vdm_bench: no local server available (busybox missing?) -- skipping "
                     "throughput bench.\n");
        return 0;  // not a failure of the engine; nothing to measure against
    }

    Engine eng;
    DownloadSpec spec;
    spec.url = srv.url("/payload.bin");
    spec.save_path = dir + "/out.bin";
    if (segments)
        spec.segments = segments;

    const auto cpu0 = sample_rusage().cpu_s;
    const auto t0 = std::chrono::steady_clock::now();
    auto r = run_one(eng, std::move(spec), 300s);
    const auto t1 = std::chrono::steady_clock::now();
    const auto ru1 = sample_rusage();

    if (!r.has_value()) {
        std::fprintf(stderr, "vdm_bench throughput: download failed: %s\n",
                     r.error().to_string().c_str());
        return 1;
    }

    const double wall_s = std::chrono::duration<double>(t1 - t0).count();
    const double cpu_s = ru1.cpu_s - cpu0;
    const double mbps = (r.value().bytes * 8.0 / 1'000'000.0) / wall_s;
    const double cpu_pct = wall_s > 0 ? (cpu_s / wall_s) * 100.0 : 0.0;

    std::fprintf(stderr, "throughput: %llu bytes in %.2fs\n", (unsigned long long)r.value().bytes,
                wall_s);
    report("throughput", mbps, "Mbps");
    report("cpu", cpu_pct, "% of one core");
    report("peak RSS", ru1.peak_rss_kb / 1024.0, "MiB");

    int rc = 0;
    if (require_mbps > 0 && mbps < require_mbps) {
        std::fprintf(stderr, "FAIL: %.2f Mbps < required %ld Mbps\n", mbps, require_mbps);
        rc = 1;
    }
    if (max_cpu_pct > 0 && cpu_pct > max_cpu_pct) {
        std::fprintf(stderr, "FAIL: %.2f%% CPU > allowed %ld%%\n", cpu_pct, max_cpu_pct);
        rc = 1;
    }
    return rc;
}

int cmd_load(const Args &a) {
    const int tasks = static_cast<int>(a.get_long("--tasks", 20));
    const std::uint64_t task_size = a.get_size("--task-size", "4M");
    const long require_rss_kb = a.get_long("--require-rss-kb", 0);
    // 0 => engine default (default_segments=8, docs/04 §8's actual DoD scenario). The
    // ctest-registered smoke run overrides this down -- see CMakeLists.txt for why.
    const auto segments_override = static_cast<std::uint32_t>(a.get_long("--segments", 0));

    std::string dir = tmp_workdir();

    // docs/04 §8's scenario is default_segments=8 per task, so up to tasks*8 concurrent
    // segments (160, at the default --tasks 20) -- that's the point: the RSS ceiling is
    // about *concurrent* segment buffers under a realistic multi-segment spread, not one
    // connection per task. busybox httpd (LocalServer, used for the throughput bench)
    // could not sustain that many concurrent connections reliably: reproducible hangs past
    // a 120s per-task wait at --tasks 20 --task-size 4M, though not always at 2M -- some
    // connections simply never got serviced. tools/testserver's threaded server (already
    // exercised at real concurrency by the hostile-mode suite in engine_test.cpp) doesn't
    // have that ceiling, so it's the transport here despite being the slower-per-request
    // choice noted in support/local_server.hpp -- for this bench "slower" is actually
    // wanted anyway (see below).
    //
    // Pace via testserver's own `throttled` mode rather than the engine's
    // rate::RateLimiter: on loopback even 160 segments would otherwise race to completion
    // before there's any concurrent overlap to measure RSS against, and pacing externally
    // avoids a separate, real finding -- rate::RateLimiter::set_global_limit() under this
    // much segment contention was observed to starve a couple of tasks for 120s+ instead of
    // completing in the few seconds the rate implies (single shared TokenBucket, no
    // fairness ordering across peek/commit races -- see docs/adr/0016). Throttle per
    // connection, not per task: each segment is its own connection, so divide the
    // per-task rate across default_segments to land total task duration in the same
    // ballpark regardless of how many segments the engine actually opens.
    const std::uint64_t assumed_segments = segments_override ? segments_override : 8;
    const std::uint64_t per_conn_bps =
        std::max<std::uint64_t>(1, task_size / (5 * assumed_segments));  // ~5s per task
    vdm::bench::TestServerProc srv(per_conn_bps);
    if (!srv.available()) {
        std::fprintf(stderr,
                     "vdm_bench: tools/testserver unavailable -- skipping load bench.\n");
        return 0;
    }

    Engine eng;  // default Config: default_segments=8, max_active_segments=32 (docs/04 §8)

    std::vector<std::promise<Result<DownloadOutcome>>> proms(tasks);
    std::vector<std::future<Result<DownloadOutcome>>> futs;
    std::vector<std::atomic<bool>> fired(tasks);
    futs.reserve(tasks);
    for (auto &p : proms) futs.push_back(p.get_future());

    std::vector<DownloadHandle> handles;
    handles.reserve(tasks);
    const auto t0 = std::chrono::steady_clock::now();
    for (int i = 0; i < tasks; ++i) {
        DownloadSpec spec;
        spec.url = srv.url("/throttled/file/" + std::to_string(task_size));
        spec.save_path = dir + "/out" + std::to_string(i) + ".bin";
        if (segments_override)
            spec.segments = segments_override;
        DownloadCallbacks cbs;
        cbs.on_finished = [&proms, &fired, i](Result<DownloadOutcome> r) {
            if (!fired[i].exchange(true))
                proms[i].set_value(std::move(r));
        };
        handles.push_back(eng.start(std::move(spec), std::move(cbs)));
    }

    // TSan's per-access instrumentation overhead is heavy enough (observed: a couple of
    // stragglers past 120s at --tasks 20 --task-size 2M, no TSan report -- just slow, not
    // stuck) that a tight per-task budget here isn't testing the engine, it's testing the
    // sanitizer. 300s per straggler is still bounded, just generous enough that "slow under
    // instrumentation" and "actually wedged" stay distinguishable.
    const auto task_timeout = std::chrono::seconds(a.get_long("--task-timeout-s", 300));
    int failures = 0;
    for (int i = 0; i < tasks; ++i) {
        if (futs[i].wait_for(task_timeout) != std::future_status::ready) {
            std::fprintf(stderr, "task %d: timed out\n", i);
            ++failures;
            continue;
        }
        auto r = futs[i].get();
        if (!r.has_value()) {
            std::fprintf(stderr, "task %d: %s\n", i, r.error().to_string().c_str());
            ++failures;
        }
    }
    const auto t1 = std::chrono::steady_clock::now();
    const auto ru = sample_rusage();

    std::fprintf(stderr, "load: %d tasks, %d failed, %.2fs wall\n", tasks, failures,
                std::chrono::duration<double>(t1 - t0).count());
    report("peak RSS", ru.peak_rss_kb / 1024.0, "MiB");

    if (failures > 0)
        return 1;
    if (require_rss_kb > 0 && ru.peak_rss_kb > require_rss_kb) {
        std::fprintf(stderr, "FAIL: peak RSS %ld KiB > allowed %ld KiB\n", ru.peak_rss_kb,
                     require_rss_kb);
        return 1;
    }
    return 0;
}

int cmd_heap_profile(const Args &a) {
    const int tasks = static_cast<int>(a.get_long("--tasks", 20));
    const std::uint64_t task_size = a.get_size("--task-size", "4M");
    const auto segments_override = static_cast<std::uint32_t>(a.get_long("--segments", 0));
    const int top_n = static_cast<int>(a.get_long("--top", 20));

    std::string dir = tmp_workdir();
    const std::uint64_t assumed_segments = segments_override ? segments_override : 8;
    const std::uint64_t per_conn_bps =
        std::max<std::uint64_t>(1, task_size / (5 * assumed_segments));
    vdm::bench::TestServerProc srv(per_conn_bps);
    if (!srv.available()) {
        std::fprintf(stderr,
                     "vdm_bench: tools/testserver unavailable -- skipping heap-profile.\n");
        return 0;
    }

    // On from before Engine's own construction, so its fixed cost (http worker threads,
    // curl_multi handles, ...) shows up in the profile too, not just what the tasks add.
    g_site_tracking.store(true, std::memory_order_relaxed);
    Engine eng;

    std::vector<DownloadHandle> handles;
    handles.reserve(tasks);
    for (int i = 0; i < tasks; ++i) {
        DownloadSpec spec;
        spec.url = srv.url("/throttled/file/" + std::to_string(task_size));
        spec.save_path = dir + "/out" + std::to_string(i) + ".bin";
        if (segments_override)
            spec.segments = segments_override;
        handles.push_back(eng.start(std::move(spec), DownloadCallbacks{}));
    }

    // Wait for peak RSS to stop growing (5 consecutive 200ms samples with no increase, or
    // a generous cap) -- that's the moment whose live-allocation breakdown answers "where
    // does the RSS actually go", the same question core/docs/m7-baseline.md's ~70 MB
    // number raises against ADR 0012's ~45-50 MB estimate.
    long last_rss = 0, stable_for = 0;
    for (int i = 0; i < 300 && stable_for < 5; ++i) {
        std::this_thread::sleep_for(200ms);
        long rss = sample_rusage().peak_rss_kb;
        stable_for = (rss <= last_rss) ? stable_for + 1 : 0;
        last_rss = rss;
    }

    // Cross-check against the budget's own accounting, independent of the allocator-site
    // snapshot below: this is what caught the real find here (core/docs/m7-baseline.md,
    // docs/adr/0012) -- sum(effective_segments) tracked budget.active well past
    // budget.total before the SegmentBudget fix these numbers are now reported alongside.
    std::uint32_t total_effective_segments = 0;
    for (auto &h : handles) total_effective_segments += h.progress().effective_segments;
    auto eb = eng.segment_budget().budget();
    std::fprintf(stderr,
                "segment budget: %u live across tasks, engine reports total=%u active=%u "
                "starved=%u\n",
                total_effective_segments, eb.total, eb.active, eb.tasks_starved);

    // Turn tracking off *before* touching g_site_stats: the snapshot below allocates its
    // own vector storage through the very same overridden operator new. With tracking
    // still on and g_track_mu already held for the snapshot, that allocation would recurse
    // into track_alloc() and try to lock g_track_mu again -- self-deadlock on a
    // non-recursive mutex (hit this exact hang while first wiring this subcommand up:
    // g_in_tracker only guards the tracker's own bookkeeping containers, not other code
    // that allocates while holding g_track_mu directly).
    g_site_tracking.store(false, std::memory_order_relaxed);
    std::vector<std::pair<void *, SiteStats>> sites;
    {
        std::lock_guard lk(g_track_mu);
        sites.assign(g_site_stats.begin(), g_site_stats.end());
    }
    std::sort(sites.begin(), sites.end(), [](const auto &x, const auto &y) {
        return x.second.live_bytes > y.second.live_bytes;
    });

    std::fprintf(stderr, "heap-profile: peak RSS %.2f MiB, %d tasks, %llu segments assumed\n",
                last_rss / 1024.0, tasks, (unsigned long long)assumed_segments);
    std::fprintf(stderr, "%-10s %-10s  %s\n", "live KiB", "#allocs", "call site");
    std::uint64_t total_live = 0;
    for (auto &[site, st] : sites) total_live += st.live_bytes;
    std::fprintf(stderr, "total tracked live: %.2f MiB across %zu sites\n\n",
                total_live / 1024.0 / 1024.0, sites.size());
    std::vector<void *> top_addrs;
    for (int i = 0; i < top_n && i < static_cast<int>(sites.size()); ++i) top_addrs.push_back(sites[i].first);
    auto labels = resolve_sites_via_addr2line(top_addrs);
    for (std::size_t i = 0; i < top_addrs.size(); ++i) {
        auto &[site, st] = sites[i];
        std::fprintf(stderr, "%-10.1f %-10llu  %s\n", st.live_bytes / 1024.0,
                    (unsigned long long)st.alloc_count, labels[i].c_str());
    }

    // Handles/Engine go out of scope here: ~Engine quiesces every task (waits for its
    // workers to actually drain -- see the quiesce() fix) rather than waiting for the
    // downloads to run to completion, which isn't the point of this subcommand.
    return 0;
}

int cmd_alloc_check(const Args &a) {
    const std::uint64_t size = a.get_size("--size", "128M");
    const double window_s = a.get_double("--window-s", 2.0);
    const double budget_per_s = a.get_double("--budget-per-s", 5.0);

    std::string dir = tmp_workdir();

    // busybox httpd over loopback is fast enough that even a --size in the hundreds of MB
    // can complete in well under --window-s (measured: 64M in ~0.06s) -- there'd be no
    // steady-state middle to sample. Pace the transfer with tools/testserver's `throttled`
    // mode instead of the engine's own rate::RateLimiter: the limiter's precision
    // pause/resume path allocates a timer node + std::function on every throttle event (see
    // support/testserver_client.hpp), which would be exactly the kind of cost this bench
    // exists to catch -- pacing external to the engine keeps the sample honest.
    const std::uint64_t target_bps = std::max<std::uint64_t>(1, size / std::max(1.0, window_s * 6));
    vdm::bench::TestServerProc srv(target_bps);
    if (!srv.available()) {
        std::fprintf(stderr,
                     "vdm_bench: tools/testserver unavailable -- skipping alloc-check.\n");
        return 0;
    }

    Engine eng;
    DownloadSpec spec;
    spec.url = srv.url("/throttled/file/" + std::to_string(size));
    spec.save_path = dir + "/out.bin";
    // testserver throttles each connection independently, not the aggregate -- with the
    // default multi-segment split the N parallel connections would finish in ~1/N of the
    // time target_bps was sized for. Pin to one segment so the pacing math above holds.
    spec.segments = 1;

    std::promise<Result<DownloadOutcome>> p;
    auto f = p.get_future();
    std::atomic<bool> fired{false};
    DownloadCallbacks cbs;
    cbs.on_finished = [&](Result<DownloadOutcome> r) {
        if (!fired.exchange(true))
            p.set_value(std::move(r));
    };
    auto h = eng.start(std::move(spec), std::move(cbs));

    // Ramp-up: let the probe, segment split, and first buffer fills happen (all legitimate
    // allocation) before we start counting. Bail out if it finishes (or fails) before we
    // ever get a steady-state window to sample -- too small a --size for --window-s.
    for (int i = 0; i < 500 && h.progress().downloaded == 0; ++i) {
        if (f.wait_for(0s) == std::future_status::ready) {
            std::fprintf(stderr,
                         "vdm_bench: download finished during ramp-up -- use a bigger "
                         "--size or a smaller --window-s\n");
            return 2;
        }
        std::this_thread::sleep_for(10ms);
    }

    const std::uint64_t before = g_alloc_count.load(std::memory_order_relaxed);
    std::this_thread::sleep_for(std::chrono::duration<double>(window_s));
    const std::uint64_t after = g_alloc_count.load(std::memory_order_relaxed);

    if (f.wait_for(0s) == std::future_status::ready) {
        std::fprintf(stderr,
                     "vdm_bench: download finished during the sampling window -- use a "
                     "bigger --size or a smaller --window-s\n");
        return 2;
    }

    const std::uint64_t delta = after - before;
    const double budget = budget_per_s * window_s + 5;  // +5: fixed slack for one-off events
    std::fprintf(stderr, "alloc-check: %llu allocations in %.2fs (budget %.0f)\n",
                (unsigned long long)delta, window_s, budget);
    if (static_cast<double>(delta) > budget) {
        std::fprintf(stderr,
                     "FAIL: allocation count scales with the transfer, not just periodic "
                     "bookkeeping -- something on the write path is allocating.\n");
        return 1;
    }
    return 0;
}

void usage() {
    std::fprintf(stderr,
                 "usage: vdm_bench <throughput|load|heap-profile|alloc-check> [options]\n"
                 "  throughput   [--size 5G] [--segments N] [--require-mbps N] "
                 "[--max-cpu-pct N]\n"
                 "  load         [--tasks 20] [--task-size 4M] [--segments N] "
                 "[--require-rss-kb N] [--task-timeout-s 300]\n"
                 "  heap-profile [--tasks 20] [--task-size 4M] [--segments N] [--top 20]\n"
                 "  alloc-check  [--size 128M] [--window-s 2] [--budget-per-s 5]\n");
}

}  // namespace

int main(int argc, char **argv) {
    if (argc < 2) {
        usage();
        return 2;
    }
    std::string cmd = argv[1];
    Args args(argc, argv, 2);
    if (cmd == "throughput")
        return cmd_throughput(args);
    if (cmd == "load")
        return cmd_load(args);
    if (cmd == "heap-profile")
        return cmd_heap_profile(args);
    if (cmd == "alloc-check")
        return cmd_alloc_check(args);
    usage();
    return 2;
}