sunchao commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4105107272
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -148,37 +295,90 @@ impl MemoryPool for CometFairMemoryPool {
additional: usize,
) -> Result<(), DataFusionError> {
if additional > 0 {
- let mut state = self.state.lock();
- let num = state.num;
- let limit = self
- .pool_size
- .checked_div(num)
- .expect("overflow in checked_div");
- // We use state.used instead of reservation.size() because
DataFusion 53+
- // calls pool.try_grow() before incrementing the reservation's
atomic size,
- // so reservation.size() would not include prior grows.
- let used = state.used;
- if limit < used + additional {
- return resources_err!(
- "Failed to acquire {additional} bytes where {used} bytes
already reserved ({} bytes overcommitted) and the fair limit is {limit} bytes,
{num} registered",
- self.spark.overcommit()
- );
- }
-
- // A partial grant is handed back and refused, which triggers
spilling in the caller.
- if let Err(refusal) = self.spark.try_acquire(additional)? {
- return resources_err!(
- "Failed to acquire {} bytes plus {} bytes overcommitted,
only got {} bytes. Reserved: {} bytes",
- additional,
- refusal.overcommit,
- refusal.granted,
- state.used
- );
- }
- state.used = state
- .used
- .checked_add(additional)
- .expect("overflow in checked_add");
+ // Checking the fair limit and reserving the bytes is one atomic
step, so concurrent
+ // grows can never jointly exceed pool_size / num. The blocking
JVM calls then run
+ // without any lock held, and the reservation rolls back if the
JVM does not back it.
+ {
+ let mut state = self.state.lock();
+ let num = state.num;
+ let limit = self
+ .pool_size
+ .checked_div(num)
+ .expect("overflow in checked_div");
+ // The pool tracks one total across every consumer and checks
the fair limit
+ // against that total, not against this reservation's own size.
+ let used = state.used;
+ match used.checked_add(additional) {
+ Some(total) if total <= limit => state.used = total,
+ _ => {
+ return resources_err!(
+ "Failed to acquire {additional} bytes where {used}
bytes already reserved ({} bytes overcommitted) and the fair limit is {limit}
bytes, {num} registered",
+ self.spark.overcommit()
+ );
+ }
+ }
+ }
+
+ // The anchor comes after the local limit check, so a grow the
pool rejects itself
+ // never makes a JVM call, and before the real request, so the
byte is held before
+ // the balance can reach zero. The JVM call can panic inside its
JNI frame; the
+ // optimistic reservation must not outlive either call, or the
leaked bytes poison
+ // the task-shared pool for every other consumer.
+ match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+ self.take_missing_anchor()
+ })) {
+ Ok(Ok(())) => {}
+ Ok(Err(e)) => {
+ self.settle_acquire(additional, 0);
+ return Err(e.into());
+ }
+ Err(panic) => {
+ self.settle_acquire(additional, 0);
+ std::panic::resume_unwind(panic);
+ }
+ }
+ // Spark is asked for the request plus any outstanding overcommit,
and a full grant
+ // repays the overcommit. A short grant stays with Spark until
this pool hands it
+ // back below, so the bytes can stay charged meanwhile.
+ let refusal = match
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+ self.spark.try_acquire_leaving_a_short_grant(additional)
Review Comment:
[P1] [P1] Prevent `showMemoryUsage` from blocking short-grant rollback
Allowing these calls to overlap exposes a JVM lock cycle.
`CometTaskMemoryManager.acquireMemory` releases Spark's task monitor after
acquiring memory, then reacquires it in `showMemoryUsage` when the grant is
short. In a 100-byte executor pool, let other tasks hold 82 and 1 bytes and
this pool hold its anchor. A 30-byte request receives 16. Before its diagnostic
reacquires the monitor, release the other task's 1 byte and start a second
10-byte request. The second request parks below its new 25-byte minimum while
holding the task monitor. The first blocks in `showMemoryUsage`, so it cannot
return to Rust and release the 16-byte short grant that would satisfy the
second request. Both requests fit the local fair limit with two consumers.
Expected behavior is to return the short grant, spill, and let the second
request proceed. Instead, the task can hang until an unrelated task frees
memory. The previous fair-pool mutex excluded this interleaving. Could the
bridge avoid reacquiring the task monitor after obtaining the grant, or cover
acquisition and diagnostics with one monitor scope while keeping releases
independent?
Evidence: Bounded reproduction:
`/tmp/pr5613-bb273df5-jvm-probe/ShortGrantProbe.scala`, compiled with the
unchanged head's `CometTaskMemoryManager.java` against Spark 4.1.3. A logging
latch schedules the interleaving and is released before observing the failure.
Output in `result.log`: `head: first=BLOCKED, second=WAITING, held=16+anchor`.
Stacks identify `TaskMemoryManager.showMemoryUsage:325` and
`ExecutionMemoryPool.acquireMemory:142`. Releasing 50 bytes from the unrelated
task rescues both threads. A control emulating the previous serialization
completes without that release and ends with zero balance. Source inspection
confirms the same monitor structure across all five checked Spark versions.
--
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]
