sunchao commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4242214184
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -152,45 +211,85 @@ impl MemoryPool for CometFairMemoryPool {
additional: usize,
) -> Result<(), DataFusionError> {
if additional > 0 {
- let mut state = self.state.lock();
- let num = state.consumers.len();
- let limit = self
- .pool_size
- .checked_div(num)
- .expect("overflow in checked_div");
- let consumer_used = *state.consumer_used(reservation);
- if limit < consumer_used.saturating_add(additional) {
- return resources_err!(
- "Failed to acquire {additional} bytes where this consumer
already holds {consumer_used} bytes and the fair limit is {limit} bytes, {num}
registered ({} bytes overcommitted)",
- self.spark.overcommit()
- );
- }
- // The shares alone do not bound the pool's total, because a
consumer keeps what it
- // reserved before another consumer registered.
- let used = state.used;
- if self.pool_size < used.saturating_add(additional) {
- return resources_err!(
- "Failed to acquire {additional} bytes where {used} bytes
already reserved ({} bytes overcommitted) and the pool limit is {} bytes",
- self.spark.overcommit(),
- self.pool_size
- );
+ let consumer = reservation.consumer().id();
+ // Checking both limits and charging the bytes is one atomic step,
so concurrent grows
+ // can never jointly take a consumer past its share or the pool
past pool_size. The
+ // blocking JVM calls then run without any lock held, and the
charge rolls back if
+ // the JVM does not back it.
+ {
+ let mut state = self.state.lock();
+ let num = state.consumers.len();
+ let limit = self
+ .pool_size
+ .checked_div(num)
+ .expect("overflow in checked_div");
+ let consumer_used = *state.consumer_used(consumer);
+ if limit < consumer_used.saturating_add(additional) {
+ return resources_err!(
+ "Failed to acquire {additional} bytes where this
consumer already holds {consumer_used} bytes and the fair limit is {limit}
bytes, {num} registered ({} bytes overcommitted)",
+ self.spark.overcommit()
+ );
+ }
+ // The shares alone do not bound the pool's total, because a
consumer keeps what
+ // it reserved before another consumer registered.
+ let used = state.used;
+ if self.pool_size < used.saturating_add(additional) {
+ return resources_err!(
+ "Failed to acquire {additional} bytes where {used}
bytes already reserved ({} bytes overcommitted) and the pool limit is {} bytes",
+ self.spark.overcommit(),
+ self.pool_size
+ );
+ }
+ state.charge(consumer, additional);
}
- // 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
- );
+ // 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. The JVM
call can panic inside
+ // its JNI frame; the optimistic reservation must not outlive it,
or the leaked bytes
+ // poison the task-shared pool for every other consumer.
+ let refusal = match
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+ self.spark.try_acquire_leaving_a_short_grant(additional)
Review Comment:
[P2] Coordinate concurrent acquisitions’ overcommit coverage. Running this
call unlocked lets two `try_grow` calls snapshot the same debt. With one byte
of debt and three bytes free, the first `try_grow(1)` acquires two bytes. If it
pauses after Spark returns but before `repay`, a second `try_grow(1)` also
requests two bytes, although only one remains free. Below Spark’s minimum
share, that request parks. The first then clears the debt, but the second
remains blocked on its stale request and cannot reach the surplus-return logic.
Expected behavior is successful acquisition of the available byte. Instead, a
default `fair_unified` task can stall until another release supplies
unnecessary capacity. Base serialization makes the second request one byte and
complete. Please coordinate in-flight debt coverage while keeping releases able
to progress.
Evidence: Bounded regression
`review_repaid_debt_does_not_park_an_exact_fit_sibling`, using unchanged head
production pool code: both local and Spark pools contain 100 bytes. Other tasks
hold 83 and 1. `grow(17)` receives 16, creating debt 1. Release the third
task’s byte and two bytes from the other task. Gate the first `try_grow(1)`
after its two-byte Spark grant, start the second `try_grow(1)`, then resume the
first. Head reports `requests=[17, 2, 2] debt=0 free=1 second_completed=false`;
pristine base reports `requests=[17, 2, 1] second_completed=true`. A freshly
compiled head Java bridge against real Spark 4.1.3 confirms the two-byte
request waits with task usage 18, other-task usage 81 and one byte free; the
one-byte control completes immediately. Releasing one external byte rescues the
waiter. Artifacts:
`/tmp/pr5613-review-final-6a90-8jdrim1x/native/head-probe.log`,
`native/base-probe.log`, and `jvm/debt-probe.log`.
--
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]
