sunchao commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r3928405605
##########
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:
[P1] Avoid retaining the byte a waiting acquire may need
Could you avoid withholding one byte until `finish_acquire`? This can
deadlock the default off-heap pool. In a 1 GiB Spark execution pool, let
another task hold 900 MiB, let this task hold 100 MiB, and let a second native
consumer for this task request 100 MiB. Spark parks that request because this
task is below its 1/(2N) minimum share. When the holder frees 100 MiB, this
branch sends only 100 MiB - 1. On wake, Spark computes `toGrant = 100 MiB - 1`;
because that is short of the request and `curMem + toGrant = 100 MiB` is still
below 256 MiB, it waits again. The deferred byte is paid only by
`finish_acquire`, which cannot run while this acquire is waiting.
I reproduced this against unmodified Spark 4.1.3 `ExecutionMemoryPool`:
retaining one byte left the grower in `WAITING`, and freeing one additional
byte from the other task let it complete. The new stub test misses this because
it grants the full request after any release without reapplying Spark's
free-memory and minimum-share checks. Could you preserve the task entry without
withholding capacity needed by the waiter?
--
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]
