andygrove opened a new pull request, #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934
## Which issue does this PR close?
Relates to #4576. This is the first of the pieces extracted from the closed
prototype in #4582, and it deliberately stops short of enforcement, so it
does
not close that issue.
## Rationale for this change
Comet's memory pool counts *declared reservations*: bytes an operator
explicitly
asked for. A lot of real allocation never goes through it — Arrow builders,
expression kernels producing intermediates, decompression buffers, Parquet
metadata, `object_store` buffers, tokio's own machinery. Pool reservations
are
therefore a lower bound on Comet's footprint, and the size of the gap is
workload-dependent.
Today that gap is invisible at runtime. Diagnosing an out-of-memory report
means
reasoning about it indirectly, and `spark.comet.exec.memoryPool.fraction`
asks
operators to hand-tune a haircut for a quantity nobody can measure.
#4582 tried to both measure the gap and enforce on it. Reviewing it
convinced me
the enforcement half needs a question answered first — how well does a
tracked
byte balance actually track RSS on real workloads? — and that question is
much
easier to answer if the measurement lands on its own. So this PR is the
measurement only.
## What changes are included in this PR?
Behind a new, off-by-default `alloc-accounting` cargo feature:
- `native/core/src/alloc_accounting.rs`: `AccountingAllocator<A>` wraps
whichever
global allocator the build selected and maintains one signed process-wide
byte
balance, exposed by `current_balance()`.
- `native/core/src/lib.rs`: installs the wrapper over jemalloc / mimalloc /
system. The cfgs are mutually exclusive, so a build without the feature is
exactly the previous arrangement — no wrapper, no per-allocation work.
- `native/core/src/execution/jni_api.rs`: reports the balance as the
`native_allocated` tracing metric, logged in the same place as
`jemalloc_allocated` and alongside the per-thread pool reservations it is
meant
to be compared against.
- `tracing.md`: documents the feature and the new metric.
It is observability only. It never rejects an allocation, never panics, and
does
not touch the memory pool. Two things follow from that which are worth
calling
out:
Because it cannot fail an allocation, `realloc` accounts *after* delegating
and
only on success. The prototype had to account before delegating, because
panicking after `inner.realloc` would leave a caller unwinding with a stale
pointer — a soundness constraint that simply does not exist here, and
dropping it
also removes the over-count on a failed realloc.
Per-thread deltas are batched and flushed at 64 KiB, so the common path is a
thread-local add-and-compare rather than an atomic read-modify-write. The
prototype leaked up to 64 KiB of accounting every time a thread died, which
matters because the blocking pool churns on tokio's idle timeout.
`ThreadDrift`'s
destructor settles the remainder on exit. That is slightly more delicate
than it
looks: touching a destructor-bearing thread-local can itself allocate on
first
use, so `track()` keeps a destructor-free re-entrancy flag and settles
re-entrant
calls straight into the shared balance, and uses `try_with` so an allocation
during thread teardown cannot panic inside the allocator.
### What this does not do
The balance counts `Layout` bytes, not resident pages, so it excludes
allocator
fragmentation, jemalloc's retained pages, `mmap`ed regions, and anything a C
dependency allocates through libc `malloc`. It is a lower bound on RSS —
just a
much tighter one than pool reservations. It is also process-wide, not
per-task.
Whether it is a good enough proxy to enforce on is exactly what I would like
to
learn from it before proposing that.
## How are these changes tested?
Unit tests in `alloc_accounting.rs` cover the settle/flush helper and the
negative-balance clamp, and two tests drive real allocations through the
installed allocator: one asserts an 8 MiB allocation moves the reported
balance,
and one asserts that drift from exited threads reaches the shared balance.
The thread-exit test is mutation-checked. Each worker allocates a
sub-threshold
buffer and hands ownership back before exiting, so the matching free happens
on
the main thread after the worker is gone and the destructor is the only path
by
which those bytes can be counted. Neutering the destructor fails it with
"balance moved 0 bytes, expected at least 1048576". An earlier version of the
test, where each thread freed what it allocated, passed under the same
mutation
and was replaced.
`cargo clippy --all-targets -- -D warnings` and the full native test suite
pass
under `default`, `alloc-accounting`, `jemalloc,alloc-accounting`, and
`mimalloc,alloc-accounting`.
Not yet measured: the per-allocation overhead of the wrapper when the
feature is
on. That is worth a benchmark before anyone considers enabling it by default,
and I have not done it.
--
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]
