You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/09/23 20:35:21 UTC
[PR] perf: make one thread-local access per tracked allocation [datafusion-comet]
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: make one thread-local access per tracked allocation [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #6166:
URL: https://github.com/apache/datafusion-comet/pull/6166#discussion_r4096913811
##########
native/core/src/alloc_accounting.rs:
##########
@@ -42,30 +42,64 @@ const SETTLE_THRESHOLD: isize = 64 * 1024;
/// another flushes the matching positive one.
static BALANCE: AtomicIsize = AtomicIsize::new(0);
+/// The phases of a thread's [`ThreadState`].
+///
+/// A thread starts `UNREGISTERED`. Its first tracked allocation moves it through `REGISTERING`,
+/// while [`SETTLE_ON_EXIT`]'s destructor is registered, to `REGISTERED`, where deltas accumulate in
+/// the thread's drift. [`SettleOnExit::drop`] moves it to `EXITED` when the thread ends. In every
+/// phase but `REGISTERED`, deltas go straight to the shared balance: while registering, because
+/// registration can itself allocate on some platforms; after exiting, because nothing would
+/// settle the drift.
+const UNREGISTERED: u8 = 0;
+const REGISTERING: u8 = 1;
+const REGISTERED: u8 = 2;
+const EXITED: u8 = 3;
+
+/// A thread's accounting state, kept in one thread-local so that tracking an allocation costs one
+/// thread-local access.
+///
+/// `libcomet` is a shared library that the JVM loads with `dlopen`, so its thread-locals use the
+/// general-dynamic TLS model: every access calls `__tls_get_addr` in the dynamic loader. That
+/// call is most of the wrapper's cost on allocation-heavy queries, so the fast path makes exactly
+/// one. The state is const-initialized and has no destructor, which also spares the fast path the
+/// lazy-initialization check a thread-local with a destructor needs, and leaves it readable while
+/// the thread's other thread-local destructors run.
Review Comment:
Two claims here don't hold up. By the profile in the description, TLS is about a third of the wrapper's overhead (1.25 of the 3.9 points this PR's build adds), not most of it. And the `phase` check replaces std's lazy-initialization check one for one. Both compile to a byte load, a compare and a branch, so the fast path skips nothing.
Also, `__tls_get_addr` is specific to x86_64 Linux. "Calls into the dynamic loader" holds on every target. I'd cut this down to the invariant: one const-initialized thread-local with no destructor, so `track` makes one TLS access.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -42,30 +42,64 @@ const SETTLE_THRESHOLD: isize = 64 * 1024;
/// another flushes the matching positive one.
static BALANCE: AtomicIsize = AtomicIsize::new(0);
+/// The phases of a thread's [`ThreadState`].
+///
+/// A thread starts `UNREGISTERED`. Its first tracked allocation moves it through `REGISTERING`,
+/// while [`SETTLE_ON_EXIT`]'s destructor is registered, to `REGISTERED`, where deltas accumulate in
+/// the thread's drift. [`SettleOnExit::drop`] moves it to `EXITED` when the thread ends. In every
+/// phase but `REGISTERED`, deltas go straight to the shared balance: while registering, because
+/// registration can itself allocate on some platforms; after exiting, because nothing would
+/// settle the drift.
+const UNREGISTERED: u8 = 0;
+const REGISTERING: u8 = 1;
+const REGISTERED: u8 = 2;
+const EXITED: u8 = 3;
Review Comment:
`REGISTERING` can't be reached with Rust 1.93 or later, since registering a thread-local destructor no longer allocates through the global allocator (see the summary). In a probe on 1.94 and 1.95, registration never re-entered `track`. Dropping it removes this phase, the rationale here and at L149-152, and the `REGISTERING` half of `deltas_bypass_the_drift_outside_the_registered_phase`.
The three phases left would read better as a `Copy` enum in the `Cell`, like `CopyMode` and `ScanIoStoreRole` elsewhere in `native/core`:
```rust
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Phase {
Unregistered,
Registered,
Exited,
}
```
##########
native/core/src/alloc_accounting.rs:
##########
@@ -360,28 +406,91 @@ mod tests {
/// A thread's remaining drift must reach the shared balance when the thread exits.
///
- /// This drops a `ThreadDrift` holding a drift directly, rather than injecting one into a real
- /// thread's `LOCAL_DRIFT` and letting the thread exit. With the wrapper installed, the thread's
- /// teardown allocates, and those allocations flush an oversized drift through `track` before
- /// the destructor runs, so a thread-exit test would pass without the destructor. That the
- /// destructor runs when a thread exits is the `thread_local!` guarantee; what needs testing is
- /// that it settles the drift. The injected amount is far larger than any real allocation, and
- /// is taken back out afterwards.
+ /// This injects a drift into a registered thread's state and drops a `SettleOnExit` directly,
+ /// rather than letting the thread exit. With the wrapper installed, the thread's teardown
+ /// allocates, and those allocations flush an oversized drift through `track` before the
+ /// destructor runs, so a thread-exit test would pass without the destructor. Nothing between
+ /// the injection and the drop allocates, so only the destructor can move the drift. That the
+ /// destructor runs when a registered thread exits is the `thread_local!` guarantee; what needs
+ /// testing is that it settles the drift. The drop marks the thread exited, so the test runs on
+ /// a thread of its own. The injected amount is far larger than any real allocation, and is
+ /// taken back out afterwards.
#[test]
- fn dropping_a_thread_drift_settles_it() {
+ fn dropping_the_exit_hook_settles_the_drift() {
Review Comment:
Nit: the tests call `SettleOnExit` the "exit hook" and the description calls it the "exit guard". Using `SettleOnExit` throughout, for example `dropping_settle_on_exit_settles_the_drift`, would keep one name for it.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -100,23 +134,35 @@ fn track(delta: isize) {
return;
}
- // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will
- // clear it, so this frame must only settle and return.
- if IN_TRACK.with(|in_track| in_track.replace(true)) {
- BALANCE.fetch_add(delta, Ordering::Relaxed);
- return;
- }
+ STATE.with(|state| match state.phase.get() {
+ REGISTERED => settle(&state.drift, delta),
+ UNREGISTERED => register_and_track(state, delta),
+ _ => {
+ BALANCE.fetch_add(delta, Ordering::Relaxed);
+ }
+ })
Review Comment:
LLVM lays out the `_` arm as the fall-through, so the hot `REGISTERED` case is a taken branch in all four `GlobalAlloc` methods (checked on aarch64). Testing `REGISTERED` alone here and moving the `EXITED` case into the `#[cold]` function makes the hot case the fall-through and leaves one branch in `track`:
```rust
STATE.with(|state| {
if state.phase.get() == REGISTERED {
settle(&state.drift, delta)
} else {
track_slow(state, delta)
}
})
```
In `register_and_track`, the `else` at L163-165 can't run. `try_with` fails only once `SETTLE_ON_EXIT` has been destroyed, and by then `SettleOnExit::drop` has set `EXITED`, so `track` never comes back here. `registered` is also tested twice (L158-161). One `if` covers it.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -360,28 +406,91 @@ mod tests {
/// A thread's remaining drift must reach the shared balance when the thread exits.
///
- /// This drops a `ThreadDrift` holding a drift directly, rather than injecting one into a real
- /// thread's `LOCAL_DRIFT` and letting the thread exit. With the wrapper installed, the thread's
- /// teardown allocates, and those allocations flush an oversized drift through `track` before
- /// the destructor runs, so a thread-exit test would pass without the destructor. That the
- /// destructor runs when a thread exits is the `thread_local!` guarantee; what needs testing is
- /// that it settles the drift. The injected amount is far larger than any real allocation, and
- /// is taken back out afterwards.
+ /// This injects a drift into a registered thread's state and drops a `SettleOnExit` directly,
+ /// rather than letting the thread exit. With the wrapper installed, the thread's teardown
+ /// allocates, and those allocations flush an oversized drift through `track` before the
+ /// destructor runs, so a thread-exit test would pass without the destructor. Nothing between
+ /// the injection and the drop allocates, so only the destructor can move the drift. That the
+ /// destructor runs when a registered thread exits is the `thread_local!` guarantee; what needs
+ /// testing is that it settles the drift. The drop marks the thread exited, so the test runs on
+ /// a thread of its own. The injected amount is far larger than any real allocation, and is
+ /// taken back out afterwards.
#[test]
- fn dropping_a_thread_drift_settles_it() {
+ fn dropping_the_exit_hook_settles_the_drift() {
const INJECTED: isize = 1 << 40;
let _guard = serial();
- let before = BALANCE.load(Ordering::Relaxed);
- drop(ThreadDrift(Cell::new(INJECTED)));
- let moved = BALANCE.load(Ordering::Relaxed) - before;
+ let (moved, phase) = std::thread::spawn(|| {
+ track(1);
+ let before = BALANCE.load(Ordering::Relaxed);
+ STATE.with(|state| state.drift.set(state.drift.get() + INJECTED));
+ drop(SettleOnExit);
+ let moved = BALANCE.load(Ordering::Relaxed) - before;
+ let phase = STATE.with(|state| state.phase.get());
+ track(-1);
+ (moved, phase)
+ })
+ .join()
+ .unwrap();
BALANCE.fetch_sub(INJECTED, Ordering::Relaxed);
assert!(
moved >= INJECTED / 2,
- "a dropped thread drift never reached the shared balance: balance moved {moved} \
- bytes, expected at least {}",
+ "a dropped exit hook never settled the thread's drift: balance moved {moved} bytes, \
+ expected at least {}",
INJECTED / 2
);
+ assert_eq!(
+ phase, EXITED,
+ "a dropped exit hook did not mark the thread exited"
+ );
+ }
+
+ /// A thread's first tracked delta registers its exit hook, after which deltas accumulate in
+ /// the thread's drift. With the wrapper installed, the thread's own allocations have already
+ /// done this by the time the closure runs, so only the end state is checked.
+ #[test]
+ fn a_tracked_delta_registers_the_exit_hook() {
+ std::thread::spawn(|| {
+ track(1);
+ assert_eq!(STATE.with(|state| state.phase.get()), REGISTERED);
+ track(-1);
+ })
+ .join()
+ .unwrap();
+ }
+
+ /// Outside the `REGISTERED` phase, a delta must go straight to the shared balance: while
+ /// registering, because the drift is not yet settled on exit, and after exiting, because it
+ /// never will be.
+ #[test]
+ fn deltas_bypass_the_drift_outside_the_registered_phase() {
Review Comment:
Once `REGISTERING` goes, only the `EXITED` case is left. It could fold into `dropping_the_exit_hook_settles_the_drift`: after the drop, call `track(INJECTED)` and check that the drift stays put and the balance moves. That covers `EXITED` through the real transition instead of a hand-set phase, and leaves one `INJECTED`.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -360,28 +406,91 @@ mod tests {
/// A thread's remaining drift must reach the shared balance when the thread exits.
///
- /// This drops a `ThreadDrift` holding a drift directly, rather than injecting one into a real
- /// thread's `LOCAL_DRIFT` and letting the thread exit. With the wrapper installed, the thread's
- /// teardown allocates, and those allocations flush an oversized drift through `track` before
- /// the destructor runs, so a thread-exit test would pass without the destructor. That the
- /// destructor runs when a thread exits is the `thread_local!` guarantee; what needs testing is
- /// that it settles the drift. The injected amount is far larger than any real allocation, and
- /// is taken back out afterwards.
+ /// This injects a drift into a registered thread's state and drops a `SettleOnExit` directly,
+ /// rather than letting the thread exit. With the wrapper installed, the thread's teardown
+ /// allocates, and those allocations flush an oversized drift through `track` before the
+ /// destructor runs, so a thread-exit test would pass without the destructor. Nothing between
+ /// the injection and the drop allocates, so only the destructor can move the drift. That the
+ /// destructor runs when a registered thread exits is the `thread_local!` guarantee; what needs
+ /// testing is that it settles the drift. The drop marks the thread exited, so the test runs on
+ /// a thread of its own. The injected amount is far larger than any real allocation, and is
+ /// taken back out afterwards.
#[test]
- fn dropping_a_thread_drift_settles_it() {
+ fn dropping_the_exit_hook_settles_the_drift() {
const INJECTED: isize = 1 << 40;
let _guard = serial();
- let before = BALANCE.load(Ordering::Relaxed);
- drop(ThreadDrift(Cell::new(INJECTED)));
- let moved = BALANCE.load(Ordering::Relaxed) - before;
+ let (moved, phase) = std::thread::spawn(|| {
+ track(1);
+ let before = BALANCE.load(Ordering::Relaxed);
+ STATE.with(|state| state.drift.set(state.drift.get() + INJECTED));
+ drop(SettleOnExit);
+ let moved = BALANCE.load(Ordering::Relaxed) - before;
+ let phase = STATE.with(|state| state.phase.get());
+ track(-1);
+ (moved, phase)
+ })
+ .join()
+ .unwrap();
BALANCE.fetch_sub(INJECTED, Ordering::Relaxed);
assert!(
moved >= INJECTED / 2,
- "a dropped thread drift never reached the shared balance: balance moved {moved} \
- bytes, expected at least {}",
+ "a dropped exit hook never settled the thread's drift: balance moved {moved} bytes, \
+ expected at least {}",
INJECTED / 2
);
+ assert_eq!(
+ phase, EXITED,
+ "a dropped exit hook did not mark the thread exited"
+ );
+ }
+
+ /// A thread's first tracked delta registers its exit hook, after which deltas accumulate in
+ /// the thread's drift. With the wrapper installed, the thread's own allocations have already
+ /// done this by the time the closure runs, so only the end state is checked.
+ #[test]
+ fn a_tracked_delta_registers_the_exit_hook() {
Review Comment:
This can't fail now that the wrapper is always installed. The thread is already `REGISTERED` before the closure starts, so `register_and_track` never runs, and a probe showed the test still passes without `track(1)`. I'd delete it, or set the phase back to `UNREGISTERED` first so it goes through the registration path.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: make one thread-local access per tracked allocation [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #6166:
URL: https://github.com/apache/datafusion-comet/pull/6166#issuecomment-5821703602
Thanks @mbutrovich and @comphead. I've gone with @comphead's suggestion. The PR now only deletes `IN_TRACK` from `main`'s version, so `LOCAL_DRIFT.try_with` is the only thread-local access, and std's own state for `LOCAL_DRIFT` does what the hand-rolled phase did. The inline comments on `STATE` and `SettleOnExit` no longer apply, since that code is gone.
I checked the variant on both platforms before switching. On macOS it makes one TLV thunk call in each `GlobalAlloc` method, as the first revision did, and an interleaved A/B over the `alloc_overhead` cases put it level with the first revision and about 1.5 ns per alloc/free pair ahead of `main`. For x86_64 Linux I compiled a standalone copy with the flags Cargo uses for `libcomet`, which gives one `__tls_get_addr` call per method, against two on `main`. That needed checking on its own, because a cdylib-only build of the same code uses the local-dynamic TLS model and makes a single call even on `main`. `libcomet` is also an rlib, which keeps these thread-locals exported and so general-dynamic. It also means Cargo never applies the release profile's `lto = "thin"` to `libcomet`. I haven't re-run Q21 on this revision, and the description says which numbers come from the first one.
The tests pass on the new head with the system allocator, jemalloc and mimalloc, the description is updated for #6162, and CI is running again now that Preflight got past the mirror.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: make one thread-local access per tracked allocation [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6166:
URL: https://github.com/apache/datafusion-comet/pull/6166#discussion_r4098144746
##########
native/core/src/alloc_accounting.rs:
##########
@@ -360,28 +406,91 @@ mod tests {
/// A thread's remaining drift must reach the shared balance when the thread exits.
///
- /// This drops a `ThreadDrift` holding a drift directly, rather than injecting one into a real
- /// thread's `LOCAL_DRIFT` and letting the thread exit. With the wrapper installed, the thread's
- /// teardown allocates, and those allocations flush an oversized drift through `track` before
- /// the destructor runs, so a thread-exit test would pass without the destructor. That the
- /// destructor runs when a thread exits is the `thread_local!` guarantee; what needs testing is
- /// that it settles the drift. The injected amount is far larger than any real allocation, and
- /// is taken back out afterwards.
+ /// This injects a drift into a registered thread's state and drops a `SettleOnExit` directly,
+ /// rather than letting the thread exit. With the wrapper installed, the thread's teardown
+ /// allocates, and those allocations flush an oversized drift through `track` before the
+ /// destructor runs, so a thread-exit test would pass without the destructor. Nothing between
+ /// the injection and the drop allocates, so only the destructor can move the drift. That the
+ /// destructor runs when a registered thread exits is the `thread_local!` guarantee; what needs
+ /// testing is that it settles the drift. The drop marks the thread exited, so the test runs on
+ /// a thread of its own. The injected amount is far larger than any real allocation, and is
+ /// taken back out afterwards.
#[test]
- fn dropping_a_thread_drift_settles_it() {
+ fn dropping_the_exit_hook_settles_the_drift() {
const INJECTED: isize = 1 << 40;
let _guard = serial();
- let before = BALANCE.load(Ordering::Relaxed);
- drop(ThreadDrift(Cell::new(INJECTED)));
- let moved = BALANCE.load(Ordering::Relaxed) - before;
+ let (moved, phase) = std::thread::spawn(|| {
+ track(1);
+ let before = BALANCE.load(Ordering::Relaxed);
+ STATE.with(|state| state.drift.set(state.drift.get() + INJECTED));
+ drop(SettleOnExit);
+ let moved = BALANCE.load(Ordering::Relaxed) - before;
+ let phase = STATE.with(|state| state.phase.get());
+ track(-1);
+ (moved, phase)
+ })
+ .join()
+ .unwrap();
BALANCE.fetch_sub(INJECTED, Ordering::Relaxed);
assert!(
moved >= INJECTED / 2,
- "a dropped thread drift never reached the shared balance: balance moved {moved} \
- bytes, expected at least {}",
+ "a dropped exit hook never settled the thread's drift: balance moved {moved} bytes, \
+ expected at least {}",
INJECTED / 2
);
+ assert_eq!(
+ phase, EXITED,
+ "a dropped exit hook did not mark the thread exited"
+ );
+ }
+
+ /// A thread's first tracked delta registers its exit hook, after which deltas accumulate in
+ /// the thread's drift. With the wrapper installed, the thread's own allocations have already
+ /// done this by the time the closure runs, so only the end state is checked.
+ #[test]
+ fn a_tracked_delta_registers_the_exit_hook() {
+ std::thread::spawn(|| {
+ track(1);
+ assert_eq!(STATE.with(|state| state.phase.get()), REGISTERED);
+ track(-1);
+ })
+ .join()
+ .unwrap();
+ }
Review Comment:
It was `thread_exit_settles_remaining_drift`, on macOS, before I merged #6162. That test only ran in the build without the wrapper, where the test thread's `track(1)` was its first tracked delta, so skipping the registration made it fail. #6162 removed that build, and the direct drop the merge put in its place can't see the registration, as you found. With @comphead's version there's no registration of our own left to test: `thread_local!` registers `LOCAL_DRIFT`'s destructor on first access, and `dropping_a_thread_drift_settles_it` covers what the destructor does.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -100,23 +134,35 @@ fn track(delta: isize) {
return;
}
- // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will
- // clear it, so this frame must only settle and return.
- if IN_TRACK.with(|in_track| in_track.replace(true)) {
- BALANCE.fetch_add(delta, Ordering::Relaxed);
- return;
- }
+ STATE.with(|state| match state.phase.get() {
+ REGISTERED => settle(&state.drift, delta),
+ UNREGISTERED => register_and_track(state, delta),
+ _ => {
+ BALANCE.fetch_add(delta, Ordering::Relaxed);
+ }
+ })
Review Comment:
I measured this before switching. The first revision with this change, and with `REGISTERING` dropped, came out 0.3 to 0.6 ns per alloc/free pair faster than the `IN_TRACK`-less version at 4 KiB and up on aarch64. On x86_64 the two differ by one compare and a taken branch, with the same number of prologue and epilogue instructions. That didn't seem worth keeping the hand-rolled state for.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: make one thread-local access per tracked allocation [datafusion-comet]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #6166:
URL: https://github.com/apache/datafusion-comet/pull/6166#discussion_r4096765206
##########
native/core/src/alloc_accounting.rs:
##########
@@ -360,28 +406,91 @@ mod tests {
/// A thread's remaining drift must reach the shared balance when the thread exits.
///
- /// This drops a `ThreadDrift` holding a drift directly, rather than injecting one into a real
- /// thread's `LOCAL_DRIFT` and letting the thread exit. With the wrapper installed, the thread's
- /// teardown allocates, and those allocations flush an oversized drift through `track` before
- /// the destructor runs, so a thread-exit test would pass without the destructor. That the
- /// destructor runs when a thread exits is the `thread_local!` guarantee; what needs testing is
- /// that it settles the drift. The injected amount is far larger than any real allocation, and
- /// is taken back out afterwards.
+ /// This injects a drift into a registered thread's state and drops a `SettleOnExit` directly,
+ /// rather than letting the thread exit. With the wrapper installed, the thread's teardown
+ /// allocates, and those allocations flush an oversized drift through `track` before the
+ /// destructor runs, so a thread-exit test would pass without the destructor. Nothing between
+ /// the injection and the drop allocates, so only the destructor can move the drift. That the
+ /// destructor runs when a registered thread exits is the `thread_local!` guarantee; what needs
+ /// testing is that it settles the drift. The drop marks the thread exited, so the test runs on
+ /// a thread of its own. The injected amount is far larger than any real allocation, and is
+ /// taken back out afterwards.
#[test]
- fn dropping_a_thread_drift_settles_it() {
+ fn dropping_the_exit_hook_settles_the_drift() {
const INJECTED: isize = 1 << 40;
let _guard = serial();
- let before = BALANCE.load(Ordering::Relaxed);
- drop(ThreadDrift(Cell::new(INJECTED)));
- let moved = BALANCE.load(Ordering::Relaxed) - before;
+ let (moved, phase) = std::thread::spawn(|| {
+ track(1);
+ let before = BALANCE.load(Ordering::Relaxed);
+ STATE.with(|state| state.drift.set(state.drift.get() + INJECTED));
+ drop(SettleOnExit);
+ let moved = BALANCE.load(Ordering::Relaxed) - before;
+ let phase = STATE.with(|state| state.phase.get());
+ track(-1);
+ (moved, phase)
+ })
+ .join()
+ .unwrap();
BALANCE.fetch_sub(INJECTED, Ordering::Relaxed);
assert!(
moved >= INJECTED / 2,
- "a dropped thread drift never reached the shared balance: balance moved {moved} \
- bytes, expected at least {}",
+ "a dropped exit hook never settled the thread's drift: balance moved {moved} bytes, \
+ expected at least {}",
INJECTED / 2
);
+ assert_eq!(
+ phase, EXITED,
+ "a dropped exit hook did not mark the thread exited"
+ );
+ }
+
+ /// A thread's first tracked delta registers its exit hook, after which deltas accumulate in
+ /// the thread's drift. With the wrapper installed, the thread's own allocations have already
+ /// done this by the time the closure runs, so only the end state is checked.
+ #[test]
+ fn a_tracked_delta_registers_the_exit_hook() {
+ std::thread::spawn(|| {
+ track(1);
+ assert_eq!(STATE.with(|state| state.phase.get()), REGISTERED);
+ track(-1);
+ })
+ .join()
+ .unwrap();
+ }
Review Comment:
The description says removing the registration makes the thread-exit test fail. I couldn't reproduce that on macOS. With line 157 changed to `let registered = true;`, so `SettleOnExit` is never registered, `cargo test -p datafusion-comet --lib alloc_accounting` still passes all 7 tests, with both the system allocator and `--features jemalloc`. Removing `track(1)` at line 424 also leaves them passing. This test can't catch it either, since as its doc says, the thread has already registered by the time the closure runs.
Which change did you check, and on which platform? If nothing fails without the registration, could a test observe it from a destructor that runs after `SettleOnExit`'s? One way is a raw `pthread` that touches a probe thread-local with a destructor before its first tracked allocation. Destructors run in reverse registration order, so the probe runs last and can record whether `STATE.phase` was `EXITED` by then.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: make one thread-local access per tracked allocation [datafusion-comet]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich merged PR #6166:
URL: https://github.com/apache/datafusion-comet/pull/6166
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: make one thread-local access per tracked allocation [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #6166:
URL: https://github.com/apache/datafusion-comet/pull/6166#issuecomment-5824270908
@comphead I've run Q21 on this revision now, with all four builds re-run in one session on the same Linux box. The description has the full tables.
| Build | Q21 median | vs no wrapper |
| --- | --- | --- |
| No wrapper | 30.21 s | |
| Wrapper, `main` | 31.66 s | +4.8% |
| Wrapper, first revision | 31.01 s | +2.6% |
| Wrapper, this revision | 30.85 s | +2.1% |
The two revisions are within each other's run-to-run spread, and `__tls_get_addr` is 1.38% of samples for both, against 2.60% on `main`.
I also checked your x86_64 estimate against the release `libcomet.so`. The state check is the same in both revisions: a byte load, a compare and a taken branch. This revision saves one more callee-saved register than the first revision as pushed (r15, r14 and rbx, against r14 and rbx), because std's inlined `initialize` keeps the thread-local's address live across the destructor registration. In `dealloc` and `realloc` that register takes the place of the first revision's alignment push, so the instruction count doesn't change. `alloc` and `alloc_zeroed` get two more instructions, a `sub rsp` and a `pop`. So on x86_64 it costs one register rather than two push/pop pairs, and only those two methods pay for it. That fits about 0.2 points more `__rust_alloc` self time in the profile, and it doesn't show in Q21. It also refines what I said about the prologue earlier: against the first revision as pushed, the count only matches in `dealloc` and `realloc`.
On the DuckDB idea, I've opened #6213. Of the 4.2 points the wrapper still adds, `__tls_get_addr` is 1.3 and the rest is in `__rust_alloc` and `__rust_dealloc` themselves. So 1.3 points is roughly what dropping the thread-local could save, before paying for `sched_getcpu` and an atomic add on every call.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: make one thread-local access per tracked allocation [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #6166:
URL: https://github.com/apache/datafusion-comet/pull/6166#issuecomment-5818673550
Thanks @andygrove that looks promising
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org