sunchao commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r3910247103
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -124,19 +162,25 @@ impl MemoryPool for CometFairMemoryPool {
fn shrink(&self, _reservation: &MemoryReservation, subtractive: usize) {
if subtractive > 0 {
- let mut state = self.state.lock();
- // We don't use reservation.size() here because DataFusion 53+
decrements
- // the reservation's atomic size before calling pool.shrink(), so
it would
- // reflect the post-shrink value rather than the pre-shrink value.
- if state.used < subtractive {
- panic!(
- "Failed to release {subtractive} bytes where only {} bytes
tracked by pool",
- state.used
- )
+ {
+ let mut state = self.state.lock();
+ // We don't use reservation.size() here because DataFusion 53+
decrements
+ // the reservation's atomic size before calling pool.shrink(),
so it would
+ // reflect the post-shrink value rather than the pre-shrink
value.
+ if state.used < subtractive {
+ panic!(
+ "Failed to release {subtractive} bytes where only {}
bytes tracked by pool",
+ state.used
+ )
+ }
+ state.used -= subtractive;
}
Review Comment:
[P2] Keep waiting Spark acquisitions registered during a full release
Could you handle Spark's same-task wait/release contract before allowing
this release to overlap an acquire? With `fair_unified`, two native consumers
can now enter the same `CometTaskMemoryManager` concurrently. In a 100-unit
executor pool, let another task hold 90 and this task hold 10. This task's next
10-unit grow waits below Spark's 1/(2N) minimum. Freeing its last 10 units on
another native thread removes its entry from
`ExecutionMemoryPool.memoryForTask` and wakes the grower. The grower then
[indexes the removed
entry](https://github.com/apache/spark/blob/7c14a3c28b141cc97a330c4d0f5d2a6da7267f85/core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala#L114)
and throws `NoSuchElementException: key not found`. Spark's release bypasses
the task monitor held by the waiting acquire, so that monitor does not prevent
this interleaving. The Rust provisional reservation does not keep Spark's entry
alive.
I reproduced the failure using unchanged Spark 3.5.9 pool source with only
logging/annotation/memory-mode scaffolding. A scheduling control modeling the
previous serialization completed after the other task freed its memory. The
relevant map lifecycle is also present in 4.0.4 source. This was a component
probe plus JNI source tracing, not a full Comet query reproduction. Please make
full releases safe while grows are pending and add a regression that exercises
Spark's memory manager.
--
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]
