Try fix leak

Author: Vedin

datafusion/physical-plan/src/joins/grace_hash_join/exec.rs

@@ -1163,11 +1163,9 @@ fn process_partition_chunk(
         reservation.try_grow(actual_reserved - reserved_bytes)?;
     } else if reserved_bytes > actual_reserved {
         let delta = reserved_bytes - actual_reserved;
-        let available = reservation.size();
-        if delta > available {
-            reservation.shrink(available);
-        } else {
-            reservation.shrink(delta);
+        if reservation.try_shrink(delta).is_err() {
+            // If accounting drifted, free whatever is left to avoid leaking
+            let _ = reservation.free();
         }
     }
     reserved_bytes = actual_reserved;
@@ -1190,11 +1188,8 @@ fn process_partition_chunk(
         partitions[i].spill_batch_auto(&part_batch, &request_msg, reservation)?;
     }
 
-    let available = reservation.size();
-    if reserved_bytes > available {
-        reservation.shrink(available);
-    } else {
-        reservation.shrink(reserved_bytes);
+    if reservation.try_shrink(reserved_bytes).is_err() {
+        let _ = reservation.free();
     }
     Ok(())
 }
@@ -1272,11 +1267,8 @@ impl PartitionWriter {
         let total_accounted = self.buffered_bytes;
         let coalesced = concat_batches(&self.schema, &self.buffer)?;
         self.buffer.clear();
-        let available = reservation.size();
-        if total_accounted > available {
-            reservation.shrink(available);
-        } else {
-            reservation.shrink(total_accounted);
+        if reservation.try_shrink(total_accounted).is_err() {
+            let _ = reservation.free();
         }
         self.buffered_bytes = 0;
 

datafusion/physical-plan/src/joins/grace_hash_join/stream.rs

@@ -990,20 +990,13 @@ impl GraceHashJoinStream {
                             };
                             if bytes_to_free > 0 {
                                 let mut res = reservation.lock();
-                                let available = res.size();
-                                if bytes_to_free > available {
-                                    // Don't panic on accounting drift; free what we can and log.
-                                    let shrink = available;
-                                    if shrink > 0 {
-                                        res.shrink(shrink);
-                                    }
+                                if res.try_shrink(bytes_to_free).is_err() {
+                                    let freed = res.free();
                                     debug!(
-                                        "Grace hash join reservation underflow: attempted to free {}, but reservation size is {}",
-                                        human_readable_size(bytes_to_free),
-                                        human_readable_size(available)
+                                        "Grace hash join reservation underflow: freed {} after shrink failure (requested {})",
+                                        human_readable_size(freed),
+                                        human_readable_size(bytes_to_free)
                                     );
-                                } else {
-                                    res.shrink(bytes_to_free);
                                 }
                             }
                             *left_fut = None;
@@ -1223,18 +1216,13 @@ impl GraceHashJoinStream {
                                 };
                                 if bytes_to_free > 0 {
                                     let mut res = self.reservation.lock();
-                                    let available = res.size();
-                                    if bytes_to_free > available {
-                                        if available > 0 {
-                                            res.shrink(available);
-                                        }
+                                    if res.try_shrink(bytes_to_free).is_err() {
+                                        let freed = res.free();
                                         debug!(
-                                            "Grace hash join reservation underflow on stream completion: attempted to free {}, but reservation size is {}",
-                                            human_readable_size(bytes_to_free),
-                                            human_readable_size(available)
+                                            "Grace hash join stream completion freed {} after shrink failure (requested {})",
+                                            human_readable_size(freed),
+                                            human_readable_size(bytes_to_free)
                                         );
-                                    } else {
-                                        res.shrink(bytes_to_free);
                                     }
                                 }
                                 if let Some(start) = current_join_start.take() {