analyze_trace.rs over-counts pool reservations by summing per-thread counters #5990

Describe the bug

analyze_trace.rs compares the native allocation counter against "the sum of per-thread Comet memory pool reservations" and reports points where allocation exceeds the pool total. The pool total is computed at native/common/src/bin/analyze_trace.rs:164:

let pool_total: u64 = pool_by_thread.values().sum();

That sum is not the total reserved. The per-thread counters come from total_reserved_for_thread (native/core/src/execution/jni_api.rs:216), which sums the pools registered on one thread and deduplicates shared pools only within that thread:

let mut seen = HashSet::new();
pools.values().filter_map(|p| {
    let ptr = Arc::as_ptr(p) as *const ();
    seen.insert(ptr).then(|| p.reserved())
}).sum::<usize>()

Task-shared pools are shared across threads, so every thread that references one reports its full reservation. Summing across threads therefore multiplies a task-shared pool's reservation by the number of threads referencing it.

Evidence

Measured on TPC-H SF100 (2 executors x 8 cores, Spark 4.1.1, spark.memory.offHeap.size=16g, tracing enabled, jemalloc,alloc-accounting), analysing the per-thread sum the same way this tool does:

• reserved / native ratio reached 500x. A pool cannot have reserved 500 times what the allocator has handed out.
• 27% of samples reported reserved greater than native allocation.
• Reserved was 0 at the median, because threads with no registered pools contribute nothing and early samples see an empty map.
• "Peak gap" landed on samples where reserved happened to be 0, so the reported peak was just peak allocation rather than a gap.

Replacing the per-thread sum with a single process-wide total over the distinct entries of TASK_SHARED_MEMORY_POOLS changed the same traces substantially: max gap fell from 1694 MB to 560 MB, and the max ratio from 366x to 41x.

Impact

The tool's headline output, the points where allocation exceeds the pool total, is unreliable, and so is any gap figure derived from it. That includes at least one conclusion already recorded on #5934 ("moments with >1 GB native vs ~0 pool reservations"), which was produced this way.

This is a documented tool, not a scratch binary: the contributor guide describes it and gives the command to run it (docs/source/contributor-guide/tracing.md). It is used to diagnose exactly the accounting questions tracked in #4576, so wrong numbers here are actively misleading.

Suggested fix

Emit one process-wide total and analyse against that, rather than summing per-thread counters:

1. Add a helper that sums reserved() over the distinct entries of TASK_SHARED_MEMORY_POOLS (keyed by task attempt, so each pool is counted once).
2. Log it as a single counter at the same point native_allocated is emitted (jni_api.rs ~1103), so both halves of a sample are true at the same instant.
3. Have analyze_trace.rs anchor on that counter instead of pool_by_thread.values().sum().

Step 3 matters on its own: the counters are emitted in the order jemalloc, native, per-thread, total, so anchoring a sample on native_allocated pairs it with the previous group's total. That one-sample lag alone accounted for the difference between a 1694 MB and a 560 MB peak gap in the numbers above.

I have a working 23-line patch for steps 1 and 2 from the investigation in #4576 and can turn it into a PR.

Additional context

Follows from the measurements posted on #4576. The per-thread counters remain useful for their original purpose, which is per-thread attribution in a trace viewer; the defect is only in summing them to obtain a process total.