Re: [PR] fix: enforce fair memory share per reservation [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on code in PR #5466:
URL: https://github.com/apache/datafusion-comet/pull/5466#discussion_r3905935196


##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -187,4 +191,29 @@ impl MemoryPool for CometFairMemoryPool {
     fn reserved(&self) -> usize {
         self.state.lock().used
     }
+
+    fn memory_limit(&self) -> MemoryLimit {
+        MemoryLimit::Finite(self.pool_size)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::execution::memory_pool::UnboundedMemoryPool;
+
+    #[test]
+    fn fair_share_uses_requesting_reservation_and_reports_pool_limit() {
+        let backing: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default());
+        let other = MemoryConsumer::new("other").register(&backing);
+        let requesting = MemoryConsumer::new("requesting").register(&backing);
+        other.grow(10);
+        requesting.grow(6);
+
+        assert_eq!(fair_limit_exceeded(32, 2, &requesting, 10), None);
+        assert_eq!(fair_limit_exceeded(32, 2, &requesting, 11), Some((6, 16)));
+
+        let pool = CometFairMemoryPool::new(Arc::new(Global::null()), 32);

Review Comment:
   This test only reaches the free function with a hardcoded `num` of 2, so nothing exercises `try_grow` itself. If someone reverted `try_grow` back to passing `state.used`, this would still pass.
   
   The `other` reservation also has no bearing on either assertion, since `fair_limit_exceeded` never looks at other consumers. The test would behave identically without it, so it does not really reproduce the 10 plus 6 case in its name.
   
   Would it be possible to register both consumers on the `CometFairMemoryPool` so that `num` comes from the pool's own bookkeeping rather than a literal? I realize `try_grow` is awkward to reach because `acquire` needs a live JVM. If putting the JNI calls behind a small trait is more than you want to take on here, then a comment explaining why the test stops at the helper would at least make the gap visible to the next person reading it.
   
   On the `memory_limit()` half, I checked and `Global::null()` is safe to drop without a JVM in jni 0.22.4 because `Drop` short-circuits on a null reference, so no concern there.



##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -142,21 +152,15 @@ impl MemoryPool for CometFairMemoryPool {
 
     fn try_grow(
         &self,
-        _reservation: &MemoryReservation,
+        reservation: &MemoryReservation,
         additional: usize,
     ) -> Result<(), DataFusionError> {
         if additional > 0 {
             let mut state = self.state.lock();
             let num = state.num;
-            let limit = self
-                .pool_size
-                .checked_div(num)
-                .expect("overflow in checked_div");
-            // We use state.used instead of reservation.size() because DataFusion 53+
-            // calls pool.try_grow() before incrementing the reservation's atomic size,
-            // so reservation.size() would not include prior grows.
-            let used = state.used;
-            if limit < used + additional {
+            if let Some((used, limit)) =
+                fair_limit_exceeded(self.pool_size, num, reservation, additional)
+            {
                 return resources_err!(
                     "Failed to acquire {additional} bytes where {used} bytes already reserved and the fair limit is {limit} bytes, {num} registered"
                 );

Review Comment:
   I want to back up @sunchao's point about losing the aggregate bound, with two specifics that I think make it more concrete.
   
   The first is that there is a documentation claim this breaks. `docs/source/user-guide/latest/tuning.md` lines 66 to 67 say the shared pool "ensures that the combined memory usage stays within the per-task limit". After this change nothing in the fair pool enforces that. `state.used` becomes pure bookkeeping and the only real ceiling left is Spark's `TaskMemoryManager`, which applies the full off-heap size and knows nothing about `spark.comet.exec.memoryPool.fraction`.
   
   The second is that the sibling-reservation case is not hypothetical for the operators we actually run. `new_empty()` and `split()` show up in `sorts/sort.rs:644`, `sorts/streaming_merge.rs:248`, `sorts/stream.rs:187` and `:253`, and `aggregates/row_hash.rs:1368`. Each sibling shares one consumer registration but carries its own size counter, so each one independently gets a full `pool_size / num` allowance. A single sort can hold several multiples of its intended share.
   
   In fairness, upstream `FairSpillPool` has the same late-registration overshoot property, so this is not a hazard we are inventing. But our pool takes a `pool_size` and after this change uses it only as a divisor. Could you keep the corrected per-reservation check and add `state.used + additional > pool_size` next to it? If you would rather not, then I think `tuning.md` lines 66 to 70 need updating in this PR, since the `pool_size / num_reservations` wording on line 70 also gets less accurate once the denominator is registrations and the numerator is per-reservation.



##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -44,6 +43,17 @@ struct CometFairPoolState {
     num: usize,
 }
 
+fn fair_limit_exceeded(
+    pool_size: usize,
+    num: usize,
+    reservation: &MemoryReservation,
+    additional: usize,
+) -> Option<(usize, usize)> {
+    let used = reservation.size();
+    let limit = pool_size.checked_div(num).expect("overflow in checked_div");

Review Comment:
   Since this is moving into a helper anyway, could it use `.unwrap_or(pool_size)` the way `FairSpillPool` does with `.checked_div(state.num_spill).unwrap_or(spill_available)`? The current message describes the wrong failure, because `checked_div` returns `None` on division by zero rather than overflow, and a panic in here aborts the process instead of surfacing as a Spark error. The helper is now a free function that is directly callable with `num = 0`, so it is a bit easier to hit than it was before.



-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org