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: [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]
