mbutrovich commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4360808927
##########
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:
> Running this call unlocked lets two `try_grow` calls snapshot the same
debt.
As I read it, the fair pool has this race only because this PR lets two
acquires run at once. With acquires serialized in the pool, which Spark already
does per task, two `try_grow` calls can't snapshot the same debt, and the fair
pool doesn't need the claim ledger. The unified pool's version of this race is
already on `main`.
--
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]
