Skip to content

perf: make one thread-local access per tracked allocation - #6166

Open
andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:perf-alloc-accounting-tls
Open

andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:perf-alloc-accounting-tls

Conversation

@andygrove

@andygrove andygrove commented Sep 23, 2026

Copy link
Copy Markdown
Member

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.

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 dlopened library, or a design without thread-locals.

libcomet is loaded with dlopen, so its thread-locals use the
general-dynamic TLS model and every access calls __tls_get_addr. track()
made three accesses per allocation and per free (the re-entrancy flag
twice, and the drift cell with its lazy-destructor state check), which
cost about 5% on TPC-H Q21.

Keep the drift and a phase in one const-initialized thread-local with
no destructor, so the fast path is one access, a phase check and an add.
The exit-time settle moves to a zero-sized guard thread-local that is
registered once per thread, on its first tracked allocation, with the
phase marking the re-entrant allocations registration can make. After
the guard's destructor runs, deltas go straight to the shared balance.

Closes apache#6165.
@github-actions github-actions Bot added enhancement New feature or request performance labels Sep 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf: allocation accounting wrapper costs ~5% on TPC-H Q21 from thread-local lookups in the dlopened library

1 participant