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