dwsmith1983 commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4242442472
##########
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:
> Please coordinate in-flight debt coverage while keeping releases able to
progress.
Done in 8cf17395e. `SparkMemory` now keeps the overcommit and the part of it
that in-flight requests have already asked Spark for in one small ledger,
behind a mutex held only for the arithmetic and never across a Spark call, so a
release still repays and moves on. Each `try_grow` claims only the debt nobody
else is asking for: in your case the second asks Spark for its own byte and
completes while the first is still paused, and the requests come out as 17, 2,
1. A refusal, an error or a panic hands the claim back for the next `try_grow`,
and bytes a release repaid while Spark was granting them still go back to Spark.
The unified pool reached the same code without a lock on main, so it gets
the fix too. Your scenario is now a test in `fair_pool.rs`. One window remains,
and the memory guide describes it: a release that repays the debt while a
claimed request is parked in Spark leaves that request larger than needed,
since Spark cannot shrink a parked request.
--
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]
