dwsmith1983 commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r3930692784
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -60,26 +100,102 @@ impl CometFairMemoryPool {
task_memory_manager_handle: Arc<Global<JObject<'static>>>,
pool_size: usize,
) -> CometFairMemoryPool {
+ Self::with_bridge(
+ Box::new(JniTaskMemoryBridge {
+ task_memory_manager_handle,
+ }),
+ pool_size,
+ )
+ }
+
+ fn with_bridge(bridge: Box<dyn TaskMemoryBridge>, pool_size: usize) ->
CometFairMemoryPool {
Self {
- task_memory_manager_handle,
+ bridge,
pool_size,
- state: Mutex::new(CometFairPoolState { used: 0, num: 0 }),
+ state: Mutex::new(CometFairPoolState {
+ used: 0,
+ num: 0,
+ jvm_held: 0,
+ pending_acquires: 0,
+ deferred_release: 0,
+ paying_deferred: false,
+ }),
+ deferred_done: Condvar::new(),
}
}
fn acquire(&self, additional: usize) -> CometResult<i64> {
- let handle = self.task_memory_manager_handle.as_obj();
- JVMClasses::with_env(|env| unsafe {
- jni_call!(env,
- comet_task_memory_manager(handle).acquire_memory(additional as
i64) -> i64)
- })
+ self.bridge.acquire(additional)
}
fn release(&self, size: usize) -> CometResult<()> {
- let handle = self.task_memory_manager_handle.as_obj();
- JVMClasses::with_env(|env| unsafe {
- jni_call!(env,
comet_task_memory_manager(handle).release_memory(size as i64) -> ())
- })
+ self.bridge.release(size)
+ }
+
+ /// Debits a release from the JVM-side balance and returns how much to
hand back now. A
+ /// release that would zero the balance while acquires are in flight keeps
one byte back,
+ /// because Spark drops the task's accounting entry at zero and a parked
acquire then indexes
+ /// the missing entry. Blocking instead could deadlock: the waiter may
need this very memory.
+ /// The n-1 bytes freed here still wake Spark's waiter; the single held
byte only matters
+ /// in a pool small enough that one byte decides the fair-share threshold,
and even there
+ /// the deferred payoff releases it as soon as in-flight acquires drain.
+ fn plan_release(state: &mut CometFairPoolState, bytes: usize) -> usize {
+ state.jvm_held = state
+ .jvm_held
+ .checked_sub(bytes)
+ .expect("released more bytes than the JVM side holds");
+ if bytes > 0 && state.jvm_held == 0 && state.pending_acquires > 0 {
Review Comment:
Fixed in db1f1bc6e. The deferral is gone: every release now goes to the JVM
whole, so a parked acquire sees the full freed amount and Spark's minimum-share
check passes. The task entry is kept alive by a permanent anchor instead. The
first acquire asks for one extra byte and the pool holds it until it drops, so
the balance never reaches zero mid task. That means the task stays in Spark's
active-task set for the pool's lifetime and retains one byte, which the header
comment now states.
Your scenario is pinned as
`min_share_wait_is_granted_after_the_holder_frees_its_memory_in_full`. The stub
now models `ExecutionMemoryPool.acquireMemory` from 4.1.3 with the per-task
entry lifecycle, the 1/N and 1/(2N) shares, the wait loop, and `notifyAll` on
release, with a bounded wait that fails the test instead of hanging. It failed
on the previous head with the waiter timing out and passes now.
--
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]
