sunchao commented on code in PR #5991:
URL: https://github.com/apache/datafusion-comet/pull/5991#discussion_r4031375527
##########
native/core/src/execution/jni_api.rs:
##########
@@ -1107,6 +1107,13 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_executePlan(
&exec_context.tracing_memory_metric_name,
total_reserved_for_thread(exec_context.rust_thread_id) as u64,
);
+ // Process-wide total, emitted here so it is true at the same
instant as the
+ // allocation counter above. The per-thread counter cannot be
summed across threads:
+ // it reports a shared pool's full reservation once per
referencing thread.
+ log_memory_usage(
+ "comet_memory_reserved_total",
+ crate::execution::memory_pools::total_reserved_across_tasks()
as u64,
+ );
Review Comment:
### Correctness
[P2] Include non-task-shared pools in the authoritative total
`total_reserved_across_tasks()` only sees pools created through
`acquire_task_shared_pool`. The supported on-heap `greedy`, `fair_spill`,
`greedy_global`, `fair_spill_global`, and `unbounded` branches in
`create_memory_pool` bypass that registry. A traced executor using one of those
modes therefore emits zero here despite live reservations. The analyzer now
treats that zero as authoritative and discards the valid per-thread
reservation, reporting tracked allocations as excess. With allocation and a
per-thread reservation both 100 MiB and this total zero, the base analyzer
reports zero excess while this head reports 100 MiB. Please include every
supported pool with identity deduplication, or avoid publishing an
authoritative process total for configurations the registry does not cover.
##########
native/core/src/execution/memory_pools/task_shared.rs:
##########
@@ -94,6 +94,22 @@ impl Drop for TaskSharedMemoryPool {
}
}
+/// Total bytes reserved across every live task-shared pool, process-wide.
+///
+/// Each registry entry is one task attempt's pool, so summing the entries
counts each pool exactly
+/// once. This is deliberately not the sum of the per-thread
`thread_NNN_comet_memory_reserved`
+/// tracing counters: each of those reports the full reservation of a shared
pool once per thread
+/// that references it, so adding them across threads multiplies a task-shared
pool by its thread
+/// count.
+pub(crate) fn total_reserved_across_tasks() -> usize {
+ TASK_SHARED_MEMORY_POOLS
+ .lock()
+ .values()
+ .filter_map(Weak::upgrade)
+ .map(|pool| pool.reserved())
Review Comment:
### Correctness
[P1] Release the registry lock before dropping upgraded pool references
The mutex guard lives through this iterator chain, while each upgraded `Arc`
is dropped at the end of the `map` closure. If another task releases its last
external reference after `Weak::upgrade`, that temporary becomes the last owner
and `TaskSharedMemoryPool::drop` tries to lock `TASK_SHARED_MEMORY_POOLS`
again. The tracing thread deadlocks holding the process-wide registry, blocking
subsequent pool acquisition/release too. This is reachable when one traced task
samples while another finishes. An isolated Rust probe using these exact
total/acquire/Drop bodies hangs on that schedule. Retaining an external owner
or collecting the upgraded references and releasing the registry guard before
summing both complete. Please ensure all upgraded references are dropped
outside the registry lock and add a last-owner sampling race test.
##########
native/common/src/bin/analyze_trace.rs:
##########
@@ -161,12 +178,13 @@ fn main() {
// After each allocated or pool update, check the current state. A
comparison needs one
// sample of each side: an observed zero reservation is a real value
that allocation can
// exceed, so only the absence of any pool sample defers the check.
- let pool_total: u64 = pool_by_thread.values().sum();
+ let pool_total: u64 = pool_total_counter.unwrap_or_else(||
pool_by_thread.values().sum());
if pool_total > peak_pool_total {
peak_pool_total = pool_total;
}
- if source.is_some() && !pool_by_thread.is_empty() && latest_allocated
> pool_total {
+ let have_pool_sample = pool_total_counter.is_some() ||
!pool_by_thread.is_empty();
+ if source.is_some() && have_pool_sample && latest_allocated >
pool_total {
Review Comment:
### Correctness
[P2] Pair allocation samples with the corresponding new pool total
Once `pool_total_counter` is set, this still checks every
allocation/per-thread event against the previous total. A valid producer
sequence is `(allocated=0, per-thread=0, total=0)`, then a busy-poll
reservation event of 100 MiB, followed by `(allocated=100 MiB, per-thread=100
MiB, total=100 MiB)` at executePlan completion. The base reports zero excess
for this sequence. This head records a false 100 MiB peak at the
allocation/per-thread events and never retracts it when the matching total
arrives. The legacy trace with the total events removed also reports zero
excess. Please evaluate the process-wide comparison only after the
corresponding total arrives, with appropriate sample association for
interleaved threads, rather than combining a new allocation with an old
reservation total.
--
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]
