andygrove opened a new pull request, #6166: URL: https://github.com/apache/datafusion-comet/pull/6166
## Which issue does this PR close? Closes #6165. ## Rationale for this change `libcomet.so` is loaded by the JVM with `dlopen`, so its `thread_local!` variables use the general-dynamic TLS model. Every access calls `__tls_get_addr` in the dynamic loader. The allocation accounting wrapper's `track()` made three thread-local accesses on every allocation and every free: - `IN_TRACK.replace(true)` - `LOCAL_DRIFT.try_with(...)`, which also checks the lazily registered destructor's state - `IN_TRACK.set(false)` With the wrapper installed, this cost about 5% on TPC-H SF100 Q21. The `alloc_overhead` benchmark does not show it, because it links `comet` as an rlib into an executable, where a thread-local access is a single instruction. The wrapper is behind the `alloc-accounting` feature on `main`. #6162 installs it in every build, so this should land first. ## What changes are included in this PR? `alloc_accounting.rs` only. - **One thread-local on the fast path.** The per-thread drift and a phase (unregistered, registering, registered, exited) share one `const`-initialized thread-local with no destructor. Tracking an allocation is now one thread-local access, a phase check and an add. No lazy-initialization check is needed, and the state stays readable while the thread's other thread-local destructors run. - **A separate exit guard.** Settling the remaining drift at thread exit moves to `SettleOnExit`, a zero-sized thread-local whose destructor does the settle. It is registered once per thread, on the thread's first tracked allocation, in the `registering` phase. Registering a destructor can allocate on some platforms, and that allocation re-enters `track` and goes straight to the shared balance. - **After exit.** Once the guard's destructor has run, the phase is `exited`, and deltas from later destructors go straight to the shared balance rather than to a drift nothing would settle. ## How are these changes tested? Unit tests, run in the default build, with `alloc-accounting`, and with `jemalloc,alloc-accounting`: - The existing thread-exit test now registers the guard with a first tracked delta before injecting a drift. Removing the registration makes it fail. - A new test checks that a tracked delta registers the exit guard. - A new test checks that deltas in the `registering` and `exited` phases bypass the drift and reach the shared balance. TPC-H SF100 Q21 on a 32-core Linux box (Spark 4.1.1, 2 executors x 8 cores, glibc allocator). These builds have the wrapper installed by default, as in #6162, with this change applied on top. Each build ran 3 times, alternating, with 3 iterations per run. The table uses the median of iterations 2 and 3. | Build | Q21 median | vs base | | --- | --- | --- | | No wrapper | 30.24 s | | | Wrapper, before this PR | 31.64 s | +4.6% | | Wrapper, this PR | 30.97 s | +2.4% | All builds returned the same result. `perf` over the same query, summed over both executors: | Build | `__rust_alloc` / `__rust_dealloc` self | `__tls_get_addr` | Total | | --- | --- | --- | --- | | No wrapper | 0.48% | 0.08% | 0.56% | | Wrapper, before this PR | 4.97% | 2.63% | 7.60% | | Wrapper, this PR | 3.13% | 1.33% | 4.46% | The remaining cost is the one `__tls_get_addr` per call, and `alloc` calling the backend and accounting on return instead of tail-calling it. The shared atomic is not a factor: its `lock add` instructions account for well under 0.1% of samples. Removing the remaining lookup would need the initial-exec TLS model, which is nightly-only and unsafe for a `dlopen`ed library, or a design without thread-locals. -- 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]
