andygrove opened a new pull request, #5991:
URL: https://github.com/apache/datafusion-comet/pull/5991
## Which issue does this PR close?
Closes #5990.
## Rationale for this change
`analyze_trace` compares the process-wide allocation counter against the
total memory reserved by Comet's pools, and reports points where allocation
exceeds it. It derives that total by summing the per-thread
`thread_NNN_comet_memory_reserved` counters:
```rust
let pool_total: u64 = pool_by_thread.values().sum();
```
That is not the total reserved. Those counters come from
`total_reserved_for_thread`, which sums the pools registered on one thread and
deduplicates shared pools only *within* that thread. Task-shared pools are
shared *across* threads, so every thread referencing one reports its full
reservation, and summing across threads multiplies it by the number of
referencing threads.
Measured on TPC-H SF100 (2 executors x 8 cores, Spark 4.1.1, 16g off-heap,
tracing enabled): the reported total reached **500x** the real allocation,
**27%** of samples claimed more memory reserved than the allocator had handed
out, and the median reported total was **0** because threads with no registered
pools contribute nothing. The tool's "peak gap" then landed on samples where
the total happened to be zero, so it reported peak allocation rather than a gap.
The tool is documented in the contributor guide and is used to investigate
exactly the memory accounting questions tracked in #4576, so wrong numbers here
are actively misleading. Correcting the source changed the same traces
substantially: peak gap fell from 1694 MB to 560 MB.
## What changes are included in this PR?
- `total_reserved_across_tasks()` in
`native/core/src/execution/memory_pools/task_shared.rs` sums `reserved()` over
the distinct entries of `TASK_SHARED_MEMORY_POOLS`. The registry is keyed by
task attempt, so each pool is counted exactly once regardless of how many
threads reference it.
- `jni_api.rs` emits that as a single `comet_memory_reserved_total` counter,
at the same point `native_allocated` is already emitted, so both halves of a
sample are true at the same instant. The per-thread counters are unchanged,
since they remain useful for per-thread attribution in a trace viewer.
- `analyze_trace.rs` prefers `comet_memory_reserved_total` when the trace
carries it. Traces without it still analyze, using the per-thread sum, but the
tool now says plainly that those numbers over-count.
Two details worth flagging for review:
- The total must be matched before the existing
`name.contains("comet_memory_reserved")` branch, which would otherwise also
match `comet_memory_reserved_total` and fold the process-wide counter into the
per-thread map.
- Anchoring matters independently of the sum. The counters are emitted in
the order allocated, per-thread, total, so pairing a sample on the allocation
counter alone picks up the previous group's total. That lag alone accounted for
part of the difference between a 1694 MB and a 560 MB peak gap when this was
measured offline.
## How are these changes tested?
- Two new unit tests for `total_reserved_across_tasks()`. The first covers
the case this bug is about: a task-shared pool referenced by several execution
contexts is counted once, not once per reference. The second covers release,
that dropping the last reference removes the pool from the total. Both take a
serial lock, because the registry is process-wide and the crate's tests run in
parallel, so without it the two tests observe each other's reservations. #5212
notes there are currently no unit tests for any of the pools.
- Verified end to end by running the built `analyze_trace` against a real
TPC-H SF100 trace, once as recorded and once with the new counter stripped out
to exercise the fallback. Same trace, same allocation counter:
| pool total source | peak pool total |
| --- | --- |
| `comet_memory_reserved_total` | 2080.3 MB |
| sum of 8 per-thread counters | 2896.2 MB |
Peak `native_allocated` on that trace was 2100.3 MB, so the per-thread sum
invents about 816 MB of reservation that was never held. The fallback path
prints the over-count warning and still produces a report, so existing traces
remain readable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
